Lightweight Spectral Statistical Correction for Fast Sampling in Flow Models

Diffusion and flow matching models achieve remarkable quality in various generation tasks. However, solving the flow ODE to a good precision typically requires many network evaluations, leading to high latencies. Multiple lightweight methods have been proposed for fast sampling with pretrained models. However, the vast majority of them focus on minimizing the truncation error of the underlying ODE solver, and do not attempt to directly control the distributional shift caused by taking large discretization steps. In this paper, we demonstrate that adding a simple fast analytical statistical correction after each solver step can significantly improve the quality of the generated images. Specifically, we present Spectral Statistical Correction (SSC) - a linear transformation applied independently to each DCT coefficient of the latent state so as to optimally correct its distribution. SSC can be learned from a small number of images without requiring backpropagation through the model. Furthermore, it can be applied on top of other fast sampling methods. As we show, SSC consistently leads to substantial improvements even over strong samplers, allowing to obtain high quality results with as little as 5 NFEs with modern flow models.