Technical Summary: Quality-Aware Modulation for Diffusion Transformers

Problem Statement

Modern text-to-image diffusion models, particularly Diffusion Transformers (DiT) like Stable Diffusion 3.5 (SD3.5), rely on conditioning signals—specifically timestep embeddings and pooled prompt embeddings—to modulate the denoising process via Adaptive Layer Normalization (AdaLN). While these signals effectively communicate the current noise level and semantic intent, they lack quality-aware information. Consequently, the modulation mechanism cannot directly react to sample-specific failure modes, such as error accumulation, incorrect prompt interpretation, or hallucinated features that emerge as the latent image evolves. The baseline conditioning is independent of the current latent state ($\partial M_{base} / \partial x_t = 0$), preventing the model from correcting fidelity issues based on the intermediate image's actual quality.

Methodology

The authors propose the Quality Representation Module (QRM), a lightweight transformer-based adapter designed to inject quality-aware signals into the diffusion process without modifying the pretrained backbone.

Architecture

QRM is implemented as a transformer decoder that predicts additive modulation offsets ($M_{qrm}$) to the baseline AdaLN parameters ($M_{base}$).

• Inputs: The module receives the current latent image ($x_t$), pooled prompt embedding ($p$), timestep embedding ($t$), and the baseline modulation parameters ($M_{base}$).
• Processing:
  1. The latent image $x_t$ is projected into latent tokens via a learned $1\times1$ convolution and adaptive pooling.
  2. Baseline modulation parameters are normalized and projected to form query tokens ($Z_0$).
  3. A key-value memory ($KV$) is constructed from the latent tokens, prompt, and timestep embeddings.
  4. Through cross-attention and self-attention layers, the query tokens attend to the $KV$ memory to generate a refined representation.
  5. The output is projected back to the native AdaLN widths to produce the offset vector $M_{qrm}$.
• Application: The final modulation applied to the transformer blocks is $M_t = M_{base,t} + M_{qrm,t}$.

Training Strategy (ReFL)

QRM is trained using a Reward Feedback Learning (ReFL) strategy:

• Frozen Components: The diffusion backbone (denoiser, text encoders, VAE), the sampler, and the reward model remain frozen. Only QRM parameters are updated.
• Objective: The module learns to maximize a scalar reward score derived from a frozen reward model (e.g., HPSv2.1 or CLIPScore) applied to an image decoded from an intermediate denoising step.
• Loss Functions: The paper utilizes either a reward-maximization loss or a margin-seeking loss, which pushes the QRM-modulated image to outperform the baseline by a specific margin.
• Timestep Selection: To improve efficiency and focus on structural formation, QRM is activated only during a predefined subset of high-noise timesteps (e.g., $t \in \{50, \dots, 26\}$), rather than the entire diffusion schedule.

Inference

During inference, the trained QRM module acts as a plug-in adapter. It predicts modulation offsets based on the current latent state and adds them to the baseline parameters. This allows the model to steer the denoising trajectory toward higher-quality states without retraining the backbone or altering the sampling schedule.

Key Contributions

1. QRM Architecture: Introduction of a lightweight transformer module that learns quality-aware representations from existing model inputs (latent, prompt, timestep) to refine modulation parameters.
2. Non-Invasive Integration: A mechanism to augment diffusion transformer modulation with latent image quality signals without modifying the underlying diffusion backbone, VAE, or sampling schedule.
3. Reward-Guided Modulation: Demonstration that predicted modulation updates ($M_{qrm}$) can effectively refine the denoising process at each timestep, improving prompt adherence and image fidelity.
4. Comprehensive Evaluation: Extensive ablation studies analyzing the effects of QRM architecture size, active timestep ranges, and training reward metrics on image quality.

Experimental Results

The authors evaluated QRM using Stable Diffusion 3.5 (SD3.5) on the Parti-Prompts benchmark.

• Quantitative Performance:
  • Compared to the SD3.5 baseline, QRM trained with HPSv2.1 as the reward signal achieved the highest scores among all evaluated methods (including prompt refinement, fine-tuning, and ReNeg approaches).
  • Improvements: +0.20 in Aesthetic score and +0.82 in HPSv2.1 score.
  • When trained with CLIPScore, QRM improved human-preference alignment (HPSv2.1) while maintaining comparable perceptual quality.
• Qualitative Analysis:
  • Visualizations of denoising trajectories show that applying QRM during early, high-noise timesteps leads to recognizable semantic structures emerging significantly earlier than in the baseline.
  • Applying QRM only in middle or late stages yielded smaller improvements, confirming that global structure is established early in the diffusion process.
• Ablation Findings:
  • Capacity: Larger QRM models optimized more strongly for the specific training reward (e.g., HPSv2.1) but showed slight trade-offs in other metrics (e.g., CLIPScore).
  • Timestep Range: Modulation during early timesteps ($50 \to 26$) was consistently more effective for alignment than modulation in later stages.
  • Reward Metric: The choice of reward model (CLIPScore vs. HPSv2.1) directly shaped the perceptual biases learned by QRM.

Significance and Claims

The paper claims that modulation-level interventions provide a scalable and effective pathway for aligning large diffusion models with human-preference signals. Unlike methods that require expensive fine-tuning of the entire backbone (e.g., LoRA or full fine-tuning), QRM achieves consistent quality improvements by operating entirely through additive modulation updates conditioned on the evolving latent representation.

The authors position QRM as a flexible, quality-aware adapter that can be steered by different reward signals and applied to various diffusion transformer architectures. They acknowledge that QRM inherits the biases of the training reward models and suggest that future work could explore multi-reward training and adaptive timestep selection to further enhance robustness and generalization. The work demonstrates that significant gains in image fidelity and prompt alignment can be achieved without altering the pretrained weights of state-of-the-art models.