Arex sampler improves few-step flow matching for image generation

AREX: Affine-Residual Exponential Integrator for Few-Step Sampling in Flow Matching

Artificial IntelligenceMachine Learning

Summary

Generating images or text-based images often involves a process called flow matching, which can be slow or need many steps. The authors created AREX, a fast method that breaks down the problem into a predictable mathematical part and a smaller, learned part to speed up image generation. This method works without retraining the existing models and produces better-quality images with fewer steps. AREX uses math to handle the simpler part exactly and only learns from what’s left over, making the process more efficient.

What this means in practice

Authors

Shizheng Lin, Soon Hoe Lim, N. Benjamin Erichson

Abstract

We introduce AREX, a training-free sampler for pretrained flow matching models that uses the target mean and covariance to capture an analytically tractable part of the sampling dynamics. We show that the velocity field of the moment-matched Gaussian target is the $L^2$-optimal affine approximation to the marginal velocity field. This motivates decomposition of the learned dynamics into an affine component over the whole sampling path, determined by the first two target moments, and a neural residual term. AREX keeps the affine component and integrates it using an explicit matrix-valued propagator. In turn, we only require to integrate over the residual term. This differs from scalar exponential integrators, which analytically handle only isotropic linear dynamics. Across image and text-to-image generation tasks, AREX consistently improves sample fidelity in the few-step sampling regime without retraining the underlying model.
View PDFOpen arXiv