Optimizing DDPM Sampling with Shortcut Fine-Tuning
Ying Fan, Kangwook Lee
ICML 2023
cs.LG
2023-01-31
Fine-tune the sampler with reinforcement learning instead of imitating the reverse process. On CIFAR-10, 10-step sampling reaches FID 2.28, beating the 1000-step model's 3.03.
Generating one image with a DDPM takes around a thousand denoising steps, which is the model's main drawback at deployment. The standard way to cut the step count is almost always the same: reproduce the reverse denoising process as faithfully as possible with fewer steps, using better noise estimates, smarter sub-sampling schedules, or non-Gaussian noise. Whatever the tweak, the reverse process is treated as a fixed reference to imitate.
The authors reframe generation as a control problem: pure noise is the start state, a natural image is the goal, and the reverse process is really a demonstration trajectory (imitation learning) — useful, but rarely optimal. If a demonstration is not optimal, why cling to it? The question becomes whether the sampler can find a shorter shortcut on its own, by driving the generated distribution directly toward the data distribution instead of fitting each reverse step.
The yardstick is an integral probability metric (IPM), a family of distribution distances that includes Wasserstein, total variation, and MMD. The objective is simply to minimize the IPM between the generated distribution pθ0 and the data distribution q0, with no constraint that each step stay faithful to the reverse process.
For the gradient the paper offers two routes. The first, SFT, differentiates through the whole T-step sampling chain GAN-style, which inherits RNN-style troubles: vanishing and exploding gradients, and heavy memory. The second, SFT-PG, is the core contribution. Theorem 4.1 proves that gradient descent on the IPM is equivalent to policy gradient. The trick is to cast diffusion sampling as a finite-horizon MDP: each conditional pθt(xt|xt+1) is the policy, the state transition is the identity, and reward is paid only at the final step, equal to the critic's score on the output image. The gradient then uses the critic's value rather than its gradient, so a wider family of critics is admissible and long-chain vanishing or exploding gradients disappear. The price is the usual high variance of policy gradient, controlled with a baseline function V in the style of REINFORCE and GAE.
Two results keep training stable. Theorem 4.2 gives a surrogate of the IPM that guarantees monotonic improvement over several generator updates inside a trust region, in the spirit of TRPO; in practice this means a small learning rate with gradient-norm clipping. For the critic, the authors find they can reuse the baseline loss to regularize the critic's value (baseline regularization, B) instead of the WGAN-style gradient penalty (GP); B is looser, tracks changes in the generator more sensitively, and works better empirically.
On the toy swiss-roll set, a DDPM trained with T=10 is fine-tuned, and the 10-step fine-tuned sampler beats even the 1000-step model:
| Model | W2(p0,q0) x10-2 (lower is better) |
| DDPM, 10 steps | 8.29 |
| DDPM, 100 steps | 2.36 |
| DDPM, 1000 steps | 1.78 |
| SFT-PG (B), 10 steps | 0.64 |
The real weight is on image benchmarks. Starting from a pretrained 1000-step DDPM on CIFAR-10 (32x32) and CelebA (64x64), fixing FastDPM's 10-step schedule and fine-tuning only the mean network:
| Method (all 10 steps) | CIFAR-10 FID | CelebA FID |
| Naive DDPM sub-sampling | 34.76 | 36.69 |
| FastDPM | 29.43 | 28.98 |
| Analytic-DPM | 22.94 | 28.99 |
| SN-DDPM | 16.33 | 20.60 |
| SFT-PG (B) | 2.28 | 2.01 |
| 1000-step model (reference) | 3.03 | 3.26 |
With only 10 steps, fine-tuned FID not only leaves the strongest fast sampler of the time (SN-DDPM) far behind, it also sits below the 1000-step full model. Against deterministic DDIM-family samplers, SFT-PG at NFE=10 gives the best FID in the table (2.28); at NFE=8 it reaches 2.64, on par with progressive distillation at 2.57. On compute, progressive distillation takes about a day on 8 TPUv4 chips for CIFAR-10, while SFT-PG takes about 6 hours on 4 RTX 2080Ti.
This is the earliest known work to train a diffusion model with reinforcement learning. Set in 2023, the lasting contribution is the reframing. Stop treating the reverse process as a yardstick to reproduce; it is just a demonstration trajectory that may not be optimal, and once you optimize the resulting distribution directly the sampler can grow its own shortcut. That idea keeps getting rediscovered on the diffusion-plus-RL line, which is why the paper is being discussed again now.
For practitioners, the practical half is that on CIFAR-10-scale images, 10 steps buy quality beyond the 1000-step model, at lower compute than distillation. The sober half is that this is 32x32 imagery and 2023-scale model sizes, with no validation on large text-to-image. Porting it straight into a Stable Diffusion-class pipeline is premature.
The authors acknowledge that fine-tuning adjusts only the mean network μθ; the variance is fixed, and learning it is left to future work. Even though gradients are no longer tracked at sampling time, collecting training trajectories still requires T' full inference steps, so sampling remains slower than a one-shot GAN.
Several things are also not fully nailed down. Policy gradient is inherently high-variance; the baseline tames it but the result is sensitive to how well the baseline network fits. If the generated distribution starts too far from the data, training can settle into a sub-optimum. The critic is a small hand-built CNN, so signal quality depends on it rather than on a ready-made strong critic like a pretrained-feature MMD. Most importantly, all results stop at the CIFAR-10/CelebA tier, so generality is untested.