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

2026-10-02

AI Summary

This study addresses the challenge of efficiently integrating dynamics in few-step sampling for flow matching models by proposing a training-free sampler. For the first time, this work introduces an L2-optimal affine approximation into the sampling process, decomposing the dynamics into affine and residual components via target moments. The affine component is solved exactly using matrix-valued propagation operators and explicit exponential integration techniques, overcoming the limitations of conventional scalar integrators that are restricted to isotropic linear dynamics. Experiments demonstrate that the proposed sampler significantly improves sample fidelity in few-step generation for both unconditional image synthesis and text-to-image tasks without requiring model retraining.

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.

Problem

• Flow Matching
• Few-Step Sampling
• Sample Fidelity
• Training-Free Sampler

Innovation

• Flow Matching
• Training-free Sampling
• Exponential Integrator
• Affine-Residual Decomposition
• Few-Step Generation