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

Shizheng Lin; Soon Hoe Lim; N. Benjamin Erichson

stat.MLOct 2, 2026

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

Authors: Shizheng Lin, Soon Hoe Lim, N. Benjamin Erichson

Organizations: Department of Mathematics, KTH Royal Institute of Technology · Nordita, KTH Royal Institute of Technology and Stockholm University · International Computer Science Institute · Lawrence Berkeley National Laboratory

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.

Figures & tables

Figure 1: Few-step text-to-image generation with AREX. Qualitative samples from the pretrained FLUX.1-dev model using AREX-C with only $8$ network evaluations. AREX analytically propagates the moment-matched affine component of the learned flow and numerically integrates only the residual. All samples are generated without retraining the model.

Figure 2: Unconditional sampling on the two-moons target at NFE $4$ . The Euler-1000 solution is shown as a high-accuracy reference. At the same four-evaluation budget, AREX recovers the target geometry better than the competing training-free samplers.

Sampler	NFE 2	NFE 4	NFE 8	NFE 16	NFE 32
Euler	166.24	51.88	17.46	9.15	6.28
DPM-Solver++	166.24	43.18	11.48	6.74	4.99
UniPC	166.24	36.53	10.49	7.14	5.29
A-FloPS	145.63	29.27	10.29	7.14	6.02
CAB	145.05	23.77	8.91	5.85	4.23
DPM-Solver-v3	112.04	24.69	6.95	3.92	3.55

Table 1: FID results for CIFAR-10. AREX is best from NFE 2 to 16; DPM-Solver-v3 is best at NFE 32.

Sampler	NFE 2	NFE 4	NFE 8	NFE 16	NFE 32
Euler	154.29	53.98	11.02	4.17	2.65
DPM-Solver++	154.28	31.37	4.20	2.56	2.12
UniPC	154.28	22.08	3.57	2.22	2.05
A-FloPS	129.73	20.74	3.34	2.39	2.39
CAB	132.94	12.37	4.04	2.26	2.02
STORK-4	132.95	12.59	3.98	2.48	2.19

Table 2: FID results for ImageNet-256, with SiT-XL/2, class-conditional with CFG at SiT’s published scale

w=1.5

. AREX-C-AB3 is best at NFE

2

4

and

8

. At NFE

16

, UniPC, DPM-Solver-v3, CAB and AREX-C-AB3 lie within

0.13

FID of each other, and at NFE

32

every method but Euler and the one-node rule lies within

0.37

FID.

Figure 3: Samples generated by different sampling schemes with SANA-0.6B, MJHQ-30K prompts, $w=4.5$ , at NFE 4, 6 and 8, using the configurations of Table 3 .

Sampler	NFE 4	NFE 5	NFE 6	NFE 7	NFE 8	NFE 9	NFE 10
DPM-Solver++(2M)	23.39	14.96	11.59	9.84	8.85	8.24	7.99
UniPC	25.25	15.47	12.06	10.31	9.36	8.69	8.38
A-FloPS	62.48	33.88	21.72	15.68	12.68	10.98	9.91
CAB	60.00	48.88	41.68	32.20	23.99	18.05	13.45
STORK-4	73.34	37.30	16.33	9.68	7.70	7.07	6.70
DPM-Solver-v3	10.56	8.80	7.95	7.58	7.11	6.98	6.87

Table 3: FID results for MJHQ-30K, with SANA-0.6B 512px, CFG

w=4.5

. STORK-4 is run in its published configuration at every NFE; NFE 4–6 lie below the range reported in its paper. Among the samplers in the table, AREX-C-AB2 is best at every NFE, including NFE 10.

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Model	$\lambda_{\max}$	$\kappa$	AREX	Euler	best baseline
SANA-0.6B / MJHQ-30K	$76$	$8.8$	$8.2$	–	$10.6$ (DPM-Solver-v3)
I-CFM / CIFAR-10	$221$	$14.9$	$12.2$	$51.9$	$23.2$ (STORK-4)
SiT-XL/2 / ImageNet	$266$	$16.3$	$8.2$	$54.0$	$9.8$ (DPM-Solver-v3)

Appendix

Table 4: The propagator gain

\kappa=\sqrt{1+\lambda_{\max}}

, computed from the covariance used by the sampler, together with the FID at NFE 4.

Parameter	CIFAR-10	ImageNet-256	MJHQ-30K
Model	I-CFM (TorchCFM)	SiT-XL/2	SANA-0.6B, $512$ px
Space	pixel	latent (SD-VAE)	latent
Dimension $d$	$3072$	$4096$	$8192$
Precision	fp32	bf16	bf16
Conditioning	none	class, CFG $w=1.5$	text, CFG $w=4.5$
Conditions	–	$1000$ classes, cycled	MJHQ prompts

Appendix

Table 5: Settings of the three benchmarks.

Method	Implementation	Options selected on the held-out split
Euler	forward Euler on the learned field	–
DPM-Solver++ Lu et al. (2025)	DPMSolverMultistepScheduler ( dpmsolver++ , flow prediction)	solver order ( $1$ , $2$ or $3$ ); factory order $2$ on SANA
UniPC Zhao et al. (2023)	UniPCMultistepScheduler , same configuration	solver order ( $1$ , $2$ or $3$ ); factory order $2$ on SANA
A-FloPS Jin et al. (2026)	authors’ AEuler , constraint $\lambda\in[-1,1]$	time grid (unshifted or shifted), once per pipeline
CAB Roy and Nair (2026)	released scheduler	solver order ( $2$ or $3$ )
STORK-4 Tan et al. (2026)	authors’ scheduler	$(s,\text{Taylor order})\in\{5,9,14\}\times\{1,2\}$ on CIFAR-10 and SiT-XL/2; published configuration ( $s=9$ , first order) on SANA

Appendix

Table 6: Baselines, their implementations, and the discrete options that are selected per NFE on the held-out split. Everything else is the configuration of the paper or the default of the code.

Method	Option	CIFAR-10	SiT-XL/2	SANA-0.6B
DPM-Solver++	solver order	$1,\,2,\,3,\,3,\,3$	$1,\,2,\,3,\,3,\,3$	$2$ (factory)
UniPC	solver order	$1,\,2,\,3,\,3,\,3$	$1,\,2,\,3,\,3,\,3$	$2$ (factory)
CAB	solver order	$2,\,3,\,3,\,3,\,3$	$2,\,2,\,2,\,2,\,2$	$2$
A-FloPS	grid shift (per pipeline)	$1$	$3$	$1$
STORK-4	$(s,\text{Taylor order})$	$(14,1)$ , $(9,2)$ , $(9,2)$ , $(9,1)$ , $(9,1)$	$(5,2)$ , then $(14,2)$	$(9,1)$ (published)
DPM-Solver-v3	order	$2,\,2,\,3,\,3,\,3$	$2,\,2,\,3,\,3,\,3$	$2$

Appendix

Table 7: Options selected per NFE on the held-out split. Entries are listed for NFE

2

4

8

16

32

on CIFAR-10 and SiT-XL/2, and for NFE

4

10

on SANA, where one value serves every NFE. “Factory” and “published” mark configurations that were fixed rather than selected. The pseudo-order option of DPM-Solver-v3 was never selected.

Model	Rule	NFE 2	NFE 4	NFE 8	NFE 16	NFE 32
CIFAR-10	one-node	$0.5$	$0.35$	$0.35$	$0.2$	$0.1$
	AB2	$0.2$	$0.1$	$0.2$	$0.2$	$0.35$
	AB3	$0.2$	$0.2$	$0.2$	$0.2$	$0.35$
SiT-XL/2	one-node	$0.2$	$0.2$	$0.2$	$0.2$	$0.2$
	AB2	$0.05$	$0.05$	$0.05$	$0.05$	$0.1$
	AB3	$0.05$	$0.05$	$0.1$	$0.05$	$0.1$

Appendix

Table 8: Selected node position

\rho

for each backbone, rule, and NFE.

Model	Rule	NFE	$\rho=0.05$	$0.1$	$0.2$	$0.35$	$0.5$
CIFAR-10	AB2	4	15.37	14.11	14.72	18.11	26.02
	AB2	8	8.52	7.92	7.39	8.69	13.13
	AB3	4	18.34	15.85	14.23	16.27	24.07
	AB3	8	11.22	9.54	7.58	7.96	12.33
SiT-XL/2	AB2	4	22.34	23.48	28.01	40.24	–
	AB2	8	16.33	16.39	17.19	19.90	–

Appendix

Table 9: Selection of the node position

\rho

. FID on the held-out split (CIFAR-10

n=10{,}000

, SiT-XL/2

n=2{,}500

, SANA

n=1{,}000

) with the residual profile in place. The selected value in each row is in bold. These are small-sample FIDs, so only differences within a row are meaningful.

Model	$d$	batch	Network	AREX step	overhead	DPM++ step
CIFAR-10, I-CFM (fp32)	3,072	1000	0.140	0.0021	1.5%	0.0003
SiT-XL/2, $w=1.5$	4,096	125	2.12	0.0058	0.3%	0.0027
SANA-0.6B, $w=4.5$	8,192	25	3.72	0.041	1.1%	0.010

Appendix

Table 10: Per-step cost, ms per sample, one GH200. "Network" is one evaluation as the benchmark makes it (a CFG pair where the model is guided); "AREX" is everything an AREX-AB3 step adds to it. For SiT-XL/2 and SANA-0.6B, we use bf16 and CFG pair.

	AREX: profile $\phi$		AREX: moments		DPM-Solver-v3: EMS
Model	evals	wall-clock	evals	wall-clock	evals (+JVP)	wall-clock
CIFAR-10	$1.0\times 10^{5}$	$14$ s	$0^{\dagger}$	–	$1.2\times 10^{6}$	$27$ min
SiT-XL/2, $w=1.5$	$1.0\times 10^{5}$	$3.6$ min	$0^{\dagger}$	$3$ min	$2.6\times 10^{5}$	$380$ min
SANA-0.6B	$5.1\times 10^{4}$	$3.2$ min	$1.0\times 10^{6}\,{}^{\ddagger}$	$62$ min	$2.6\times 10^{5}$	$160$ min

Appendix

Table 11: One-time cost per checkpoint: network evaluations (a CFG pair counts as one) and wall-clock on one GH200. AREX’s residual profile

\phi

is the network cost the method requires. Its moments (with the conditional means and

\alpha

) are statistics of the target: from the training set they cost no network evaluation (CIFAR-10 on I-CFM and ImageNet on SiT). On SANA we could not access the training set, so we estimated them by self-sampling to substitute the training set, listed on its own ( † training set, ‡ self-sampled; with the training set the pass would be one VAE-encoder run.); the covariance accumulation, eigendecomposition and weight tables add

0.1

–

0.7

s. DPM-Solver-v3: the EMS pass at the authors’ settings. Euler, DPM-Solver++, UniPC, STORK, A-FloPS and CAB have no such cost.

Model	Rule	NFE 2	NFE 4	NFE 8	NFE 16	NFE 32
CIFAR-10	one node	54.63	26.13	11.45	5.34	3.66
	AB2	39.24	12.20	5.39	3.50	3.64
	AB3	39.24	12.35	5.58	3.60	3.69
SiT-XL/2	one node	134.79	34.14	9.70	4.14	2.72
	AB2	84.15	9.58	3.23	2.38	2.17
	AB3	84.15	8.18	3.02	2.24	2.12

Appendix

Table 12: Order of the quadrature method, rows are from the main tables.

Model	Integrator	NFE 2	NFE 4	NFE 8	NFE 16	NFE 32
CIFAR-10	Euler	166.24	51.88	17.46	9.15	6.28
	AB2	145.06	29.23	10.46	6.98	5.19
	AB3	145.06	25.52	9.30	6.21	4.67
	Heun	412.44	295.30	119.29	48.64	18.07
	RK4	412.44	198.53	76.15	26.79	7.99
	AREX-AB2	39.24	12.20	5.39	3.50	3.64

Appendix

Table 13: Classical integrators on the full field at equal NFE, against the AREX row of the main tables. CIFAR-10 and SiT-XL/2 at 50k samples, SANA-0.6B at 30k. Heun and RK4 take their stages inside chunks of two and four consecutive grid intervals, a shorter final chunk uses the rule of its length (Euler for one interval, Heun for two, Kutta’s third order rule for three), as for the ETD-RK rows of Table 17 . Thus the RK4 entry at NFE 2 is one Heun step, and at NFE 5, 6, 7, 9 and 10 on SANA the last chunk is shorter,

Model	Sampler	NFE 2	NFE 4	NFE 8	NFE 16	NFE 32
CIFAR-10	AREX (one node), flat	59.62	29.04	13.21	6.23	4.01
	AREX (one node), profile	54.63	26.13	11.45	5.34	3.66
	AREX-AB2, flat	55.73	23.67	9.07	4.10	3.63
	AREX-AB2, profile	39.24	12.20	5.39	3.50	3.64
	AREX-AB3, flat	55.73	21.23	8.06	3.88	3.73
	AREX-AB3, profile	39.24	12.35	5.58	3.60	3.69

Appendix

Table 14: The residual profile

\phi

( 13 ): "profile" is the sampler of the main tables, "flat" removes

\phi

from every weight.

Model	Rule	NFE 2	NFE 4	NFE 8	NFE 16	NFE 32
CIFAR-10	AB2	145.06	29.23	10.46	6.98	5.19
	AB2 $+\phi$	216.19	32.26	7.22	5.12	4.56
	AB3	145.06	25.52	9.30	6.21	4.67
	AB3 $+\phi$	216.19	29.55	6.54	4.60	4.15
	AREX-AB2	39.24	12.20	5.39	3.50	3.64
SiT-XL/2	AB2	132.95	15.74	4.19	2.46	2.12

Appendix

Table 15: Classical Adams–Bashforth rules on the full field with (

+\psi

) and without the magnitude profile of the full field, at equal NFE; the AREX row of the main tables for reference.

	SiT-XL/2		SANA-0.6B
Covariance, mean	NFE 4	NFE 8	NFE 4	NFE 8
pooled $\Sigma_{1}$ , $m_{1}$	12.58	3.18	9.49	6.38
$\bar{S}$ , $m_{1}$	12.54	3.21	7.81	6.44
$\bar{S}$ , $\mu(c)$	8.18	3.02	8.15	6.33

Appendix

Table 16: The ingredients ablation of AREX-C. Pooled

\Sigma_{1}

with

m_{1}

is plain AREX;

\bar{S}

with

\mu(c)

is AREX-C. SiT-XL/2: AREX-C-AB3; SANA-0.6B: AREX-C-AB2.

Residual integrator	NFE 2	NFE 4	NFE 8	NFE 16	NFE 32
AREX-AB2, flat (1 NFE/step)	55.73	23.67	9.07	4.10	3.63
AREX-AB3, flat (1 NFE/step)	55.73	21.23	8.06	3.88	3.73
ETD-RK2m (2)	106.63	31.06	11.66	4.46	3.59
ETD-RK2 (2)	91.89	26.75	10.49	4.20	3.67
ETD-RK3 (3)	91.89	75.62	14.06	15.58	3.63
ETD-RK4 (4)	91.89	29.43	7.93	3.95	3.94

Appendix

Table 17: Alternative integrators for the residual at equal total NFE on CIFAR-10, without the residual profile; the AREX rows are the flat rows of Table 14 . ETD-RK

k

: exponential Runge–Kutta with

k

equispaced stages (stage residuals interpolated and integrated against the propagator; ETD-RK2m freezes

r

at the midpoint stage, the DPM-Solver-2 rule); IF-RK

k

: classical

k

-stage tableau on the residual after the affine part has been transported out (Lawson); AREX+STORK: STORK-4 on the residual with the affine transport re-based at every step,

(s,\text{order})

selected on the same held-out split as STORK. When NFE is not a multiple of

k

, the last step uses the remaining stages. At NFE

2

, every rule with

k\geq 3

therefore reduces to its two-stage version; ETD-RK3 is also affected at NFE 4, 8 and 16 and is shown for completeness.

Sampler	Metric	NFE 2	NFE 4	NFE 8	NFE 16	NFE 32
Euler	sFID	200.0	46.0	10.5	6.1	4.9
	IS	4.8	60.7	180.5	241.0	264.3
	P / R	.058/.123	.365/.440	.702/.505	.779/.562	.796/.590
DPM-Solver++	sFID	200.0	22.7	7.4	5.5	4.7
	IS	4.8	98.3	233.6	265.7	275.9
	P / R	.058/.123	.517/.455	.793/.546	.801/.582	.804/.588

Appendix

Table 18: SiT-XL/2 on ImageNet-256,

50

k samples, classifier-free guidance

w=1.5

: sFID, Inception score (IS), precision (P), and recall (R) for the samplers of Table 2 , in the same configurations. Bold marks the best value of each metric at each NFE.

Figure 4: CIFAR-10, I-CFM: four noise draws at NFE 2, 4 and 8, every sampler in the configuration of Table 1 . Same initial noise down each column group.

Figure 5: As Figure 4 , four further noise draws.

Figure 6: As Figure 4 , four further noise draws.

Figure 7: Samples generated by different sampling schemes using the same pretrained SiT-XL/2 backbone with CFG $w=1.5$ for the classes Egyptian cat and volcano . AREX consistently produces samples with higher fidelity, more coherent global structure and finer semantic detail than the other schemes, while retaining these improvements as the integration budget increases from NFE $2$ to $8$ .

Figure 8: As Figure 7 , classes balloon and space shuttle.

Figure 9: As Figure 7 , classes cheeseburger and valley.

Figure 10: SANA-0.6B, MJHQ-30K prompts, $w=4.5$ , every baseline of Table 3 : prompts "raining night, backlight, in a Chinese street, people walking with umbrellas" and "botanical drawing, Psilocybe cubensis , realistic" at NFE 4, 6 and 8, one initial noise per column group.

Figure 11: As Figure 10 , prompts "ultra realistic photo of a dark-blue Toyota RAV4 Hybrid passing by, daylight city background” and comic-style nature scene overlooking a lake with forest, northern-lights colour theme".

Figure 12: As Figure 10 , prompts "loaded nachos with two layers of chili, melted cheese, salsa and jalapeños, hyper realistic" and "a laboratory located in the jungle".

Flow-Based Generative Modeling

Flow-based generative modeling learns probability transports from simple noise to data for fast sampling, density estimation, and controllable generation.

See all 563 papers

Sharpen Your Flow: Sharpness-Aware Sampling for Flow Matching
May 12, 2026Aditi Gupta, Soon Hoe Lim, Annan Yu +1Flow Matching Flow-Based Generative Modeling
Lawrence Berkeley National Laboratory · International Computer Science Institute · 1Lawrence Berkeley National Laboratory, 2International Computer Science Institute +4
Isokinetic Flow Matching for Pathwise Straightening of Generative Flows
Apr 6, 2026Tauhid KhanFlow Matching Flow-Based Generative Modeling
Training-Free Refinement of Flow Matching with Divergence-based Sampling
Apr 6, 2026Yeonwoo Cha, Jaehoon Yoo, Semin Kim +3Flow Matching T2I Generation
KAIST

May 12, 2026J/K move · Enter open · S save

Sharpen Your Flow: Sharpness-Aware Sampling for Flow Matching

Aditi Gupta, Soon Hoe Lim, Annan Yu, N. Benjamin Erichson

Lawrence Berkeley National Laboratory · International Computer Science Institute · 1Lawrence Berkeley National Laboratory, 2International Computer Science Institute+4

Flow Matching Flow-Based Generative Modeling

Open paper AskPDF arXiv

Flow matching models generate samples by numerically integrating a learned velocity field, with each integration step requiring a neural network evaluation. Fast generation therefore requires using a small fixed evaluation budget effectively: the key question is not only how to integrate the flow, but where the sampler should spend its steps. We propose SharpEuler, a training-free sampler that profiles a pretrained model offline by estimating where the learned velocity field changes most rapidly along calibration trajectories. This finite-difference estimate defines a solver-aware sharpness profile, which is smoothed and converted by a quantile transform into a timestep grid for any desired inference budget. At test time, sampling remains ordinary Euler integration with the same number of model evaluations as a uniform schedule. We justify SharpEuler using three principles: a numerical principle identifying trajectory acceleration as the leading source of Euler discretization error, a variational principle deriving sharpness-based power-law timestep densities, and a statistical guarantee showing that the finite-sample calibrated sampler is stable at the terminal distribution level. Our experiments show that SharpEuler improves sample quality at fixed budgets, reducing inter-mode leakage and increasing mode coverage.

Abstract

Flow-Based Generative Modeling

Explore similar work

Sharpen Your Flow: Sharpness-Aware Sampling for Flow Matching

Isokinetic Flow Matching for Pathwise Straightening of Generative Flows

Training-Free Refinement of Flow Matching with Divergence-based Sampling