Test Time Scaling of Diffusion Model via Flow Matching Corrector

Sheng-Yen Chou, Yiding Chen, Shikun Sun, Wen Sun

Keywords: Diffusion Model, ML, Test-Time Scaling

Abstract: Scalable self-improving capability is a desirable property of generative models, as it enables performance improvement with additional computational resources rather than requiring more training data. Existing approaches typically rely on external reward signals to fine-tune generative models and improve the generation. In this paper, we propose the Scalable Self-Improving Correction (SSI-Corr) framework, which requires neither new training data nor external rewards. Instead, SSI-Corr trains a corrector that directly aligns the sample generation process with the target distribution. Our method is supported by theoretical analysis and scales effectively with available computational resources. In the experiment, we demonstrate that SSI-Corr improves FID scores by 27\% and 14.3 \% on a pre-trained unconditional CIFAR10 DDPM with ancestral and DDIM samplers respectively.

Primary Area: generative models

Submission Number: 5310