Inference-Time Self-Aligned Drifting for Few-Step Flow Matching
Abstract
Flow Matching (FM) provides an efficient formulation for continuous-time generative modeling via straight optimal transport (OT) trajectories. However, during accelerated sampling, these models face a fundamental variance-fidelity dilemma. Under few-step deterministic Ordinary Differential Equation (ODE) integration, the rigid OT trajectories suffer from mean concentration, leading to severe variance collapse and the loss of high-frequency details. Conversely, traditional Stochastic Differential Equations (SDEs) inject blind, exogenous Brownian noise to recover variance, but this isotropic exploration incurs inherent numerical attrition over long integration paths, bottlenecking peak fidelity. To resolve this, we introduce Self-Aligned Drifting (SADrift) to mitigate these challenges without external priors. SADrift is an inference-time framework that augments flow matching inference with self-interacting dynamics, without additional training. By maintaining a momentum history of the sampling trajectory, SADrift extracts the \emph{kinematic residual} to generate a deterministic extrapolative field. This mechanism forces the trajectory to efficiently explore the local tangent space, effectively decoupling exploratory variance into targeted deterministic dispersion and a minimal stochastic buffer. To safely anchor these dispersed trajectories and prevent out-of-distribution divergence, the residual drift is analytically integrated into a micro-time Ornstein-Uhlenbeck (OU) process, providing a mathematically sound, mean-reverting geometric safeguard. Requiring exactly zero additional Neural Function Evaluations (NFE), SADrift empowers a frozen base model to achieve a highly competitive FID of 2.00 at just 20 NFE on ImageNet-256. By successfully bypassing the numerical attrition of the 250-step SDE baseline (FID 2.06), SADrift establishes a new Pareto frontier for variance injection. Beyond image synthesis, we further validate SADrift on fast video generation paradigms (e.g., TurboDiffusion with Wan2.1-1.3B), demonstrating its universal capability to unlock latent generative quality across different modalities and acceleration regimes.