Technical Summary: Diversify Diffusion with Temperature Sampling and Variance-Corrective Time Shifting

Problem Statement

Diffusion models faithfully reproduce their training distributions but inherently inherit data imbalances, making rare or under-represented modes difficult to reach during sampling. While a natural inference-time remedy is to sample from a high-temperature target distribution $p^{(\gamma)}_0(x) \propto p_0(x)^\gamma$ (where $0 < \gamma < 1$) to flatten dominant modes and boost rare ones, naive application of this concept fails in practice.

Directly scaling the learned score function by $\gamma$ (i.e., using $\gamma \nabla_x \log p_t(x)$) achieves the macroscopic goal of reweighting modes but induces a microscopic variance expansion. This excess variance breaks the strict correspondence between the score of the noised tempered marginal and the reverse diffusion process, resulting in off-manifold, blurry, and low-quality samples. Existing methods to enhance diversity often require retraining, auxiliary classifiers, or parallel particle filtering, which limits their applicability or fidelity.

Methodology: Variance-Corrective Time Shifting

The authors propose variance-corrective time shifting, a training-free mechanism that decouples mode reweighting from variance inflation. The core insight is that while the score of the tempered distribution $\nabla_x \log p^{(\gamma)}_t(x)$ cannot be obtained by simply scaling the original score, it can be accurately approximated by querying the network at a shifted timestep $\tilde{t}$.

Theoretical Derivation

1. The Mismatch: Direct tempering of the noisy marginal $p_t(x)^\gamma$ differs from the noisy marginal of the tempered data $p^{(\gamma)}_t(x)$. Under the assumption of well-separated modes, direct tempering effectively scales the variance of the Gaussian components by $1/\gamma$, whereas the true tempered marginal maintains the original variance schedule.
2. The Correction: To recover the desired tempered marginal without variance inflation, the method queries the pre-trained network at a shifted time $\tilde{t}$ such that the noise level satisfies $\sigma^2_{\tilde{t}} = \gamma \sigma^2_t$.
3. The Score Estimator: The score for the tempered distribution is approximated as:
   $$\nabla_x \log p^{(\gamma)}_t(x) \approx \gamma \left( -\frac{\epsilon_\theta(x, \tilde{t})}{\sigma^2_{\tilde{t}}} \right)$$
   This effectively cancels the variance expansion while preserving the mode reweighting effect.

Coarse-to-Fine Control

The method allows for a schedule $\gamma(t)$ applied at different stages of the reverse diffusion process:

• Early Stages (High Noise): Applying $\gamma(t) < 1$ here drives compositional diversity, steering sampling particles toward different attractors (modes).
• Middle/Late Stages (Low Noise): Applying $\gamma(t) < 1$ here drives local appearance variation while preserving the global composition and mode identity.

Key Contributions

1. Identification of Variance Expansion: The paper identifies that naive score scaling in diffusion models breaks the reverse diffusion process due to unintended variance inflation, degrading sample quality.
2. Variance-Corrective Time Shifting: A zero-cost, training-free correction that recovers the desired tempered marginal by querying the network at a shifted timestep, enabling effective high-temperature sampling.
3. Practical Diversity Knob: The method transforms temperature sampling into a practical tool for enhancing diversity in pre-trained diffusion and flow-matching backbones without retraining.
4. Stage-Specific Control: The authors demonstrate that the timing of the temperature intervention allows for coarse-to-fine control, distinguishing between compositional diversity (early) and local variation (late).

Experimental Results

The method was evaluated across various architectures (DiT, Stable Diffusion 1.5/3.5/XL, Motion Diffusion Models) and modalities (images, motion).

• Diversity vs. Fidelity: The method consistently improves diversity (measured by Vendi score) with minimal cost to sample quality and condition fidelity (measured by VQA score or classifier accuracy).
• Comparison:
  • vs. CADS: The proposed method maintains better prompt fidelity, whereas CADS often degrades alignment to the condition.
  • vs. Feynman-Kac (FK): FK struggles with high temperatures ($\gamma < 1$) in the authors' experiments, often yielding unusable samples, whereas the proposed method remains stable.
  • vs. Direct Tempering/TSR: Without time shifting, low $\gamma$ values introduce excessive noise. The time-shifting correction maintains high sample quality while achieving the same diversity gains.
• Quantitative Performance: On DiT, SD1.5, SDXL, and SD3.5, the method achieved the best or second-best scores in both CMMD (perceptual quality) and Vendi (diversity) while maintaining high fidelity.

Significance and Claims

The paper claims that variance-corrective time shifting resolves a fundamental limitation of applying temperature sampling to diffusion models. By correcting the variance inflation inherent in naive score scaling, the method enables users to explore the "long tail" of the data distribution—generating rare or diverse modes—without sacrificing the quality or condition adherence that diffusion models are known for.

The authors position this as a significant step toward flexible inference-time control, offering a "knob" for diversity that works across different model families (including flow-matching models) without the need for retraining or auxiliary data. They acknowledge that the method is bounded by the capacity of the underlying model and that the trade-off between fidelity and diversity requires tuning per model, but emphasize that the proposed time-shifting mechanism makes this trade-off manageable and effective.