🤖 machine learning

Diffusion Fine-tuning with Rewarded Moment Matching Distillation

This paper introduces Rewarded Moment Matching Distillation (RMMD), a novel framework that unifies diffusion model distillation and reinforcement learning to achieve superior speed-quality trade-offs, demonstrated by significant improvements on both ImageNet and the high-dimensional GenCast weather forecasting model.

Original authors: Alexis Jacq, Guillaume Couairon, Valentin De Bortoli, Quentin Berthet, Arnaud Doucet, Romuald Elie

Published 2026-06-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alexis Jacq, Guillaume Couairon, Valentin De Bortoli, Quentin Berthet, Arnaud Doucet, Romuald Elie

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a master chef (the Teacher Model) who can cook a perfect, complex dish, but it takes them 59 hours to do it. They are slow, but the food is delicious and accurate. You want a sous-chef (the Student Model) who can cook the same dish in just a few minutes, but you also want them to tweak the recipe slightly to make it spicier or healthier (optimizing for a Reward).

The problem is that if you just tell the sous-chef to "make it spicier" while they are still learning the basics, they might ruin the dish entirely or forget how to cook it properly.

This paper introduces a new training method called Rewarded Moment Matching Distillation (RMMD). Think of it as a two-phase training camp that solves this problem.

Phase 1: The "Shadowing" Phase (Distillation)

First, the sous-chef spends time shadowing the master chef. They don't just watch the final dish; they watch every single step of the cooking process.

The Analogy: Imagine the chef is peeling an onion layer by layer. The student learns to predict what the onion looks like at every stage of peeling, not just the final result.
The Result: The student learns to cook a nearly perfect version of the dish in just 8 steps (instead of 59), keeping the "naturalness" and quality of the original. This is called Moment Matching.

Phase 2: The "Taste-Test" Phase (Reward Fine-Tuning)

Now that the student is a skilled cook, you want them to adjust the flavor (e.g., make it redder, or in the paper's case, make weather forecasts more accurate).

The Problem with Old Methods: Previous methods tried to teach the student to change the flavor by looking at a "re-noised" version of a dish they didn't actually cook themselves (off-policy). It's like telling a chef to improve a recipe based on a photo of a dish someone else made. The student gets confused because the ingredients don't match what they usually use.
The RMMD Solution: The paper uses an On-Policy approach. The student cooks a dish, then the teacher adds a little bit of "noise" (like sprinkling random spices) to it, and the student tries to fix it again.
The Magic Trick: While the student is trying to fix the dish to make it "spicier" (maximize the reward), the system also whispers, "Hey, don't forget how to peel the onion properly!" It uses the "Shadowing" lessons from Phase 1 as a safety net. This prevents the student from going crazy and ruining the dish just to get the reward.

Why is this special?

It's a Balancing Act: The paper shows that you can make the dish faster and better at the same time. Old methods usually forced you to choose: either keep the quality but stay slow, or get fast but ruin the taste. RMMD finds the sweet spot.
It Works on Real Science: The authors didn't just test this on pictures of cats and dogs. They applied it to GenCast, a super-complex weather forecasting model.
- The Teacher: Takes 59 steps to predict the weather for the next 12 hours.
- The RMMD Student: Takes only 8 steps (making it 7.5 times faster).
- The Result: The fast student didn't just get faster; it actually predicted the weather better than the slow master chef for 93% of the variables (like temperature and wind). It also fixed a common problem where weather models are too confident (under-dispersed), making the forecasts more reliable.

The Bottom Line

The paper claims that RMMD is a smarter way to train AI models. It teaches them to be fast without forgetting how to be accurate, and it teaches them to follow new goals (like "be redder" or "be more accurate") without breaking the fundamental rules of how they generate data.

In short: It's like training a race car driver who can drive 200 mph (fast) but still knows exactly how to stay on the track (accurate) and can navigate a new route (reward) without crashing.

Technical Summary: Diffusion Fine-tuning with Rewarded Moment Matching Distillation

Problem Statement

Diffusion models have become the standard for high-fidelity image synthesis and scientific forecasting due to their stable training objectives. However, practical deployment is hindered by the high computational cost of inference, which requires dozens to hundreds of denoising steps. While distillation techniques compress these models into few-step generators, and Reinforcement Learning (RL) fine-tuning aligns them with specific reward functions (e.g., human preferences or scientific metrics), combining these two phases remains challenging.

Existing approaches to merge distillation and reward optimization face significant limitations:

Fragility in Joint Training: Merging distillation and reward fine-tuning into a single phase often causes the generated samples to drift out of the teacher model's input distribution, invalidating the distillation signal.
Memory and Bias Issues: Fine-tuning an already distilled model via multi-step backpropagation (e.g., DRaFT) is memory-intensive. Truncating backpropagation (e.g., ReFL) introduces bias by evaluating rewards on noisy latents.
Structural Constraints: Methods like HyperNoise avoid chain backpropagation by perturbing initial noise but are structurally confined to low-frequency image changes, limiting their ability to optimize complex, high-frequency structural rewards.

Methodology: Rewarded Moment Matching Distillation (RMMD)

The authors propose RMMD, a two-phase framework that principledy connects distillation and reward fine-tuning while preserving the high-fidelity "naturalness" of advanced distillation.

Phase 1: Moment-Matching Distillation (MMD)

The process begins by distilling a base teacher model using Moment-Matching Distillation (MMD). This method matches intermediate denoising distributions along the sampling trajectory, yielding a multi-step student model that closely tracks the teacher's marginals. This student is then frozen to serve as a stable distributional reference.

Phase 2: On-Policy Reward Fine-tuning

In the second phase, the student model is fine-tuned to maximize a reward function while maintaining distributional fidelity. Key innovations include:

On-Policy Sampling: Instead of using off-policy data points, the method samples noisy states $x^{pol}_t$ from the current student's policy (by running the student for $K_{student}$ steps, re-noising the output, and stopping the gradient). This ensures the reward objective directly approximates the student's expected reward and prevents distributional drift from invalidating the training signal.
Recycled Regularization: The core contribution is reusing the Moment-Matching loss as a regularizer during fine-tuning. Rather than a heuristic penalty, this loss explicitly penalizes the discrepancy between the student's intermediate denoising distributions and those of the frozen reference model. This provides an interpretable knob (the regularization weight) to trade off reward maximization against distributional fidelity.
Single-Step Gradient: The method captures the full frequency range without multi-step unrolling by taking a single denoising step on corrupted on-policy samples and evaluating the reward.
L2 Regularization: An additional L2 term penalizes significant deviations from the initial MMD-distilled model's predictions to further stabilize training.

The total objective combines the on-policy reward gradient with the moment-matching regularizer and the L2 term:
$\nabla_\theta \mathcal{L}_{RMMD} = \mathbb{E} \left[ \left(\frac{\partial \hat{x}_0}{\partial \theta}\right)^\top \left( m_{student} - m_{teacher} - \lambda \nabla_{\hat{x}_0} R(\hat{x}_0) \right) \right] + \lambda_{reg} \mathcal{L}_{L2}$

Key Contributions

RMMD Procedure: A novel method for diffusion distillation that simultaneously maximizes a reward function, utilizing an on-policy version of moment-matching loss for regularization.
Empirical Superiority: Demonstrated that RMMD consistently outperforms single-step baselines (DI++) and multi-step competitors (DRaFT, HyperNoise) on FID-Reward Pareto fronts across diverse reward functions, ranging from pixel-level metrics to semantic alignment (CLIP).
Scientific Application: Successfully applied RMMD to GenCast, a state-of-the-art diffusion-based weather forecasting model, optimizing the Continuous Ranked Probability Score (CRPS).

Results

ImageNet Experiments

Multi-Step Advantage: Using an 8-step MMD distilled student as a starting point significantly outperformed 1-step baselines (DI++). The 8-step RMMD achieved superior trade-offs on FID-Reward Pareto fronts.
Comparison: RMMD generated better Pareto fronts than DRaFT and HyperNoise across various resolutions (64x64 and 512x512) and reward types (CLIP alignment, Inception Score, smoothness).
Qualitative Analysis: Visualizations showed that while DRaFT-1 preserved content but introduced adversarial artifacts, and DRaFT-2 suffered from reward hacking (drastic distribution shifts), RMMD successfully integrated subtle modifications to increase rewards without significantly deviating from the original data distribution.

GenCast Weather Forecasting

Speedup: The distilled model achieved a 7.5× speedup (8 steps vs. 59 function evaluations) compared to the teacher.
Accuracy: The RMMD model optimized with CRPS outperformed the teacher model on 93% of target weather variables (compared to 75% for the best MMD-only baseline).
Calibration: RMMD addressed the under-dispersion issue common in MMD distillation. While the on-policy variant showed slightly worse dispersion metrics than the off-policy variant, it provided the best overall CRPS improvements.
Generalization: Optimizing the CRPS at a 12-hour lead time resulted in consistent improvements for forecasts up to 7 days, suggesting the method improves the underlying distributional modeling rather than just overfitting to the immediate reward.

Significance and Claims

The paper claims that RMMD provides a principled solution to the "drift" problem inherent in joint distillation and fine-tuning. By recycling the moment-matching loss as a regularizer and utilizing on-policy sampling, the method allows for the optimization of complex, differentiable objectives (like CRPS in weather forecasting) without sacrificing the generative quality or stability of the distilled model.

The authors emphasize that this approach scales to complex, high-dimensional scientific domains. The successful application to GenCast demonstrates that diffusion models can be both significantly accelerated and improved in their probabilistic predictions, paving the way for more efficient and accurate diffusion-based models in scientific computing where differentiable objectives exist. The work suggests that maintaining a stable distributional reference via moment matching is crucial for effective reward fine-tuning in diffusion models.