Diversify Diffusion with Temperature Sampling and Variance-Corrective Time Shifting
2026-07-12 • Computer Vision and Pattern Recognition
Computer Vision and Pattern RecognitionMachine Learning
AI summaryⓘ
The authors explain that diffusion models tend to reproduce popular examples well but struggle with rare or unusual ones due to imbalances in their training data. They propose a new method called variance-corrective time shifting, which adjusts how the model samples at different times during inference. This fix helps increase diversity by making rare examples more accessible without hurting the overall quality or accuracy of the samples. Their approach works with existing models and lets users smoothly control diversity at different stages of the sampling process.
diffusion modelssampling temperaturescore scalingvariance inflationtime shiftingreverse diffusionmode reweightingsample diversitypretrained modelsinference
Authors
Peizhuo Li, Emre Aksan, Alexandru-Eugen Ichim, Thabo Beeler, Olga Sorkine-Hornung
Abstract
Diffusion models faithfully reproduce their training distribution, but also inherit its imbalances and leave rare or under-represented modes hard to reach. A natural inference-time remedy is to sample from the high-temperature target $p^{(γ)}_0(x) \propto p_0(x)^γ$ for $0 < γ< 1$, which flattens dominant modes and lifts rare ones. However, naive score scaling while correctly reweighting modes also inflates the per-mode variance, breaking the reverse diffusion process and degrading sample quality. We introduce variance-corrective time shifting, a training-free fix that queries the network at a shifted timestep and scales the resulting score by $γ$, canceling the variance inflation while preserving the mode reweighting. The correction turns simple temperature sampling into a practical diversity knob for pretrained diffusion and flow-matching backbones with no retraining, and we demonstrate consistent gains at minimal cost to sample quality and condition fidelity across DiT, Stable Diffusion and Motion Diffusion models. We further show that the timing of the temperature intervention enables coarse-to-fine control: high-noise stages drive compositional diversity across modes, while low-noise stages drive local appearance variation under a fixed composition.