PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shizheng Lin, Soon Hoe Lim, N. Benjamin Erichson

arXiv:2610.03483v1stat.MLcs.AIcs.LG

TL;DR

Flow matching can require many network evaluations for accurate sampling, limiting few-step generation. AREX uses target moments to propagate an affine component exactly and integrates only the residual. Across image and text-to-image tasks, it improves fidelity in the low-NFE regime without retraining.

  • Problem

    Accurately resolving a learned flow-matching velocity field can require many network evaluations, motivating efficient few-step sampling.

    Source: §1 · p. 1

  • Method

    AREX uses the target mean and covariance to separate the learned velocity into an analytically propagated affine component and a residual integrated with Adams–Bashforth quadrature.

    Source: §1 · p. 2

  • Results

    AREX consistently improves generation fidelity across CIFAR-10, ImageNet-256, and MJHQ-30K in the low-NFE regime without retraining the flow.

    Source: §1 · p. 3, §5 · p. 11

  • Takeaways & Limitations

    Few-step sampling can exploit analytically solvable structure in the flow while reserving numerical evaluations for residual dynamics.

    Source: §5 · p. 11

    The analysis uses a linear interpolation path from an independent standard-normal source to the target distribution.

    Source: §2.1 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Flow matching can require many network evaluations to resolve its learned flow, motivating training-free few-step sampling. AREX uses target moments to integrate an affine component analytically and reports better low-NFE fidelity across image and text-to-image tasks.

1 Introduction

  • Flow matching generates samples by integrating a learned velocity field, and accurate sampling can require many network evaluations, making few-step generation computationally important.Training-free samplers can be deployed on pretrained models, but scalar linear components do not capture anisotropic target geometry.Source: §1 · pp. 1–2
  • AREX separates the learned velocity into a target-moment-matched affine component, solved analytically, and a residual integrated with Adams–Bashforth quadrature using one new network evaluation per step.AREX-C extends the construction to conditional generation with shared within-condition covariance and a condition-specific mean estimated from the first network evaluation.Source: §1 · p. 2
  • The moment-matched Gaussian field is L2-optimal among affine approximations to the population velocity, and the paper also connects the residual to non-Gaussian score correction.The contributions include a closed-form affine propagator, discretization-error analysis, and a distributional error bound for AREX-ABp.Source: §1 · p. 2
  • AREX gives the best FID among compared training-free samplers at NFE 2–8 on CIFAR-10 and ImageNet-256 and NFE 4–10 on MJHQ-30K.The two-moons illustration also shows AREX recovering target geometry better than competing training-free samplers at NFE 4.Source: §1 · p. 3, Figure 2 · p. 3

2 AREX: Affine–Residual Exponential Integrator for Sampling

AREX uses the target mean and covariance to identify direction-dependent affine dynamics, propagates them exactly, and applies numerical quadrature only to the remaining residual. Its Adams–Bashforth variants reuse residual evaluations while requiring one new network evaluation per step.

2 AREX: Affine–Residual Exponential Integrator for Sampling

  • AREX replaces a Gaussian with the target’s mean and covariance to capture global geometry, solves the resulting affine flow exactly, and numerically integrates the difference from the learned velocity.For a general target, the moment-matched Gaussian field serves as the affine component, leaving non-Gaussian structure in the residual.Source: §2 · p. 3, §2.2 · p. 4

2.1 Target moments describe direction-dependent dynamics

  • Because target variances differ by direction, the source-to-target flow expands or contracts at different rates that a scalar linear term cannot capture.This motivates using the target covariance to define the direction-dependent affine component.Source: §2.1 · p. 4

2.2 A moment-matched affine flow

  • The moment-matched Gaussian velocity is affine and L2-optimal for a general target, and the covariance’s shared eigenbasis yields a closed-form propagator along its principal directions.The exact scale change along direction j from s to t is σj(t)/σj(s).Source: §2.2 · p. 4

2.3 Integrating the residual

  • Variation of constants leaves only the residual integral for numerical approximation, unlike scalar exponential samplers that use the same linear scaling in every direction.AREX instead removes the covariance-determined direction-dependent affine dynamics before applying the numerical solver.Source: §2.3 · p. 5

2.4 AREX-AB: quadrature of the residual

  • AREX-AB uses current and stored residual evaluations for higher-order quadrature while keeping one new network evaluation per step; its residual magnitude profile adjusts for systematic scale changes.The profile is estimated once per checkpoint without extra sampling evaluations, and the resulting weights are precomputed.Source: §2.4 · pp. 5–6

3 Theoretical Results

The theoretical results justify AREX through an optimal affine backbone, exact propagation of that backbone, and error control focused on the residual. They also specify how residual regularity and stage placement affect the solver’s accuracy.

3.1 Principle I: Let data geometry define the solvable backbone

  • The moment-matched affine field is the L2 orthogonal projection of the population velocity onto affine fields, while the ideal residual reflects non-Gaussian score correction and vanishes exactly for Gaussian targets.The sampler’s learned residual also includes neural approximation error.Source: §3.1 · pp. 6–7

3.2 Principle II: Propagate the solvable backbone exactly

  • The affine backbone has a closed-form propagator, so it incurs no discretization error; AREX’s leading local error omits pure-backbone terms and retains the residual’s material derivative.Unlike Euler, which discretizes the full velocity, AREX-AB applies higher-order reconstruction to the residual alone.Source: §3.2 · p. 7

3.3 Principle III: Spend numerical budget only on the residual

  • AREX-ABp’s error bound depends on residual temporal regularity and, for ρ > 0, an additional off-grid stage error from evaluating the residual at an affine-only predictor.At ρ = 0, the stage-state term vanishes and an order-p startup gives global order p; for fixed ρ > 0, the bound guarantees only first-order asymptotic accuracy in general.Source: §3.3 · p. 8

4 Empirical Results

AREX(-C) is evaluated on unconditional CIFAR-10, class-conditional ImageNet-256, and text-conditional MJHQ-30K generation. Its strongest results are in the low-NFE regime, with benchmark-specific performance depending on the sampler variant and integration order.

4 Empirical Results

  • CIFAR-10 uses an unconditional pixel-space I-CFM backbone, ImageNet-256 uses class-conditional SiT-XL/2 with CFG w = 1.5, and MJHQ-30K uses SANA-0.6B with CFG w = 4.5.The reported FID evaluations use 50K samples for CIFAR-10 and ImageNet-256, and 30K for MJHQ-30K.Source: §4.1 · p. 8, §4.2 · p. 9, §4.3 · p. 9

4.1 CIFAR-10: Unconditional Generation on Pixel Space

  • AREX is best on CIFAR-10 FID from NFE 2 to 16, while DPM-Solver-v3 is best at NFE 32.Its advantage is largest at low NFEs and diminishes as discretization is refined.Source: §4.1 · p. 8, Table 1 · p. 9

4.2 ImageNet-256: Class-Conditional Generation on Latent Space

  • AREX-C-AB3 is best on ImageNet-256 FID at NFE 2, 4, and 8; at NFE 16, four leading methods lie within 0.13 FID.At NFE 32, all methods except Euler and the one-node rule lie within 0.37 FID.Source: Table 2 · p. 9, §4.2 · p. 9

4.3 MJHQ-30K: Text-Conditional Generation with Guidance on Latent Space

  • AREX-C-AB2 is best among the Table 3 samplers on MJHQ-30K FID at every NFE, including NFE 10.AB2 also beats AB3 at every NFE; the authors relate this to the guided residual’s rapid variation along the trajectory.Source: §4.3 · pp. 9–10, Table 3 · p. 10

5 Conclusion

The conclusion presents AREX and AREX-C as training-free samplers that use data statistics for an exactly integrated affine backbone and reserve numerical work for the residual. The related work distinguishes this data-dependent structure from schedule-based, network-fitted, Gaussian-teleportation, and trajectory-straightening approaches.

5 Conclusion

  • AREX and AREX-C use data statistics to define an affine backbone, integrate it exactly, and spend the NFE budget resolving only the residual.The paper reports improved generation fidelity in the low-NFE regime across CIFAR-10, ImageNet-256, and MJHQ-30K.Source: §5 · p. 11

A Related Work

  • Unlike scalar schedule-based structure, AREX uses a data-dependent, anisotropic affine backbone determined by target covariance and applies Adams–Bashforth quadrature to its residual.Under the linear flow matching interpolant, its covariance-free first-order specialization reduces to an Euler/DDIM-type update.Source: A Related Work · p. 13
  • AREX derives its matrix-valued backbone from the L2-optimal affine approximation to the population velocity field, rather than fitting linear coefficients from a network or sampling trajectory.DPM-Solver-v3 uses empirical network statistics, while A-FloPS fits a scalar component locally from the evolving trajectory.Source: A Related Work · p. 13
  • Unlike Gaussian teleportation, which uses moment-matched Gaussian dynamics to skip an initial sampling portion, AREX retains affine dynamics throughout and integrates the non-Gaussian residual.STORK’s stability techniques target a complementary issue: numerical stability rather than identifying an analytically solvable anisotropic component.Source: A Related Work · p. 13
  • AREX replaces global flow straightness with the weaker requirement that unresolved residual dynamics be temporally simple.The paper states that AREX’s leading local error is governed by the residual’s material derivative rather than full-trajectory acceleration.Source: A Related Work · p. 13

B Conditional Sampling with AREX-C

AREX-C adapts AREX to conditional generation by using a shared within-condition covariance and a condition-specific affine center. It estimates that center from the first guided network evaluation, while calibration addresses the mismatch between its endpoint prediction and the guided terminal mean.

B Conditional Sampling with AREX-C

  • AREX-C uses the average within-condition covariance for its affine component, separating within-condition variation from variation between conditional means.The pooled terminal covariance includes both sources, but between-condition variation reflects differences between possible targets rather than variation within the target being sampled.Source: B Conditional Sampling with AREX-C · p. 14
  • A first endpoint prediction may systematically over- or underestimate the guided terminal-mean displacement because that mean depends on velocity along the entire trajectory.The method fits a single scalar calibration coefficient once per model and guidance configuration, not to a downstream evaluation metric.Source: B Conditional Sampling with AREX-C · pp. 15–16
  • AREX-C holds a source- and condition-specific affine center fixed along each trajectory, while sharing covariance-based propagators and residual quadrature weights across conditions.The shared within-condition covariance and residual magnitude profile can be calibrated once per checkpoint.Source: B Conditional Sampling with AREX-C · pp. 14–16
  • The first guided evaluation supplies both a condition-specific mean proxy and the first residual, so AREX-C adds no sampling-time network evaluation over AREX.Estimating the shared covariance, residual profile, and calibration coefficient incurs offline calibration cost.Source: B Conditional Sampling with AREX-C · pp. 16–17
  • For population conditional I-CFM with guidance w = 1, the shared within-condition construction has an exact restricted L2 interpretation.Its affine operator is the Gaussian affine field with the target covariance replaced by the average within-condition covariance.Source: B Conditional Sampling with AREX-C · p. 17
  • For guidance w ≠ 1, the guided field is generally not the conditional-expectation velocity of the original interpolant, so the projection result does not apply directly.AREX-C's calibrated affine pair is therefore evaluated empirically in that setting.Source: B Conditional Sampling with AREX-C · pp. 15, 17

C Proofs and Details of the Main Theoretical Results

The appendix formalizes the Gaussian affine field and shows it is the L2 projection of the marginal velocity onto affine vector fields. For Gaussian targets, that affine field equals the marginal velocity.

C Proofs and Details of the Main Theoretical Results

  • For a Gaussian target, the moment-matched affine field equals the marginal velocity, so the residual is zero.The appendix states the endpoint convention for defining the field before t = 1 and notes that its propagator has a continuous terminal extension even when the covariance is singular.Source: C Proofs and Details of the Main… · p. 18

C.1 Proof of Proposition 1

  • For general targets, the moment-matched Gaussian field is the L2 orthogonal projection of the marginal velocity onto the space of affine vector fields.The projection follows because the difference between the marginal velocity and the affine field is orthogonal to every affine vector field.Source: §C.1 · pp. 18–19

C.2 Proof of Proposition 2

The appendix derives a score representation and Gaussian characterization, then gives the affine propagator and error comparisons. These results explain when the affine component is exact and how its integration affects local error.

C.2 Proof of Proposition 2

  • The marginal velocity can be represented using the score of the interpolated density, and a zero residual characterizes the Gaussian-target case.For a Gaussian target, the interpolated law is Gaussian and the moment-matched Gaussian field equals the marginal velocity.Source: §C.2 · pp. 19–21

C.3 Proof of Proposition 3

  • The affine operators share an eigenbasis across time, yielding a closed-form propagator and a composition property for the affine flow.Variation of constants then represents the full learned dynamics as affine propagation plus an integrated residual contribution.Source: §C.3 · pp. 21–22

C.4 Local Error Comparison

  • AREX successively removes affine contributions from the leading local error, and with zero residual it gives the exact affine flow on any grid.Euler and frozen exponential Euler remove fewer of those affine contributions.Source: §C.4 · p. 24

C.5 Proof of Proposition 4

  • Under the stated Lipschitz and propagator bounds, the residual quadrature error is bounded by ρh^2Φ^2LxRmax.The bound controls the difference between residual values along the exact and approximate trajectories.Source: §C.5 · p. 25

C.6 Proof of Theorem 1

The appendix bounds AREX-ABp's distributional error under regularity assumptions and establishes order p for its grid-node variant. It also discusses coupling choices, sampler relations, and residual profiling.

C.6 Proof of Theorem 1

  • Under finite stability and Lipschitz assumptions, the theorem bounds terminal solver error using residual-reconstruction error and startup error.The result also gives a Wasserstein bound and a smooth-residual specialization.Source: §C.6 · pp. 26, 28

C.7 Specialization to the Left-Node ρ = 0 Case

  • Grid-node AREX-ABp has global order p relative to the exact learned ODE under the theorem's assumptions.The naive ramp is generically only second-order because its first startup step uses AREX-1.Source: §C.7 · pp. 29–30

D.1 Why Not Other Couplings?

  • Independent endpoint coupling makes the affine coefficients share a fixed eigenbasis and admits a closed-form propagator; general couplings may require estimating cross-covariance and need not commute across time.The OT cross-covariance is generally unavailable without solving or approximating the OT problem.Source: §D.1 · p. 31

D.2 Relation to DDIM and Other Samplers

  • The zero-covariance specialization coincides with an Euler step only for the specified linear interpolation; nonzero target covariance adds covariance-informed affine transport.The full AREX method also differs through its treatment of the residual integral.Source: §D.2 · pp. 31–32

D.3 The Residual Magnitude Profile

  • The residual magnitude profile factors out empirically measured time-varying RMS amplitude, leaving normalized residual variation for quadrature without adding sampling-time network evaluations.The profile is estimated from calibration trajectories and precomputed.Source: §D.3 · p. 32

D.4 AREX versus Euler: A Spectral Analysis

The covariance spectrum controls both Euler’s affine discretization error and the direction-dependent dynamics that AREX propagates exactly. Curvature of the full flow need not prevent easy integration when the remaining residual is sufficiently regular; κ diagnoses affine amplification but does not predict a specific solver improvement.

D.4 AREX versus Euler: A Spectral Analysis

  • AREX removes the affine-part discretization error that grows with covariance eigenvalues, while Euler’s leading error is controlled by the target covariance spectrum.For λj ≫1, the removed relative multiplicative error scales as π/8 h√λj; at λj = 0, Euler and AREX are both exact for the affine mode.Source: §D.4 · pp. 33–34

D.5 Relation to Curvature of the Flow

  • The relevant integration difficulty is residual curvature after removing the solvable affine dynamics, not whether the full flow is straight.For Gaussian targets the ideal marginal flow is affine, so AREX propagates it without discretization error; for general targets, curvature is bounded by excess transport cost of the independent coupling.Source: §D.5 · pp. 34, 36

D.6 A Diagnostic From the Propagator Itself

  • κ measures the largest forward amplification of the exact affine propagator, but does not by itself predict improvement over another solver.The covariance spectrum also identifies directions with strong structured dynamics; the reported models and tasks all have κ ≫1.Source: §D.6 · pp. 36–37

E.1 Experimental Settings

The experiments compare training-free samplers on three pretrained-model benchmarks, with per-NFE settings selected on held-out data. AREX also requires moment and residual-profile preparation, while its per-step overhead is small relative to a network evaluation.

E.1 Experimental Settings

  • All samplers use the same pretrained checkpoints, noise seeds, and FID code across unconditional CIFAR-10, class-conditional ImageNet-256, and text-conditional MJHQ-30K.Discrete method options are selected per NFE on held-out splits, but SANA’s split is a subset of the evaluation prompts.Source: §E.1 · pp. 37, 39

E.2 Moment and Profile Estimation Details

  • AREX’s moments are estimated from training data when available, while SANA uses self-sampling because its training set is unavailable.Its residual profile is measured once per checkpoint and pooled across guided conditions; dense covariance eigendecomposition scales as d^3, and larger-dimensional backbones were not tested due to time and computation constraints.Source: §E.2 · pp. 40–41

E.3 Wall-Clock Cost

  • AREX adds about 0.3–1.5% of a network call in per-step overhead, but requires offline moments and residual-profile evaluation before sampling.Six of the seven baselines do not require that offline cost; when training data are available, the moments cost no network evaluations.Source: §E.3 · pp. 41–42

E.4 Ablation Study

The ablations compare residual quadrature choices, full-field integration, residual profiling, and AREX-C's moment choices. They show that performance depends on the integration budget and backbone, with multistep residual rules favored at the smallest budgets.

Order of the quadrature method

  • AB2 leads on CIFAR-10 and SANA, while AB3 leads on SiT-XL/2 from NFE 4; the largest gaps occur at NFE 4.The two multistep rules are closer than the one-node rule and their relative performance depends on the backbone.Source: §E.4 · p. 42, Table 12 · p. 42

Comparison to classical integrators

  • At NFE ≤ 8, higher-order classical integration of the full field does not reach AREX; at NFE 32, full-field AB3 and RK4 slightly lead on SiT-XL/2.The comparison uses equal NFE; the gap to AREX closes from NFE 16.Source: §E.4 · p. 42, Table 13 · p. 43

The residual profile

  • The residual profile captures predictable changes in residual size across a coarse step, but does not improve full-field integration at NFE ≤ 4 on CIFAR-10 or SiT-XL/2.On SANA, full-field AB2 with ψ trails AREX-C-AB2 by 21.8 FID at NFE 4, but is close at NFE 9 and 10.Source: §E.4 · p. 42, Table 14 · p. 44

Two ingredients of AREX-C

  • On SiT-XL/2, conditional mean improves FID by 4.4 at NFE 4 and 0.19 at NFE 8; on SANA, within-condition covariance improves it by 1.68 at NFE 4.On SANA, conditional mean improves FID by 0.11 at NFE 8 but is 0.34 behind the pooled mean at NFE 4.Source: §E.4 · p. 43

Alternative integrators for the residual

  • At NFE 2, both Adams–Bashforth rules lead every alternative by more than 36 FID; from NFE 8 onward, single-step rules catch up, with ETD-RK4 ahead of both Adams–Bashforth rules at NFE 8.The exponential form is ahead of the integrating-factor form at every NFE up to 16 for both stage counts, by a wide margin for k = 4.Source: §E.4 · p. 44

E.5 Additional Metrics on ImageNet-256

Additional ImageNet-256 metrics show AREX-C's strengths vary across metrics and NFE. Qualitative panels provide matched-noise samples at several budgets for CIFAR-10, ImageNet classes, and text-to-image prompts.

E.5 Additional Metrics on ImageNet-256

  • At NFE 2, AREX-C has the best sFID, IS, precision, and recall; it has best or tied-best precision at every NFE.DPM-Solver-v3 leads sFID and IS at NFE 4; AREX-C leads IS from NFE 8 and ties DPM-Solver-v3 on sFID at NFE 8 and 16.Source: §E.5 · p. 44, Table 18 · p. 46

E.6 Qualitative Samples

  • CIFAR-10 panels compare samplers at NFE 2, 4, and 8 using the same initial noise within each column group.The panels show four noise draws across Figures 4–6.Source: Figure 4 · p. 47, Figure 5 · p. 47, Figure 6 · p. 48
  • ImageNet class samples compare schemes from NFE 2 to 8; AREX is described as producing higher fidelity, more coherent structure, and finer detail.Additional panels show SANA text-to-image samples at NFE 4, 6, and 8 for MJHQ-30K prompts.Source: Figure 7 · p. 48, Figure 8 · p. 49, Figure 9 · p. 50

Related papers

Loading 2610.03483v1…