Temporal Score Rescaling for Temperature Sampling in Diffusion and Flow Models

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Xu, Yanbo, Wu, Yu, Park, Sungjae, Zhou, Zhizhuo, Tulsiani, Shubham
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913157493030912
author Xu, Yanbo
Wu, Yu
Park, Sungjae
Zhou, Zhizhuo
Tulsiani, Shubham
author_facet Xu, Yanbo
Wu, Yu
Park, Sungjae
Zhou, Zhizhuo
Tulsiani, Shubham
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2510_01184
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Temporal Score Rescaling for Temperature Sampling in Diffusion and Flow Models
Xu, Yanbo
Wu, Yu
Park, Sungjae
Zhou, Zhizhuo
Tulsiani, Shubham
Machine Learning
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.
title Temporal Score Rescaling for Temperature Sampling in Diffusion and Flow Models
topic Machine Learning
url https://arxiv.org/abs/2510.01184