Reward-Tilted Distillation (RTDMD)

• The paper introduces a two-stage RTDMD framework that reweights a pretrained teacher distribution with rewards to enhance few-step generative models.
• RTDMD employs an AC-DMD cold start to stabilize training followed by a hybrid RL stage that jointly maximizes reward and matches the teacher distribution via a KL objective.
• Experimental results on SD3-M, SD3.5-M, and FLUX.2 models show that RTDMD outperforms traditional distillation and RL-only techniques in key metrics like CLIPScore and GenEval.

Searching arXiv for the specified RTDMD paper and closely related reward-tilted/distillation works to ground the article in current literature. arXiv search query: (Huang et al., 25 May 2026) RTDMD reward-tilted distribution matching distillation Reward-Tilted Distribution Matching Distillation (RTDMD) is a two-stage training framework for few-step diffusion and flow generators that combines distribution matching distillation with reward-guided reinforcement learning by defining a reward-tilted teacher distribution and training the student to match it (Huang et al., 25 May 2026). In this formulation, a pretrained teacher distribution $p_\psi(x)$ is reweighted by a scalar reward $r(x)$, so that high-reward samples receive larger mass while the support and structural prior of the teacher are retained. RTDMD was introduced to address two coupled limitations of prior few-step distillation: standard distillation is not reward-aware, and direct reward optimization is difficult in generators whose intermediate steps are stochastic while the final step is deterministic (Huang et al., 25 May 2026).

1. Formal objective and motivation

RTDMD defines the reward-tilted teacher as

$$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$

where $r(x)$ is a scalar reward, $\beta \ge 0$ controls how strongly reward reweights the teacher, and $Z$ is the normalization constant (Huang et al., 25 May 2026). The student generator $p_\theta$ is then trained by minimizing

$D(p_\theta \,\|\, \tilde p_\psi).$

The motivation is explicit. Classical DMD-style methods distill a teacher distribution as-is, but the teacher is not necessarily aligned with human preferences, so a faithful student can inherit both desirable and undesirable modes. At the same time, reward optimization in few-step generators is structurally awkward because intermediate transitions inject noise while the terminal mapping is deterministic; methods that optimize only one part of this process leave a portion of the reward gradient unused (Huang et al., 25 May 2026).

This reward-tilted perspective places RTDMD within a broader family of distillation methods that bias learning toward preferred regions of the teacher distribution. In streaming video generation, for example, Reward Forcing introduces Rewarded Distribution Matching Distillation (Re-DMD), which biases distribution matching toward high-motion samples using a vision-language reward and states that “Vanilla distribution matching treats every training sample equally, limiting the model’s ability to prioritize dynamic content” (Lu et al., 4 Dec 2025). RTDMD generalizes the same core principle to few-step image generation, but does so through an explicit KL objective against a reward-tilted teacher distribution (Huang et al., 25 May 2026).

2. KL decomposition and the meaning of the reward tilt

The central derivation in RTDMD follows directly from the definition of $\tilde p_\psi$: $$\log \tilde p_\psi(x)=\log p_\psi(x)+\beta r(x)-\log Z.$$ Therefore,

$r(x)$0

Since $r(x)$1 does not depend on $r(x)$2, the gradient decomposes as

$r(x)$3

The paper’s key claim is that training against the reward-tilted teacher is therefore equivalent to jointly performing distribution matching and reward maximization (Huang et al., 25 May 2026).

This decomposition clarifies what RTDMD is and is not. The method does not replace teacher distillation with an unconstrained RL objective; instead, it modifies the target distribution so that reward optimization is embedded inside a distribution-preserving objective. A plausible implication is that the reward term can improve alignment without discarding the teacher’s image prior, because the optimization is explicitly anchored by $r(x)$4.

Closely related formulations in later work preserve the same intuition while changing where reward enters the computation. Re-DMD in Reward Forcing uses normalized exponential sample weights of the form $r(x)$5 so that high-reward video samples dominate the distillation gradient (Lu et al., 4 Dec 2025). GDMD instead moves the reward from sample space to gradient space, evaluating an implicit target induced by the DMD update rather than the raw output, with the stated goal of synchronizing RL and distillation and avoiding optimization divergence (Dong et al., 21 Apr 2026). These variants differ mechanistically, but all treat reward as a way to tilt distribution matching rather than as an independent post hoc objective.

3. Stage I: Ambient-Consistent Distribution Matching Distillation

The first RTDMD stage is Ambient-Consistent Distribution Matching Distillation (AC-DMD), a cold-start distillation phase designed to produce a stable few-step student before reinforcement learning begins (Huang et al., 25 May 2026). The paper argues that standard DMD is insufficient in this setting because, under coefficient-preserving sampling (CPS) with stochastic intermediate transitions, the state at step $r(x)$6 is a noisy latent $r(x)$7 rather than a clean image. Accordingly, distribution matching is re-derived over the remaining interval $r(x)$8, conditioned on the actual intermediate latent.

For step $r(x)$9, the re-noising process is defined as

$$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$0

with

$$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$1

The generator then minimizes the reverse KL on that subinterval: $$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$2 The resulting generator gradient is

$$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$3

Here $$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$4 is the teacher score and $$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$5 is the fake or student score estimator (Huang et al., 25 May 2026).

The fake score model is trained on the same subinterval by denoising score matching: $$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$6 where

$$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$7

The paper states that this objective is unbiased for the true student marginal score (Huang et al., 25 May 2026).

AC-DMD adds a consistency regularizer because the student distribution shifts during training while $$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$8 must track it under limited updates. If the fake denoiser is optimal, its $$\tilde p_\psi(x)=\frac{p_\psi(x)\exp(\beta r(x))}{Z},$$9-prediction should satisfy

$r(x)$0

The associated penalty is

$r(x)$1

with a two-sample estimator used in practice. The final fake-score objective is

$r(x)$2

This cold-start stage is what the paper designates as AC-DMD (Huang et al., 25 May 2026).

4. Stage II: joint reinforcement, hybrid policy gradient, and SubGRPO

After cold start, RTDMD enters a reinforcement stage that jointly optimizes the distribution-matching and reward terms from the KL decomposition. The total generator update is

$r(x)$3

(Huang et al., 25 May 2026).

The few-step generator is treated as a policy over the latent trajectory

$r(x)$4

For the first $r(x)$5 steps, CPS defines Gaussian transitions

$r(x)$6

while the final step is deterministic: $r(x)$7 The reward gradient therefore decomposes into a stochastic REINFORCE-style part and a deterministic backpropagation part: $r(x)$8 This hybrid estimator is one of RTDMD’s principal differences from methods that optimize only stochastic denoising steps or only the terminal deterministic mapping (Huang et al., 25 May 2026).

To reduce variance in the stochastic part, RTDMD introduces step-subset GRPO (SubGRPO). A subset $r(x)$9 of size $\beta \ge 0$0 is chosen, and only those steps receive independent noise across trajectories in a group; all other stochastic steps share the same group noise. The corresponding gradient is

$\beta \ge 0$1

where the advantage is the group-normalized reward

$\beta \ge 0$2

The paper characterizes this as a Rao–Blackwellized variant of naive independent-noise GRPO, with lower variance at a fixed gradient budget (Huang et al., 25 May 2026).

5. Relationship to adjacent reward-aware distillation methods

RTDMD belongs to a rapidly expanding class of methods that combine teacher matching with reward guidance, but the point of contact between reward and distillation differs substantially across papers.

Method	Reward attachment point	Distinctive mechanism
RTDMD (Huang et al., 25 May 2026)	Reward-tilted teacher distribution	AC-DMD cold start; hybrid stochastic-plus-deterministic RL; SubGRPO
Re-DMD in Reward Forcing (Lu et al., 4 Dec 2025)	Sample-level exponential reweighting	Motion-quality reward from VideoAlign; tilt toward dynamic video samples
GDMD (Dong et al., 21 Apr 2026)	Gradient-induced implicit target	Reward evaluates the DMD update target rather than raw sample
Stream-R1 (Wu et al., 5 May 2026)	Rollout-level and element-level reweighting	Inter-Reliability and Intra-Perplexity with reward saliency
$\beta \ge 0$3 / GNDM (Fan et al., 30 Mar 2026)	Distillation reconceptualized as reward	Group-normalized distribution matching; PPO/GRPO-style clipping
RMMD (Jacq et al., 29 Jun 2026)	Reward fine-tuning with distillation regularizer	On-policy reward optimization plus moment-matching regularization

RTDMD differs from DMD and DMD2 because it explicitly targets $\beta \ge 0$4 rather than only distilling the teacher prior (Huang et al., 25 May 2026). It differs from RL-only methods because the generator remains anchored to the teacher distribution. It also differs from reward-weighted DMD variants in streaming video. Re-DMD, for instance, multiplies the usual DMD score-difference term by a normalized exponential reward weight and uses VideoAlign’s motion quality as the reward function with $\beta \ge 0$5, thereby emphasizing dynamic content while preserving data fidelity (Lu et al., 4 Dec 2025). Stream-R1 extends this idea beyond scalar rollout weighting by applying one shared reward model both to inter-rollout exponential weighting and to intra-rollout spatial and temporal saliency maps, producing a hierarchical reward tilt without architectural modification or additional inference cost (Wu et al., 5 May 2026).

GDMD occupies a different point in the design space. Instead of rewarding the generated sample, it interprets the DMD gradient as an implicit target tensor $\beta \ge 0$6, scores that target with a reward model, and uses the score as an adaptive weighting inside a joint DMD-plus-preference objective (Dong et al., 21 Apr 2026). The paper’s explicit claim is that this aligns RL with the distillation trajectory and removes the need for the cold-start phase required by DMDR-style training. By contrast, RTDMD retains a cold-start stage in the form of AC-DMD (Huang et al., 25 May 2026).

Another nearby formulation is $\beta \ge 0$7, which re-conceptualizes distribution matching itself as an RL reward and introduces Group Normalized Distribution Matching (GNDM). In that framework, the DMD signal is converted into a reward-like quantity $\beta \ge 0$8, normalized within groups, and optimized through a clipped GRPO-style objective. The stated benefit is a unified reward-centric view that supports adaptive weighting, multi-reward composition, and importance sampling (Fan et al., 30 Mar 2026). RMMD, finally, is not presented as RTDMD, but as Rewarded Moment Matching Distillation; it uses on-policy reward optimization with moment matching as a regularizer and is described as a reward-tilted distribution matching procedure in which the distillation loss functions as a proxy for a distribution-preserving constraint (Jacq et al., 29 Jun 2026).

A common misconception is that all reward-aware distillation methods are interchangeable. The literature indicates otherwise. Some methods tilt sample weights, some tilt rollout weights, some reward gradients or implicit targets, and some preserve fidelity through moment matching rather than KL-derived DMD terms. RTDMD is specifically the formulation based on a reward-tilted teacher distribution and a two-stage AC-DMD-plus-hybrid-RL training recipe (Huang et al., 25 May 2026).

6. Experimental regime, reported performance, and practical scope

RTDMD is evaluated on SD3-M, SD3.5-M, and FLUX.2 4B, with all models distilled to 4 inference steps (Huang et al., 25 May 2026). The reported implementation uses a student initialized from the pretrained teacher without CFG, LoRA on attention layers with rank $\beta \ge 0$9 and $Z$0, CPS sampling with $Z$1, a Stage I cold start of about $Z$2 iterations, a Stage II reinforcement phase of about $Z$3 iterations, group size 24, 48 groups per prompt for SD3 and SD3.5 and 64 for FLUX.2, AdamW, and bf16 precision (Huang et al., 25 May 2026). Training rewards are drawn from combinations of HPSv2, CLIPScore, PickScore, and GenEval, with FLUX.2 additionally using OCR Score, Aesthetic, GenEval2, ImageReward, and HPSv3 (Huang et al., 25 May 2026).

On 4-step SD3-M, RTDMD reports the best scores across all reported metrics:

• CLIPScore: 0.3161
• Aesthetic: 5.9642
• PickScore: 22.8593
• HPSv2: 0.3211
• ImageReward: 1.3024 and is described as outperforming DMD, DMD2, TDM, DMDR, GDMD, and $Z$4, while surpassing the 100-step teacher with CFG on some metrics (Huang et al., 25 May 2026).

On SD3.5-M, RTDMD reaches an overall GenEval score of 0.94, which the paper describes as competitive with much more expensive systems (Huang et al., 25 May 2026). On FLUX.2 4B, the reported results are:

• ImageReward: 1.3712
• CLIPScore: 0.3219
• Aesthetic: 5.7746
• PickScore: 23.9642
• HPSv2: 0.3516
• HPSv3: 15.5772
• GenEval: 0.9046 The paper emphasizes that the distilled 4B model surpasses the much larger FLUX.2 9B teacher on most metrics despite using only 4 inference steps (Huang et al., 25 May 2026).

The ablations are used to attribute gains to specific components. The paper reports that adding the consistency loss improves AC-DMD substantially, SubGRPO improves over naive GRPO, and adding the deterministic gradient improves over stochastic-only RL (Huang et al., 25 May 2026). These findings are consistent with the method’s conceptual structure: AC-DMD stabilizes cold start on noisy intermediate states, SubGRPO addresses variance in stochastic policy gradients, and direct reward backpropagation through the final deterministic step uses information that GRPO-style estimators alone would ignore.

The scope of RTDMD is therefore narrower and more precise than the generic phrase “reward-guided distillation” might imply. It is a reward-aware few-step distillation method in which KL minimization to a reward-tilted teacher yields a principled decomposition into teacher matching and reward maximization, instantiated through AC-DMD cold start and a hybrid reinforcement stage (Huang et al., 25 May 2026). Subsequent work suggests that this formulation is part of a broader methodological trend in which reward increasingly modifies the effective training distribution rather than acting only as an external auxiliary loss (Lu et al., 4 Dec 2025).