Track synthesis, materials and catalysis as it is published.
Search papers, labs, and topics across Lattice.
We introduce AREX, a training-free sampler for flow matching models that leverages the target mean and covariance to analytically integrate an affine component of the sampling dynamics. By decomposing the learned velocity field into an $L^2$-optimal affine approximation and a neural residual, AREX efficiently integrates the residual term using a matrix-valued propagator. This approach consistently enhances sample fidelity in few-step generation across image and text-to-image tasks without requiring model retraining.
Few-step sampling in flow matching models achieves significantly higher fidelity by analytically integrating the dominant affine component of the dynamics, leaving only a smaller neural residual to be integrated.
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.