--- title: Reward-Tilted Distillation (RTDMD) url: https://www.emergentmind.com/topics/reward-tilted-distribution-matching-distillation-rtdmd type: topic --- # Reward-Tilted Distillation (RTDMD) Searching arXiv for the specified RTDMD paper and closely related reward-tilted/distillation works to ground the article in current literature. arXiv search query: 2605.26108 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 [2605.26108]. 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 [2605.26108]. ## 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 [2605.26108]. 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 [2605.26108]. 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” [2512.04678]. 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 [2605.26108]. ## 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, \[ D(p_\theta\|\tilde p_\psi) = D(p_\theta\|p_\psi) -\beta \mathbb{E}_{x\sim p_\theta}[r(x)] +\log Z. \] Since \(Z\) does not depend on \(\theta\), the gradient decomposes as \[ \nabla_\theta D(p_\theta\|\tilde p_\psi) = \underbrace{\nabla_\theta D(p_\theta\|p_\psi)}_{\text{distribution matching}} -\beta \underbrace{\nabla_\theta \mathbb{E}_{x\sim p_\theta}[r(x)]}_{\text{reward maximization}}. \] The paper’s key claim is that training against the reward-tilted teacher is therefore equivalent to jointly performing distribution matching and reward maximization [2605.26108]. 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 \(D(p_\theta\|p_\psi)\). 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 \(\exp(r/\beta)/Z(\boldsymbol{c})\) so that high-reward video samples dominate the distillation gradient [2512.04678]. 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 [2604.19009]. 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 [2605.26108]. 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 \(k\) is a noisy latent \(\hat x_{t_k}\) rather than a clean image. Accordingly, distribution matching is re-derived over the remaining interval \([t_k,1]\), conditioned on the actual intermediate latent. For step \(k\), the re-noising process is defined as \[ x_t^{(k)}=\alpha_k(t)\hat x_{t_k}+\sigma_k(t)\epsilon, \] with \[ \alpha_k(t)=\frac{1-t}{1-t_k}, \qquad \sigma_k(t)=\frac{t-t_k}{1-t_k}, \qquad \epsilon\sim\mathcal N(0,I). \] The generator then minimizes the reverse KL on that subinterval: \[ \mathcal L_{\mathrm{gen}}^{(k)}(\theta) = \mathbb E_{t\sim \mathcal U[t_k,1]} \bigl[ \lambda_k(t)\,\mathrm{KL}(p_{\theta,t}^{(k)}\|p_{\psi,t}) \bigr]. \] The resulting generator gradient is \[ \nabla_\theta \mathcal L_{\mathrm{gen}}^{(k)} \approx - \mathbb E\!\left[ \lambda_k(t)\alpha_k(t) \bigl(s_\psi(x_t^{(k)},t)-s_\phi(x_t^{(k)},t)\bigr)^\top \frac{\partial \hat x_{t_k}}{\partial \theta} \right]. \] Here \(s_\psi\) is the teacher score and \(s_\phi\) is the fake or student score estimator [2605.26108]. The fake score model is trained on the same subinterval by denoising score matching: \[ \mathcal L_{\mathrm{fake}}^{(k)}(\phi) = \mathbb E_{t\sim\mathcal U[t_k,1],\,\hat x_{t_k},\,x_t^{(k)}} \left[ \omega_k(t)\left\| s_\phi(x_t^{(k)},t) - \nabla_x\log q_t^{(k)}(x_t^{(k)}\mid \hat x_{t_k}) \right\|_2^2 \right], \] where \[ \nabla_x\log q_t^{(k)}(x_t^{(k)}\mid \hat x_{t_k}) = \frac{\alpha_k(t)\hat x_{t_k}-x_t^{(k)}}{\sigma_k(t)^2}. \] The paper states that this objective is unbiased for the true student marginal score [2605.26108]. AC-DMD adds a consistency regularizer because the student distribution shifts during training while \(s_\phi\) must track it under limited updates. If the fake denoiser is optimal, its \(x\)-prediction should satisfy \[ \hat f_\phi(x_{t'},t') = \mathbb E_{\tilde x_{t''}\sim p_\phi(x_{t''}\mid x_{t'})} \left[ \hat f_\phi(\tilde x_{t''},t'') \right], \qquad t''