Acceleration-Aware Sampling for Few-Step Rectified Flow Models

Xiaomeng Fu, Jia Li, Yiming Hu, Yong Wang, Xi Wang, Hayden Kwok-Hay So, Xiangxiang Chu

Keywords: rectified flow models, text-to-image generation, few-step image generation

TL;DR: Design an acceleration-aware sampler to enhance image quality in few-step settings

Abstract:

Rectified flows (RFs) enable efficient, high-fidelity image synthesis by integrating a learned velocity field from noise to data.
However, in latency-constrained scenarios with few-step setting, RFs tend to produce degraded images. We identify this failure in the common Euler sampler which assumes piecewise-constant velocity and neglects inherent acceleration. Consequently, significant discretization errors occur and dominate the few-step sampling process. To address this, we introduce \textbf{A}cceleration-\textbf{A}ware \textbf{S}ampling (\textbf{A\textsuperscript{2}S}) which explicitly accounts for acceleration under the same computational cost as Euler sampler. Starting from a second-order view, we decomposes acceleration into temporal and spatial components, and compensates for both with lightweight approximations. Specifically, temporal acceleration is addressed via a time-shifted velocity evaluation. In this way, updates align with the mid-interval dynamics while preserving one forward pass per step. Meanwhile, spatial acceleration is captured by a smooth, time-dependent gain that modulates step size. As a result, A\textsuperscript{2}S is model-agnostic, plug-and-play for existing pretrained rectified-flow models and requires no retraining. Across multiple models and benchmarks, A\textsuperscript{2}S consistently improves image quality and stability in the few-step setting and remains competitive as the step count increases. Moreover, on FLUX, few-step A\textsuperscript{2}S even surpasses standard multi-step samplers in image aesthetics and text–image alignment.

Primary Area: generative models

Submission Number: 1247