Temporal Score Rescaling for Temperature Sampling in Diffusion and Flow Models

We present a mechanism to steer the sampling diversity of denoising diffusion and flow matching models, allowing users to sample from a sharper or broader distribution than the training distribution. We build on the observation that these models leverage (learned) score functions of noisy data distributions for sampling and show that rescaling these allows one to effectively control a 'local' sampling temperature. Notably, this approach does not require any finetuning or alterations to training strategy, and can be applied to any off-the-shelf model and is compatible with both deterministic and stochastic samplers. We first validate our framework on toy 2D data, and then demonstrate its application for diffusion models trained across five disparate tasks -- image generation, pose estimation, depth prediction, robot manipulation, and protein design. We find that across these tasks, our approach allows sampling from sharper (or flatter) distributions, yielding performance gains e.g., depth prediction models benefit from sampling more likely depth estimates, whereas image generation models perform better when sampling a slightly flatter distribution.

Modern generative AI systems, such as image generators, protein design models, and robotics planners, often use diffusion or flow matching models to produce samples. While these models are powerful, users have limited control over how “conservative” or “diverse” their outputs are after training. Some applications require highly reliable predictions, while others benefit from greater diversity and creativity. We introduce a simple method that lets users control this tradeoff at inference time without retraining the model. Our approach enables sampling from either a sharper or broader version of the model’s learned distribution, producing outputs that are more focused or more diverse depending on the task. The method works directly with existing pretrained models and common sampling algorithms. We evaluate our approach on image generation, depth estimation, pose prediction, robot manipulation, and protein design. Our results show that adjusting sampling diversity can improve both performance and flexibility across a wide range of applications.