Stochastic Few-step Models

Abstract

Reward alignment for flow and diffusion models, via test-time steering or fine-tuning, often relies on sampling from conditional distributions $p_{\tau\mid t}(\space\cdot\mid x_t)$ induced by stochastic dynamics. In practice, this creates a computational bottleneck requiring expensive SDE simulation. Meanwhile, recent few-step accelerations for generative flows largely target deterministic dynamics and therefore do not directly address stochastic conditional sampling. We introduce stochastic few-step models, a framework for fast sampling from SDE-defined conditional distributions by mapping the conditional SDE to a deterministic ODE with matching marginals. Building on this formulation, we show that the resulting conditional ODEs can be effectively distilled into a single few-step model, enabling efficient conditional rollouts. Experiments on Gaussian-mixture and MNIST steering show that the resulting sampler provides accurate conditional samples to improve reward steering outperforming standard denoiser heuristics.

Alex Tong

Principal Investigator

I work on improving flow models with applications to cells and proteins.