AREX: Affine-Residual Exponential Integrator for Few-Step Sampling in Flow Matching
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 L2-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.
Related
- Energy stability and error estimates for a second-order structure-preserving exponential integrator method for smectic-A liquid crystals
- Strong error analysis of a stochastic exponential integrator for SPDEs driven by fractional Brownian motion with H<1/2
- On Natural Exponential Families with Log-Affine Generalized Variance
- Improving score-based sampling via affine post-processing
- Beyond exponential tilting: sharp asymptotics and improved importance sampling for rare events
BibTeX
@misc{lin2026arex,
title = {{AREX: Affine-Residual Exponential Integrator for Few-Step Sampling in Flow Matching}},
author = {Shizheng Lin and Soon Hoe Lim and N. Benjamin Erichson},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.03483},
eprint = {2610.03483},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.03483},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - AREX: Affine-Residual Exponential Integrator for Few-Step Sampling in Flow Matching
AU - Shizheng Lin
AU - Soon Hoe Lim
AU - N. Benjamin Erichson
PY - 2026
DA - 2026/10/02
PB - arXiv
DO - 10.48550/arxiv.2610.03483
UR - https://arxiv.org/abs/2610.03483
AB - 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.
ER -