Diversify Diffusion with Temperature Sampling and Variance-Corrective Time Shifting

2026-07-12

AI Summary

Diffusion models often fail to capture rare modes in the training distribution, resulting in insufficient generation diversity. This work proposes a retraining-free variance-corrected time-shifting method combined with high-temperature sampling, which significantly enhances generation diversity while preserving sample quality and conditional fidelity. By reformulating temperature sampling as a practical mechanism for diversity control, the approach enables flexible adjustment ranging from coarse- to fine-grained levels. The effectiveness of the method is demonstrated across various pretrained diffusion and flow-matching models—including DiT, Stable Diffusion, and Motion Diffusion—highlighting its broad applicability and robustness.

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.

Problem

• diffusion models
• mode imbalance
• rare modes
• sample diversity
• distribution bias

Innovation

• temperature sampling
• variance correction
• time shifting
• diffusion models
• diversity control