Technical Summary: Diff-Instruct with Diffused Reward (Didr)

Problem Statement

Recent advances in one-step text-to-image (T2I) generators have enabled real-time synthesis by mapping latent noise directly to images in a single forward pass. However, aligning these generators with human preferences via Reinforcement Learning (RL) presents a fundamental structural challenge.

Previous methods for one-step RLHF (e.g., Diff-Instruct*, Diff-Instruct++) combine image-space reward optimization with latent-space distribution matching. These approaches typically apply preference signals (rewards) exclusively at the clean image endpoint ($x_0$) while imposing Kullback-Leibler (KL) regularization against the reference distribution across the entire noisy diffusion trajectory. The authors identify a critical structural mismatch in this paradigm: the reward acts only at the terminal state, while the regularizer spans all noise levels.

This mismatch leads to a failure mode termed terminal reward domination. Because forward noise naturally masks image details, the KL regularization penalty against mode divergence weakens significantly at high noise levels. Consequently, optimizers exploit this trajectory vulnerability, abandoning the reference distribution to collapse onto high-reward regions without incurring sufficient KL penalties. This results in improved reward scores but often at the expense of image fidelity and distributional stability.

Methodology: Diff-Instruct with Diffused Reward (Didr)

To resolve this misalignment, the authors propose Didr, a data-free, trajectory-level alignment framework derived from Integral KL (IKL) minimization. The core insight is to systematically propagate the RLHF-optimal reward-tilted clean-image distribution across all noise levels of the diffusion trajectory.

1. The Principled Trajectory Target

The method begins by establishing the correct alignment target for clean images under KL-regularized RLHF. The optimal clean-image distribution $q^*(x_0|c)$ is the reward-tilted density:
$$q^*(x_0|c) \propto q_0(x_0|c) \exp(r(x_0, c)/\tau)$$
where $q_0$ is the reference distribution, $r$ is the reward function, and $\tau$ is a temperature parameter.

Instead of applying rewards only at $x_0$, Didr "lifts" this optimal target to the full trajectory by diffusing $q^*$ through the reference forward process. This yields a family of reward-tilted marginals $\{q^*_t\}_{t \in [0, T]}$ defined at every noise level:
$$q^*_t(x_t|c) = \int q_t(x_t|x_0) q^*(x_0|c) dx_0$$
The authors prove that minimizing the Integral KL divergence between the generator's marginals and these reward-tilted marginals shares the same minimizer as the standard KL-regularized RLHF objective. This ensures that reward and regularization are balanced at every noise level, preventing terminal reward domination.

2. The Diffused Reward Score (DRS)

By differentiating the IKL objective, the authors derive the target score function. The optimal score decomposes into two components:
$$s_{target} = s_{ref} + s_r$$
where $s_{ref}$ is the reference score function, and $s_r$ is the Diffused Reward Score (DRS). The DRS acts as a reward-driven correction term defined at every noise level:
$$s_r(x_t, t, c) = \nabla_{x_t} \log \mathbb{E}_{x_0 \sim q(x_0|x_t, c)} \left[ \exp\left(\frac{r(x_0, c)}{\tau}\right) \right]$$
Unlike classifier guidance (which is a first-order approximation), the DRS retains the full log-exponential tilt, propagating preference information only where the clean image is identifiable from the noisy observation.

3. The Diffused Reward Proxy (DRP)

Since the DRS requires an intractable posterior expectation $q(x_0|x_t, c)$, the authors introduce the Diffused Reward Proxy (DRP) to make the method practical. The DRP approximates the posterior expectation using differentiable short-step denoising chains from the frozen reference model.

• Mechanism: From a noisy state $x_t$, the method runs $K$ independent $S$-step denoising chains using the frozen reference model to generate approximate clean samples $\{\hat{x}^{(k)}_0\}$.
• Gradient Estimation: The DRS is approximated as a softmax-weighted average of pathwise reward gradients:
  $$\hat{s}_r(x_t, t, c) = \frac{1}{\tau} \sum_{k=1}^K \omega^{(k)} \nabla_{x_t} r(\hat{x}^{(k)}_0, c)$$
  where weights $\omega^{(k)}$ are proportional to $\exp(r(\hat{x}^{(k)}_0, c)/\tau)$.
• Efficiency: This approach allows reward gradients to propagate stably back through the trajectory without storing reference model activations or requiring image training data.

4. Training Framework

The training process alternates between two stages:

1. Teaching Assistant (TA) Update: A TA score model ($s_\psi$) is trained via Denoising Score Matching (DSM) to track the generator's evolving marginals, surrogating the implicit generator score.
2. Generator Update: The generator ($g_\theta$) is updated using the gradient of the IKL objective, which compares the TA score against the sum of the CFG-corrected reference score and the DRP estimate.

Key Contributions

1. Identification of Terminal Reward Domination: The paper formally identifies and analyzes a structural failure mode in prior one-step RLHF methods where optimizers exploit noise levels to collapse onto high-reward regions, violating the intended RLHF balance.
2. Principled Trajectory-Level Alignment: The authors propose Didr, a framework that derives a mathematically sound trajectory-level target by propagating the RLHF-optimal clean-image density forward. This ensures reward and regularization are balanced at every noise level.
3. Analytical Derivation of DRS: The paper analytically derives the Diffused Reward Score, showing it exactly decomposes the trajectory target into a reference score plus a reward-induced correction.
4. Practical Estimator (DRP): The authors introduce the Diffused Reward Proxy, a zero-data estimator that computes the DRS via differentiable short-step denoising, enabling stable training without image data.

Empirical Results

Extensive experiments demonstrate that Didr consistently outperforms existing one-step baselines:

• Pareto Dominance: On the SDXL backbone (1024×1024), the standard Didr variant Pareto-dominates all existing one-step alignment baselines in the PickScore–FID trade-off, simultaneously surpassing multi-step models like FLUX-dev and SDXL-DPO on both axes.
• Preference Metrics: The extended training variant (Didr_longer) achieves the highest PickScore (23.9) among evaluated one-step models and improves ImageReward by 8.9% over the strongest prior baseline (Diff-Instruct*). It also leads in HPSv2.1 and Aesthetic Score.
• Generalization: Didr successfully transfers to the 6B-parameter Z-Image backbone. In a single generation step, it surpasses its 50-step teacher model on preference metrics, demonstrating scalability across architectures.
• Qualitative Improvements: Visual comparisons show Didr produces sharper fine details and more natural colors compared to prior one-step methods, avoiding the over-saturation and painterly artifacts often seen in reward-optimized baselines.

Significance and Claims

The paper claims that principled trajectory-level reward propagation is an effective alignment strategy that resolves the inherent misalignment of endpoint-only reward methods. By rigorously propagating the RLHF-optimal distribution through the diffusion process, Didr achieves the best trade-off between preference alignment and image fidelity among one-step generators.

The authors emphasize that their method requires no image training data, relying solely on text prompts and a frozen reference model. They position Didr as a solution that allows one-step generators to match or exceed the preference alignment of multi-step teachers while maintaining real-time inference speeds. The work suggests that the structural mismatch identified in prior methods is a fundamental limitation that can only be resolved by aligning the entire trajectory, not just the terminal state.