PaperlayerView exact revision on arXiv

Source-linked AI summary

Reinforcing Few-step Generators via Reward-Tilted Distribution Matching

Yushi Huang, Xiangxin Zhou, Ruoyu Wang, Chi Zhang, Jun Zhang, Tianyu Pang

arXiv:2605.26108v3cs.CV

TL;DR

Few-step generators are efficient but difficult to align with human preferences because distillation and reward optimization must handle shifting distributions and hybrid sampling dynamics. RTDMD unifies these objectives in a two-stage framework and achieves state-of-the-art quality across benchmarks with four-step sampling.

  • Problem

    Few-step generation must reduce sampling cost while addressing difficult reward alignment and noisy, shifting intermediate distributions.

  • Method

    RTDMD minimizes divergence to a reward-tilted teacher, combining AC-DMD stabilization with hybrid policy gradients for stochastic and deterministic sampling steps.

  • Results

    RTDMD achieves state-of-the-art few-step generation quality across SD3, SD3.5, and FLUX.2 under four-step sampling, surpassing full FLUX.2 9B on most benchmarks.

  • Takeaways & Limitations

    The framework provides a unified approach for preference-aligned, high-quality text-to-image generation with only four inference steps.

  • Takeaways & Limitations

    Accurate fake-score estimation remains difficult when the generator distribution shifts under limited score updates and tight compute budgets.

Abstract

from arXiv · show

Recent advances in few-step diffusion distillation have enabled efficient image generation, yet aligning these models with human preferences remains challenging. We propose Reward-Tilted Distribution Matching Distillation (RTDMD), a two-stage framework that unifies distribution matching distillation with reward-guided reinforcement learning for few-step flow generators. We show that minimizing the KL divergence to a reward-tilted teacher distribution naturally decomposes into a distribution matching term and a reward maximization term. In the first stage, we introduce Ambient-Consistent Distribution Matching Distillation (AC-DMD), which performs subinterval-wise distribution matching and augments the fake score objective with a consistency regularizer to help the fake score model track the shifting generator distribution under limited updates. In the second stage, we jointly optimize both terms: for the reward maximization term, we derive a hybrid policy gradient that combines a GRPO-style estimator for the stochastic intermediate transitions with direct reward backpropagation through the deterministic final step, and further introduce step-subset GRPO (SubGRPO) to reduce variance. Experiments on SD3, SD3.5, and FLUX.2 demonstrate that RTDMD establishes new state-of-the-art results across preference, aesthetic, and compositional metrics with only 4 inference steps, outperforming previous few-step text-to-image generation methods. Code and models are available at https://github.com/Harahan/RTDMD.

1 Introduction

RTDMD addresses the difficulty of reward-guided few-step generation with a two-stage framework that unifies distribution matching and reinforcement learning through a reward-tilted teacher distribution. Experiments show state-of-the-art quality under 4-step sampling across SD3-M, SD3.5-M, and FLUX.2 4B.

  • Motivation: Diffusion and flow-based generators produce realistic, semantically aligned images, but iterative sampling typically requires tens of denoising or flow-integration steps.This motivates distillation into more efficient few-step generators.
  • Challenges: Reward-guided few-step generation is difficult because noisy intermediate latents and shifting generator distributions require the fake score model to track a moving target under limited compute.These issues make cold-start distillation challenging.
  • Method: RTDMD minimizes KL divergence to a reward-tilted teacher distribution, decomposing the objective into distribution matching and reward maximization across two training stages.The first stage introduces Ambient-Consistent DMD as a stable cold start.
  • Results: RTDMD achieves state-of-the-art few-step generation quality under 4-step sampling on SD3-M, SD3.5-M, and FLUX.2 4B.The distilled FLUX.2 4B surpasses the full FLUX.2 9B using 50-step sampling across most benchmarks.

2 Preliminaries

Diffusion and flow-based models define probability paths from data to a Gaussian prior and sample via velocity fields governed by the PF-ODE. Flow matching uses a rectified linear path, while DMD distills pretrained flows by matching noised student and teacher marginals with an auxiliary fake velocity model.

  • Diffusion and flow-based models: Diffusion and flow-based models define a continuous probability path connecting data distribution p0 to a simple prior p1, typically a standard Gaussian.Gaussian interpolation uses xt = αtx0 + σtϵ, with αt and σt specifying the noise schedule.
  • Diffusion and flow-based models: Sampling follows the probability-flow ODE dxt/dt = v(xt, t), where the marginal velocity transports density along the probability path.The score describes local density changes, whereas the marginal velocity determines sample movement.
  • Flow matching: Flow matching adopts the rectified schedule αt = 1−t and σt = t, yielding the linear path xt = (1−t)x0 + tϵ.For fixed (x0, ϵ), the conditional velocity is ϵ − x0, and neural velocity fields are trained with the conditional flow matching loss.
  • Distribution matching distillation: Few-step distillation addresses the many function evaluations required by pretrained flow models by training a student generator whose distribution matches the teacher distribution.DMD uses reverse KL divergence but compares noised marginals in an ambient space because direct data-space optimization can be difficult when distributions have limited overlap.
  • Distribution matching distillation: DMD trains an auxiliary fake velocity field to track the current student distribution and uses its induced fake score as a surrogate for the unavailable student score.Training alternates between updating the student generator Gθ and fitting the fake model vϕ.

3 Method

RTDMD trains few-step generators by minimizing divergence to a reward-tilted teacher distribution, decomposing this objective into distribution matching and reward maximization. Its two stages combine AC-DMD for cold-start distillation with hybrid policy-gradient optimization using step-subset GRPO.

  • RTDMD objective: RTDMD up-weights high-reward regions of the teacher distribution and minimizes KL divergence to this reward-tilted target.The reward strength is controlled by β ≥ 0, while the normalizing constant Z is independent of the generator parameters.
  • RTDMD objective: The reward-tilted KL objective decomposes into a distribution matching term and a reward maximization term, motivating RTDMD’s two-stage training framework.AC-DMD provides the distribution-matching cold start, followed by joint optimization with hybrid policy gradients and step-subset GRPO.
  • Ambient-Consistent Distribution Matching Distillation: AC-DMD uses coefficient-preserving sampling to unify deterministic Euler and consistency-model samplers, with η controlling intermediate sampling stochasticity.η = 0 recovers deterministic Euler sampling, while η > 0 injects noise at each step; the generator uses K = 4 steps.
  • Ambient-Consistent Distribution Matching Distillation: AC-DMD performs distribution matching on subintervals conditioned on noisy intermediate latents and adds consistency regularization to stabilize fake-score training as the generator distribution shifts.The regularizer enforces agreement between nearby-timestep x-predictions and reduces the variance of ambient fake-score training.
  • Reward maximization: SubGRPO uses shared noise to isolate selected-step effects and typically achieves lower gradient-estimator variance than independent-noise GRPO under the same gradient sample budget.It is characterized as a Rao–Blackwellized variant of the corresponding independent-noise estimator under mild assumptions.

4 Experiments

RTDMD achieves state-of-the-art few-step results on SD3-M and scales effectively to FLUX.2 4B, while ablations show gains from AC-DMD, SubGRPO, and the hybrid policy gradient.

  • Comparison with baselines: With 4-step SD3-M generation, RTDMD achieves the best performance across all five evaluation metrics, including CLIPScore 0.3161, PickScore 22.86, and HPSv2 0.3211.It establishes a new state of the art for few-step generation against prior methods including Rdm and GDMD.
  • Scaling to more advanced models: On FLUX.2 4B, RTDMD obtains new best results on seven of nine metrics and surpasses the larger FLUX.2 9B 50-step baseline on most metrics.Its relative improvement over its own baseline is more pronounced on OCR (+0.0483 versus +0.0126) and comparable on GenEval2.
  • Effect of AC-DMD: AC-DMD consistently improves performance across tested γ values, with γ = 0.01 yielding an additional +1.58 PickScore and +0.027 HPSv2 over A-DMD.The consistency loss improves fake score estimation by enforcing coherent predictions across neighboring timesteps.
  • Effect of SubGRPO: SubGRPO consistently improves all metrics over naive GRPO by sharing noise at non-selected steps and isolating reward variation to the selected subset.Increasing the subset size from M = 1 to M = 2 raises PickScore from 23.33 to 23.45.
  • Effect of hybrid policy gradient: The hybrid policy gradient combines SubGRPO for stochastic intermediate steps with a pathwise gradient through the final deterministic transition, consistently improving performance.This addresses the omission of the final deterministic transition in existing GRPO-based diffusion methods.

5 Conclusions and Limitations · Appendix A. Related Work

RTDMD is presented as a two-stage few-step image-generation framework that combines distribution matching distillation with reward-guided reinforcement learning. The related-work discussion situates it among trajectory-based and distribution-based distillation methods and reinforcement-learning approaches for diffusion and few-step generative models.

  • 5 Conclusions and Limitations: RTDMD unifies distribution matching distillation with reward-guided reinforcement learning in a two-stage framework for few-step image generation.The first stage is AC-DMD; the second stage optimizes against a reward-tilted teacher distribution.
  • 5 Conclusions and Limitations: AC-DMD uses ambient distribution matching and a consistency regularizer to stabilize fake-score estimation.The regularizer augments fake score training in the first stage.
  • 5 Conclusions and Limitations: The second stage minimizes KL divergence to a reward-tilted teacher distribution, decomposing the objective into distribution matching and reward maximization.This decomposition motivates jointly optimizing both terms in the second stage.
  • 5 Conclusions and Limitations: For reward maximization, RTDMD derives a hybrid policy gradient combining SubGRPO for stochastic intermediate steps with direct reward backpropagation.The supplied passage identifies SubGRPO as operating on stochastic intermediate steps; the remaining mechanism is stated as direct reward backpropagation.
  • Appendix A. Related Work: Few-step distillation methods broadly comprise trajectory-based approaches and distribution-based approaches that align outputs through GAN-based or score/VSD variants.Trajectory-based examples include progressive distillation and consistency models, while DMD is identified as foundational work among distribution-based methods.
  • Appendix A. Related Work: Reinforcement learning for diffusion models includes direct reward-gradient backpropagation and multi-step decision-process formulations for preference alignment.The passage cites ReFL and DRaFT for reward gradients, and DDPO and DPOK for multi-step decision processes.
  • Appendix A. Related Work: Reinforcement learning for few-step generative models combines few-step distillation with preference alignment, including post-training approaches such as Hyper-SD and PSO.Hyper-SD adds human feedback learning to trajectory-segmented consistency distillation, while PSO fine-tunes timestep-distilled models with pairwise preferences.

Appendix B. Derivation of the Step-k Training Objective … Appendix G. Hybrid Policy Gradient

The appendices derive step-specific score-matching objectives, unify stochastic and deterministic sampling through coefficient-preserving transitions, and establish optimality and self-consistency properties for the fake score. They also provide a practical unbiased estimator for the consistency regularizer.

  • Appendix B. Derivation of the Step-k Training Objective: For training step k, the student marginal is defined by diffusing the generated state ˆx_tk forward over the remaining interval [t_k, 1].This step-specific construction underlies the reverse-KL training objective and its score-difference gradient.
  • Appendix B. Derivation of the Step-k Training Objective: The fake score must be trained on the subinterval [t_k, 1] using the generated endpoint ˆx_tk, because clean samples x_0 are unavailable when t_k > 0.Reusing the full-interval objective with mismatched conditioning would generally produce a biased estimator.
  • Appendix C. Coefficient-Preserving Sampling Formula: Coefficient-preserving sampling exactly matches signal and noise coefficients to the scheduler at time t_k while exposing stochasticity through η ∈ [0, 1].η = 0 recovers deterministic Euler ODE sampling, whereas η = 1 recovers the consistency-model scheduler, making CPS a unified formulation of both strategies.
  • Appendix D. Optimality of Ambient Denoising Score Matching: The DSM objective has the true student marginal score as its unique optimal solution.The proof reduces optimization pointwise in time and shows the minimizing predictor is the conditional mean, which equals the marginal score.
  • Appendix E. Self-Consistency of the Optimal Fake Score: An optimal x-prediction fake score satisfies self-consistency across noise levels because Gaussian interpolation induces a Markov chain and the tower rule applies.The consistency regularizer penalizes deviations from this property, constraining the fake score as the generator distribution shifts during training.
  • Appendix F. Practical Consistency Loss Estimator: The intractable consistency expectation is estimated unbiasedly by drawing two independent reverse-step samples from p_ϕ(x_t′′ | x_t′).Experiments use adjacent schedule endpoints, setting t′ = t_k and t′′ = t_{k+1}.

G.1 Derivation … Appendix I. Detailed Algorithm of RTDMD

The appendices derive a hybrid policy gradient for few-step generators with stochastic intermediate transitions and a deterministic final mapping, explain its implementation with GRPO and SubGRPO, and present the RTDMD algorithm pipeline.

  • G.1 Derivation: The hybrid policy gradient separates contributions from K −1 stochastic transitions and the deterministic final step.Stochastic transitions use log-derivatives of transition densities, while the terminal contribution directly differentiates the deterministic mapping.
  • G.2 Discussion: The stochastic gradient term accounts for parameter effects on intermediate-state distributions through trajectory log-derivatives.The derivation isolates the random intermediate states and avoids applying log-derivative identities to a deterministic step represented by a Dirac delta.
  • G.2 Discussion: The deterministic gradient term differentiates only the terminal mapping’s explicit parameter dependence, without reintroducing intermediate-state dependence already captured by the stochastic term.Reintroducing that dependence would count the same contribution twice.
  • Appendix H. Background on GRPO and Its Application to Few-step Generators: GRPO reduces policy-gradient variance by sampling N responses, normalizing advantages within each group, and using the group mean reward as a baseline instead of a value network.The method was originally developed for language model alignment and uses a clipped objective.
  • Appendix H. Background on GRPO and Its Application to Few-step Generators: Naive GRPO applies group-normalized advantages to K-step Gaussian-transition trajectories, with independent noise sampled at each stochastic step.For each prompt, trajectories receive rewards and the stochastic gradient uses Gaussian log-likelihoods.
  • Appendix H. Background on GRPO and Its Application to Few-step Generators: Independent noise across all K −1 steps produces noisy credit assignment and high gradient variance because reward differences combine perturbations from every step.SubGRPO addresses this by sharing noise at non-selected steps to isolate contributions from a chosen subset.
  • Appendix I. Detailed Algorithm of RTDMD: Appendix I presents the detailed algorithm pipeline for RTDMD in Algorithm 1.The passage identifies the algorithm presentation but does not provide further procedural details.

Appendix J. More Implementation Details

Appendix J specifies the RTDMD implementation and two-stage training procedure, including the K-step schedule, optimizer settings, score-model and generator updates, and reward-optimization details.

  • Optimization and hardware: For K = 4, the pre-shift timestep schedule is [1.0, 0.75, 0.5, 0.25], and training uses AdamW with β1 = 0.9, β2 = 0.999, and a constant learning rate.For FLUX.2 4B, Stage 1 starts from the FLUX.2 [klein] 4B checkpoint; Stage 1 uses 8 NVIDIA H20 GPUs, while Stage 2 uses 16 for FLUX.2 4B and 8 for the SD series.
  • Algorithm requirements: RTDMD requires a teacher score sψ, generator Gθ, fake score sϕ, reward r(·), K-step schedule {t_k}^K, and reward weight β.The algorithm ensures a reward-optimized K-step generator Gθ.
  • Stage I: Ambient-Consistent Distribution Matching Distillation: Stage I rolls out Gθ, samples a step index k, constructs a noised state x^(k)_t for t ∈ [t_k, 1], then updates sϕ and Gθ using Eqs. (12) and (9).This stage is labeled Ambient-Consistent Distribution Matching Distillation (cold start).
  • Stage II: Reinforcing the Few-step Generator: Stage II generates N trajectories with shared or independent noise, computes group-normalized advantages A_i, updates sϕ, and updates Gθ along ∇θL_total from Eq. (18).The stage is titled Reinforcing the Few-step Generator.

Appendix K. Discussion on the Sampling Schedule and Connection to DMD2

A-DMD performs comparably across sufficiently stochastic sampling schedules and is set to η=0.9. Its realized-state conditioning yields a scheduler-agnostic re-derivation that matches DMD2 in the CM regime without sacrificing quality.

  • Sampling schedule: η≥0.8 yields comparable performance, so the experiments use η=0.9 to provide sufficient stochasticity for diverse GRPO trajectories.The authors attribute this insensitivity to the need for adequate stochasticity in intermediate transitions for effective exploration.
  • Connection to DMD2: With η=1.0, A-DMD and DMD2 differ in training intervals: A-DMD re-noises the generated intermediate state, whereas DMD2 re-noises the predicted clean sample across [0, 1].A-DMD conditions score matching on the realized state and its subsequent interval, while DMD2 uses the generator’s x-prediction as x0.
  • Connection to DMD2: A-DMD is more general because predicted clean samples are valid under CM consistency but not under non-CM schedulers such as Euler ODE.Under non-CM schedules, treating the prediction as x0 would introduce scheduler-dependent bias in score matching.
  • Connection to DMD2: Conditioning on the realized noisy state and matching scores on [t_k, 1] makes A-DMD valid for any η∈[0, 1] and scheduler choice.The realized state is constructed as a sample from the student marginal p_θ^(k), avoiding the clean-sample assumption.
  • Connection to DMD2: A-DMD with η=1.0 converges to nearly identical final performance as standard DMD2, confirming no quality loss in the CM regime.The authors frame the re-derivation primarily as a unified, scheduler-agnostic extension rather than a quality-improving change in the CM setting.

Appendix L. GenEval Results

On the non-differentiable GenEval compositional generation benchmark, RTDMD achieves an overall score of 0.94 with only 4 NFE on SD3.5-M. Optimization relies on AC-DMD and SubGRPO, nearly matching Flow-GRPO’s 0.95 score despite its 40 NFE full multi-step model with CFG.

  • Appendix L. GenEval Results: GenEval evaluates compositional generation on SD3.5-M with a non-differentiable reward signal.Because the reward is non-differentiable, the hybrid policy gradient cannot backpropagate through the final deterministic step.
  • Appendix L. GenEval Results: Optimization relies solely on AC-DMD and SubGRPO for stochastic transitions.The final deterministic step receives no reward-gradient backpropagation under GenEval.
  • Appendix L. GenEval Results: 0.94 overall score with only 4 NFE nearly matches Flow-GRPO’s 0.95, despite Flow-GRPO using 40 NFE, a full multi-step model, and CFG.The comparison is reported for SD3.5-M against Flow-GRPO.

Appendix M. Qualitative comparison for Ablation

Appendix M presents qualitative visualizations for the ablation studies, with Figs. 6–8 covering distillation and reinforcement-learning configurations. The samples use randomly drawn prompts and noise seeds without cherry-picking, providing visual evidence that complements the main-text quantitative analysis.

  • Evaluation protocol: The visualizations are not cherry-picked and use randomly drawn prompts and noise seeds, complementing the quantitative ablation results.Figs. 6–8 provide samples for the configurations listed in Tabs. 3–4.
  • Cold-start training: Figure 6 compares AC-DMD variants across γ = 0.001, 0.005, 0.01, and 0.1 against A-DMD after cold-start training.The columns follow the configurations in Tab. 3, listed bottom-to-top.
  • Two-stage training: Figure 7 compares the same AC-DMD and A-DMD configurations after two-stage training.The columns correspond to the rows in Tab. 3 in bottom-to-top order.
  • Reinforcement-learning ablation: Figure 8 compares RTDMD configurations with M = 2, M = 2 w/o Ldet, M = 1, and w/o Ldet against GRPO and ∅.The columns correspond to the rows in Tab. 4 in bottom-to-top order.

Appendix N. More Qualitative Results

Appendix qualitative comparisons show that RTDMD produces superior visual quality and prompt adherence across diverse prompts. With identical noise inputs, it outperforms other methods in SD3-M and few-step diffusion comparisons, and generates FLUX.2 4B images in 4 NFE without classifier-free guidance.

  • Overall qualitative results: RTDMD consistently produces superior visual quality and prompt adherence compared to baselines across diverse prompts.Additional visual comparisons are provided in Figures 9–11.
  • SD3-M comparison: Using identical noise inputs, RTDMD outperforms other methods in both quality and prompt alignment for SD3-M.This comparison is reported in Figure 9 for SD3-M (Esser et al., 2024a).
  • Few-step diffusion comparison: Using identical noise inputs, RTDMD outperforms other few-step diffusion models in both quality and prompt alignment at 4 NFE.The comparison covers few-step diffusion models under a 4 NFE setting.
  • FLUX.2 4B generations: RTDMD produces visual generations on FLUX.2 4B under 4 NFE without applying classifier-free guidance.These generations are shown in Figure 11.
Loading 2605.26108v3…