Distribution Matching Distillation Meets Reinforcement Learning¶

Conference: ECCV2026
Paper: ECCV Paper
Authors: Dengyang Jiang et al.
Area: Image Generation
Keywords: distribution matching distillation, few-step generation, reinforcement learning, reward tilting, preference alignment

TL;DR¶

DMDR couples few-step diffusion distillation and preference optimization in a two-stage training process, allowing 4-step students to outperform multi-step teachers on several overall metrics, although preference gains still reduce generation diversity.

Background & Motivation¶

High-quality text-to-image models typically require repeated denoising, with each step invoking a large network, so reducing sampling steps directly affects generation cost. Distribution Matching Distillation (DMD) does not require a student to reproduce the teacher's trajectory step by step; it brings the few-step student's output distribution closer to the teacher's distribution. This addresses speed but does not automatically address user preferences: some regions of the teacher distribution produce more attractive, prompt-aligned images than others. Matching the original distribution alone provides no explicit objective for favoring higher-quality samples. Meanwhile, reinforcement learning (RL) can use reward models for preference alignment, but increasing reward does not necessarily preserve semantic capability or distribution coverage.

Simply composing these procedures also creates problems: when distillation is followed by independent RL, the few-step reference model has already lost some modes of the teacher distribution. Further optimization around this narrower reference can concentrate probability on a small set of easily rewarded images, causing reward hacking and semantic degradation. Conversely, applying RL to a multi-step model before distillation incurs expensive sampling and may lose previously acquired preference gains during compression. Direct reward back-propagation from the first distillation update is also unstable, because an immature student's incoherent images make reward feedback unreliable for identifying useful directions.

The paper therefore separates how rewards are used from when reward gradients are used: early rewards only adjust distillation weights, while direct reward optimization begins after the student stabilizes. Throughout this process, DMD is not merely an initial compression step but also a continuing constraint that preserves the teacher's generative prior. The method does not eliminate training stages; both stages contain preference information, but that information enters optimization differently. Core Idea: stabilize few-step distillation with static reward weights and dynamic distributional overlap, then jointly optimize rewards and DMD, letting preferences determine learning priorities while the teacher distribution discourages reward hacking.

Method¶

Overall Architecture¶

Training takes text prompts and noise as inputs, with a few-step generator as the central trainable model and real-score, fake-score, and reward models as supporting components. Here, a distributional score is the gradient of log density with respect to the input, not the scalar preference score returned by a reward model. The real-score estimator represents the multi-step teacher's distribution, while the fake-score estimator learns the current student's output distribution; their difference provides the DMD update signal. The fake-score estimator changes with the student and should not be understood as another fixed teacher. The paper follows backward simulation: the student first takes several sampling steps to obtain an intermediate state, then predicts a clean image from that state, training on states encountered during actual few-step sampling.

Stage 1 combines reward-tilted distillation and dynamic distributional overlap; the latter includes DynaDG and DynaRS, neither of which is an additional inference module. Stage 2 switches to joint reward optimization while retaining DMD, directly optimizing preferences instead of relying only on sample reweighting. Deployment runs only the trained few-step generator; reward and score estimators supply training supervision and are not invoked for every deployed generation.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    A["Prompt + noise"] --> G["Few-step generator<br/>Full trajectories and intermediate states"]
    subgraph S1["Stage 1: coupled distillation"]
        direction TB
        B["Reward-tilted distillation<br/>Full-trajectory rewards as weights"]
        C["Dynamic distributional overlap<br/>DynaDG + DynaRS"]
    end
    G --> B
    G --> C
    R["Reward models"] -.-> B
    T["Real- and fake-score estimators"] -.-> C
    B --> W["Weighted DMD updates<br/>Obtain a stable student"]
    C --> W
    W --> D["Joint reward optimization<br/>Direct reward + DMD"]
    R -.-> D
    T -.-> D
    D --> O["Inference: few-step generator only"]

Key Designs¶

1. Reward-tilted distillation: prioritize learning before back-propagating rewards

Standard DMD updates the generator using the teacher-student score difference, but equal update strength does not distinguish samples according to human preference. Reward-Tilted Distribution Matching (RT-DM) adds a reward-dependent static weight to this distillation signal, giving preferred samples greater influence over the update. Static means that reward gradients are not back-propagated through the weight during that generator update, not that every sample receives one constant or that rewards remain unchanged throughout training. Rewards therefore identify which outputs deserve greater emphasis, while the teacher and fake-score estimators still determine how to change the generator. The paper interprets this as learning toward a reward-tilted teacher distribution instead of uniformly retaining the teacher's existing preferences.

Rewards are evaluated on images from the student's complete few-step trajectory, such as all 4 steps, rather than relying exclusively on a single clean-image estimate from an arbitrary intermediate state. During the cold start, a complete trajectory generally yields more coherent images than a one-step prediction from an intermediate state, making feedback more informative. Since rewards only serve as weights, this complete reward-evaluation trajectory does not require a retained back-propagation graph, decoupling more reliable feedback from an expensive gradient graph. This does not mean that the teacher supplies a correct target image for each sample or that the reward model becomes a second denoising teacher. The mechanism selects distribution-matching updates according to preference and remains limited by the teacher distribution and the student's current generation quality. In Figure 5, continuing RT-DM alone eventually reaches a plateau, motivating the subsequent transition to direct reward optimization.

2. Dynamic distributional overlap: improve both teacher directions and renoised states

RT-DM determines learning priorities, but when student outputs are far from image regions familiar to the teacher, the DMD score difference itself can be unreliable. Dynamic Distribution Guidance (DynaDG) inserts trainable LoRA modules into the real-score estimator, temporarily moving the early target distribution toward the student distribution. Rather than permanently weakening the teacher, this establishes useful overlap before asking an immature student to bridge a large distribution gap. As training progresses, the teacher's LoRA scale is annealed to zero, returning the real-score estimator to the original teacher distribution. Descriptions of a fixed teacher must therefore distinguish the original frozen capability from its temporarily active adaptation branch; the early real-score estimator is not entirely unchanged.

Dynamic Renoise Sampling (DynaRS) does not move the target but changes the noise levels at which DMD compares student and teacher distributions. Early training favors high-noise states, because adding sufficient noise makes both distributions more Gaussian-like and places student and teacher in more similar regions. Optimization at these states emphasizes global structure rather than prematurely correcting fine details in images whose structure has not formed. As the student improves, this sampling bias is annealed toward uniform sampling, covering lower-noise states to learn textures and other details. The renoise time is used for training supervision and does not imply additional denoising steps at deployment. The two strategies address what target is compared and at which noise level it is compared, jointly improving distribution matching during the cold start.

3. Joint reward optimization: turn DMD into continuing regularization

Once the student generates stable images, Stage 2 directly adds a reward optimization loss instead of merely assigning larger distillation weights to high-reward samples. ReFL can update the image generator through input gradients of a frozen reward model, whereas DPO and GRPO supply different preference or policy optimization signals. The paper evaluates these as interchangeable RL branches, not as three algorithms run simultaneously in every training job. The generator continues receiving DMD updates, and the fake-score estimator must still track the changing student distribution, so this stage is not RL alone after distillation has ended. Unlike using the distilled student as the only reference, continuing DMD constrains preference updates using the richer generative prior of the multi-step teacher.

The authors explain this interaction as a reward-ascent direction combined with a direction returning toward the teacher manifold: rewards favor high-scoring regions, while distribution discrepancies provide corrective feedback. A more cautious interpretation is empirically effective soft regularization, not a hard projection that rigorously keeps every update inside the teacher's support. Figure 4 shows that independent RL can increase training rewards while reducing DPG, whereas joint optimization better preserves prompt adherence. However, this constraint cannot guarantee preservation of every mode, as the later diversity experiments demonstrate directly. An additional advantage of default ReFL in the few-step setting is that each sampling state is trained to predict a clean image directly, making reward gradients more accessible to predictions from earlier states. This differs from multi-step ReFL's difficulty with optimizing mainly late denoising steps to conserve memory, but does not establish that every implementation back-propagates through the entire trajectory.

Loss & Training¶

Following the definitions of the two losses and their coefficient in Section 3.3, the Stage 2 objective can be summarized as:

\[ \mathcal{L}_{\mathrm{total}}=\mathcal{L}_{\mathrm{dmd}}+\lambda_{\mathrm{rl}}\mathcal{L}_{\mathrm{rl}}. \]

Here, \(\mathcal{L}_{\mathrm{rl}}\) implements reward maximization under a loss-minimization convention, and \(\lambda_{\mathrm{rl}}\) controls the strength of direct preference optimization. DMD provides generator update directions, while the fake-score estimator is trained separately with a diffusion loss on student samples; denoising models predict noise and flow models predict velocity. Stage 1 instead adjusts DMD updates using RT-DM reward weights together with the two dynamic strategies; it should not be described as the direct reward loss of Stage 2. Main experiments use prompts from t2i-2M and cover SDXL-Base, SD3-Medium, and SD3.5-Large. The default reward combination is DFN-CLIP and HPSv2.1, the default RL algorithm is ReFL, and each stage runs for 1.5K steps. The transition after 1000 iterations in the mechanism study of Section 3 is not the full stage budget used in the main experiments of Section 4. The default weight in Table 5(c) is \(\lambda_{\mathrm{rl}}=0.5\); increasing it does not guarantee continued gains in general capability.

Equations (1), (2), and (4) are visibly corrupted in the extracted source; the expression above only summarizes the weighted objective identifiable from the surrounding prose. In particular, Equation (4) involves both a reward-tilted distribution and an exponential weight, but the damaged text does not establish the exact normalization, so its full derivation is not reconstructed here. The cache contains the main paper and references but not the cited appendix, so this note does not infer detailed LoRA settings, annealing schedules, full pseudocode, or user-study procedures.

Key Experimental Results¶

Main Results¶

The following selection comes from Table 2, page 11; evaluation prompts are sampled from ShareGPT-4o-Image without overlap with training prompts, and all four scores are higher-is-better. NFE denotes the number of function evaluations and should not be treated as a proportional measurement of end-to-end latency.

Backbone	Method	NFE	CLIP Score	Aesthetic Score	Pick Score	HP Score
SDXL-Base	Base-Model	50	34.7588	5.6480	22.1085	27.1477
SDXL-Base	DMDR	4	35.4940	6.0324	22.7122	32.9832
SD3-Medium	Base-Model	50	34.9025	5.5942	22.1801	28.4021
SD3-Medium	DMD2	4	33.9421	5.6137	21.7688	27.3675
SD3-Medium	DMDR	4	35.0142	5.8876	22.4892	32.1145
SD3.5-Large	Base-Model	50	35.5509	5.7014	22.4856	28.8135
SD3.5-Large	DMDR	4	35.9757	6.1541	22.9072	32.7368

On SD3-Medium, DMDR raises HP Score from DMD2's 27.3675 to 32.1145 at the same 4 NFE; this DMD2 result was reproduced by the authors. Tables 3 and 4, page 12, additionally report an increase in SD3-Medium's overall GenEval score from the teacher's 0.62 to 0.65 and overall DPG from 84.08 to 85.30. However, GenEval counting drops from 0.63 to 0.57 and position from 0.34 to 0.27, so surpassing the teacher refers to aggregate results rather than dominance on every component. The paper states that GenEval and DPG metrics were not directly optimized with RL, making them more informative about preserved capabilities than the training reward alone.

Ablation Study¶

The following is Table 5(c), page 13; all runs use SD3-Medium with 4-step sampling and vary the Stage 2 reward weight, with higher values preferred for every reported metric.

Reward weight	CLIP Score	Aesthetic Score	Pick Score	HP Score	DPG overall
0.1	34.7422	5.8508	22.3064	31.6565	85.24
0.5	35.0142	5.8876	22.4892	32.1145	85.30
1.0	35.0087	5.84430	22.4709	32.8982	85.04
2.0	35.0322	5.7806	22.1098	32.7683	82.48

Increasing the weight from 0.5 to 2.0 retains a high HP Score but reduces DPG from 85.30 to 82.48, showing why reward alone is insufficient for model selection. In Table 5(b), ReFL, DPO, and GRPO obtain DPG scores of 85.30, 85.00, and 85.08 respectively, supporting compatibility without making ReFL the winner on every metric. In Table 5(a), CLIP + HPS achieves DPG 85.30 versus 85.04 for HPS alone, while HPS alone has a higher HP Score of 32.9065 versus 32.1145 for combined rewards. Combined rewards therefore improve the balance of metrics rather than maximizing each reward model's own score.

Key Findings¶

The following is a selection from Table 7, page 15, evaluating SD3-Medium diversity on ShareGPT-4o-Image; both HPS and LPIPS are marked higher-is-better in the paper. LPIPS serves here as a proxy for variation across generations, not as the quality score of an individual image.

Metric	Base-Model	DMD2	DMDR
HPS	28.4021	27.3675	32.1145
LPIPS	0.6840	0.5832	0.5548

DMDR has higher preference scores but lower LPIPS than both controls, so mitigating reward hacking with DMD does not eliminate the diversity cost. Figures 5 and 7 support the RT-DM warm start and dynamic strategies respectively; the cache provides no reliably transcribable curve values, so exact gains are not guessed from those curves. Table 1, page 7, measures rollout time for SD3-M at 1024×1024 with local batch size 16 on an H100: few-step joint training takes 0.43–1.21 seconds versus 9.24–10.57 seconds for multi-step RL. These are rollout timings, not evidence for total training cost or deployed service latency, and choosing interval endpoints does not establish a fixed speedup.

Highlights & Insights¶

Rewards need not always enter training as gradients. RT-DM uses more informative full-trajectory feedback to prioritize updates before the student stabilizes, offering a reusable approach to preference-aware cold starts.
DynaDG and DynaRS address different aspects of distribution mismatch. One temporarily moves the target, while the other changes the noise scale at which distributions are compared, providing a more targeted intervention than merely lowering the learning rate.
Distillation can remain useful after acceleration has been achieved. Retaining the teacher distribution as a constraint helps distinguish improved preferences from higher rewards obtained by losing prompt semantics.

Limitations & Future Work¶

The authors explicitly acknowledge the quality-diversity trade-off: preference optimization concentrates probability in high-reward regions, and DMD does not fully prevent loss of mode coverage.
Manifold projection is best treated as a mechanistic interpretation rather than a rigorous guarantee; DPG degradation at larger reward weights also shows that regularization strength needs tuning.
Table 8 reports HPS values of 30.4024 and 32.4786 at weights 0.1 and 1, differing from 31.6565 and 32.8982 in Table 5(c); the text does not clearly explain the setting differences, so these should not be merged into one directly reproducible curve.
Flash-DMD and Diff-Instruct++ are excluded from the main comparison because the authors could not obtain weights or training code sufficient for fair reproduction, leaving a quantitative comparison gap with the closest methods.
The appendix is absent from the cache, preventing verification of user-study statistics and all training hyperparameters; adaptive reward weights and diversity-preserving objectives remain future directions rather than validated solutions in this paper.

DMD / DMD2: These methods primarily match a few-step distribution to its teacher; DMDR goes beyond more accurate imitation by selecting among regions with different preference quality.
Hyper-SD / PSO: These represent preference optimization after distillation; DMDR introduces reward information earlier and continues distribution matching during direct reward optimization.
Flash-DMD / Diff-Instruct++: These already combine distillation and RL, so the contribution should not be described as the first addition of two losses; it lies in the RT-DM cold start and associated dynamic strategies.
ReFL / DPO / GRPO: These are interchangeable preference-optimization branches in this framework; the transferable idea is its training schedule, not a new general-purpose policy-gradient algorithm.

Rating¶

Novelty: 4/5. The main advances are the reward-tilted cold start and dynamic overlap strategies; related work already uses joint losses.
Experimental Thoroughness: 4/5. Multiple backbones, algorithms, and diversity costs are studied, but missing closest-method comparisons and the unavailable appendix limit verification.
Writing Quality: 3/5. The two-stage motivation is clear, but strong theoretical language, cross-table discrepancies, and extraction problems require care.
Value: 4/5. The framework offers practical guidance for preference alignment in few-step generation while showing why high rewards do not imply universal capability gains.