NAGO: Noise-Aware Generative Operator via Flow Matching

Guanyu Chen; Shengze Xu; Pengwei Liu; Xingyu Ren; Pengkai Wang; Antonios Armaou; Tieyong Zeng; Dong Ni

cs.LGOct 28, 2026

NAGO: Noise-Aware Generative Operator via Flow Matching

Authors: Guanyu Chen, Shengze Xu, Pengwei Liu, Xingyu Ren, Pengkai Wang, Antonios Armaou, Tieyong Zeng, Dong Ni

Organizations: Zhejiang University · The Chinese University of Hong Kong · The Hong Kong Polytechnic University · The Pennsylvania State University

Abstract

We introduce NAGO, a noise-aware generative operator for efficient sampling from invariant measures of stochastic partial differential equations (SPDEs) in function space. Our central idea is to treat the source measure as a design variable rather than a fixed Gaussian default. We establish admissibility conditions for well-defined infinite-dimensional flows and show that the finite-band source--target covariance ratio controls the spectral stiffness of the marginal velocity field. This analysis yields an admissibility-constrained spectral-matching rule: choose a source that is sufficiently rough for well-posedness, then minimize covariance mismatch over the resolved Fourier modes. On three representative stochastic-field benchmarks, this source design improves few-step spectral accuracy and several complementary distributional and physical statistics. In a controlled Matern-field ablation, it reduces the accurate RK4 budget from about 30 steps for Gaussian functional flow matching to 3--5 steps, a $6$ -- $10\times$ reduction in neural ODE evaluations. These gains are strongest for stationary multiscale fields and depend on the metric, resolution, and guidance protocol.

Figures & tables

Figure 1: The structure of the NAGO model.

	Matérn	Gaussian	White
Smoothness $\alpha$	$>d/2$	$0$	$0$
Range parameter $\kappa$	$>0$	–	–
Normalized	–	Yes	No

Table 1: Comparison of different random fields.

Figure 2: Samples under various noise distributions and the target distribution, their spectra, and the marginal velocity fields at different frequencies for each noise type.

Figure 3: Noise-source comparison on the Matérn-field (top) and SNS (bottom) benchmarks. Columns show representative samples or spectra, error versus solver budget, cross-resolution spectra, and guided reconstruction loss, respectively.

Noise	Model/Steps	Matérn Field $(\times 10^{-3})$			Stochastic Allen–Cahn $(\times 10^{-3})$			Stochastic NS $(\times 10^{-2})$
Noise	Model/Steps	C1(5)	C2(10)	C3(50)	C1(5)	C2(10)	C3(50)	C1(5)	C2(10)
Matérn	FNO	$\mathbf{1.70}$	$\mathbf{0.638}$	$\mathbf{1.44}$	$\mathbf{0.553}$	$\mathbf{0.778}$	$\mathbf{0.612}$	$\mathbf{4.44}$	$\mathbf{4.40}$
Matérn	UNet	$312$	$342$	$275$	$30.7$	$32.1$	$32.7$	$7.27$	$7.55$
Gaussian	FNO	$113$	$5.66$	$1.52$	$7.16$	$2.31$	$0.899$	$5.11$	$9.03$
Gaussian	UNet	$2860$	$2700$	$2760$	$7.99$	$6.03$	$4.84$	$44.7$	$28.7$
White	FNO	$1.77{\times}10^{4}\,(20)$	$8930\,(50)$	$1780\,(100)$	$204$	$5.16$	$1.20$	$5220$	$2670$

Table 2: Log-MSE for different test cases. Values in parentheses indicate the number of ODE integration steps used during inference. Best results are highlighted in bold.

Steps	Matérn + FNO		Gaussian + FNO		Matérn + U-Net		Gaussian + U-Net
Steps	L2	MSE	L2	MSE	L2	MSE	L2	MSE
100	$\mathbf{8.42{\times}10^{-2}}$	$\mathbf{7.27{\times}10^{-3}}$	$1.16{\times}10^{-1}$	$2.83{\times}10^{-2}$	$6.26{\times}10^{-2}$	$4.01{\times}10^{-3}$	$3.10{\times}10^{-1}$	$9.85{\times}10^{-2}$
500	$\mathbf{1.03{\times}10^{-1}}$	$\mathbf{1.08{\times}10^{-2}}$	$1.59{\times}10^{-1}$	$2.62{\times}10^{-2}$	$2.76{\times}10^{-1}$	$7.85{\times}10^{-2}$	$1.52$	$2.37$
1000	$\mathbf{8.70{\times}10^{-2}}$	$\mathbf{7.75{\times}10^{-3}}$	$1.53{\times}10^{-1}$	$2.39{\times}10^{-2}$	$4.66{\times}10^{-1}$	$2.27{\times}10^{-1}$	$1.98$	$4.04$

Table 3: Sampling with guidance for the SNS benchmark.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Matern Fields with different $\nu,\ell$

Figure 5: Covariance of Matern Fields with different

\nu,\ell

Method	Model	Noise
Vanilla FM [ Lipman et al., ]	U-Net	Gaussian
FFM [ Kerrigan et al., 2023 ]	FNO	Gaussian
DiffusionPDE [ Huang et al., 2024 ]	FNO	Gaussian
GAN [ Goodfellow et al., 2014 ]	U-Net	–
VAE [ Kingma and Welling, 2013 ]	U-Net	–
GANO [ Rahman et al., 2022 ]	FNO	–

Appendix

Table 4: Summary of baseline methods.

Panel of Figure 6 — Figure 6: Qualitative comparison across three cases: Matérn, SAC, and SNS. The first row shows ground-truth samples, and the second row shows samples generated by NAGO .

Panel of Figure 7 — Figure 7: Matern Field generated at different resolutions. Samples generated at resolution 64 (1st column), 32 (2nd column), 16 (3rd column) with (First row) NAGO and (Second row) FFM

Panel of Figure 8 — Figure 8: Matern Field generated at different resolutions.

Panel of Figure 9 — Figure 9: Matérn cases (Generation with guidance). (a) Ground Truth. (b) Low resolution with 25% observation. Mean, standard deviation, and reconstruction MSE loss of 200 samples generated with (c, d, e) NAGO (f, g, h) gaussian noise.

Panel of Figure 10 — Figure 10: SNS cases (Generation with guidance). (a) Ground Truth. (b) Low resolution with 25% observation. Mean, reconstruction absolute error, and standard deviation of 200 samples generated with (First row) NAGO (Second row) gaussian noise.

$\nu_{0}$	0.5	1.0	1.5	2.0	3.0	6.0
lMSE	1.90e-3	2.91e-3	2.92e-3	8.94e-3	2.28e-2	6.78

Appendix

Table 5: Source-smoothness sweep on the Matérn-field target (lower is better).

Source	$\ell_{0}=0.01$	$\ell_{0}=0.05$	$\ell_{0}=0.10$	$\ell_{0}=0.20$	Gaussian
lMSE	2.995e-3	5.520e-4	1.134e-3	2.241e-3	5.888e-3

Appendix

Table 6: Source length-scale sweep at

\nu_{0}=0.5

, with a discretized Gaussian reference.

Steps	3	5	10	20	30	50
Gaussian	0.2673	0.09294	0.006363	0.005474	0.002138	0.002232
Matérn	0.001339	0.000621	0.000688	0.000850	0.001247	0.001112

Appendix

Table 7: Controlled RK4 sweep on the

64^{2}

Matérn target (lMSE, lower is better).

Method	MMD 2	std. W1	JS	$S_{2}$	$S_{4}$	KE	Enstrophy	Palinstrophy	Vort. corr.
Gaussian (10)	0.003676	0.07487	0.002247	0.1895	0.3506	0.1885	0.1786	0.2262	0.01711
NAGO (5)	0.001418	0.04575	0.000798	0.1450	0.2768	0.1330	0.1095	0.1470	0.01977

Appendix

Table 8: SNS distributional errors (lower is better).

Source	Angle error ( ∘ )	Ratio error	2D power L1	Log-power MSE
Gaussian	0.1310	0.01192	0.04611	0.002929
Isotropic Matérn	0.0815	0.01170	0.03955	0.002563

Appendix

Table 9: Stationary anisotropic Matérn target with 1,024 samples and ten RK4 steps.

Figure 11: Qualitative samples for the stationary anisotropic Matérn target. Rows show unpaired draws from the target distribution, the Gaussian-source model, and the isotropic Matérn-source model. All panels use the same symmetric color scale.

Panel of Figure 12 — Figure 12: Comparison of guided losses under different noise settings.

Panel of Figure 13 — Figure 13: Comparison of guided losses under different noise settings.

Steps	Source	Observed MSE	Mean MSE	Spectral lMSE
3	Gaussian	1.136	0.510	0.2665
3	Matérn	1.817	0.907	0.00259
100	Gaussian	0.977	0.450	0.02620
100	Matérn	1.741	0.893	0.07456

Appendix

Table 10: Fixed-total-guidance control on the Matérn target (lower is better).

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

Deep operator learning for efficient sampling from invariant measures of stochastic differential equations
Sep 10, 2026Ling Guo, Lei Li, Jingtong ZhangStochastic Differential Equations Diffusion Sampling
Department of Mathematics, Shanghai Normal University, Shanghai, China · School of Mathematical Sciences, Shanghai Jiao Tong University, Shanghai, China · Institute of Natural Sciences, MOE-LSC, Shanghai Jiao Tong University, Shanghai, China
SGNO: Spectral Generator Neural Operator for Long-Horizon PDE Prediction
Feb 21, 2026Jiayi Li, Penghao Jiang, Baiyu Chen +5PDE Surrogate Modeling Neural Operators
Guided Flow Matching for Forward and Inverse PDE Problems with Sparse Observations: Algorithm and Theory
May 25, 2026Xifeng Zhang, Jin ZhaoFlow Matching PDE Solving
School of Mathematical Science Capital Normal University Beijing, 100048, China · Academy for Multidisciplinary Studies Capital Normal University2026 Beijing, 100048, China

Sep 10, 2026J/K move · Enter open · S save

Deep operator learning for efficient sampling from invariant measures of stochastic differential equations

Ling Guo, Lei Li, Jingtong Zhang

Department of Mathematics, Shanghai Normal University, Shanghai, China · School of Mathematical Sciences, Shanghai Jiao Tong University, Shanghai, China · Institute of Natural Sciences, MOE-LSC, Shanghai Jiao Tong University, Shanghai, China

Stochastic Differential Equations Diffusion Sampling Operator Learning

Open paper AskPDF arXiv

We introduce an amortized neural sampler that combines operator learning with flow methods for sampling. It maps SDE coefficient functions to pushforwards from a reference measure to the invariant measures, enabling efficient sampling across families of stochastic differential equations. Our framework shifts traditional sampling cost to an initial training phase, after which new SDE instances require only one encoder pass and a few ODE solver steps, independent of mixing time. To handle problems in high dimensions, we use Lagrangian trajectory sensors for the coefficient functions and cross attention in the architecture. We also theoretically establish the expressivity and resolution invariance of our framework. Experiments on 1D and 2D SDE families show competitive accuracy with substantial speedups over MCMC in regimes with slow mixing, transfer across sensor counts, and demonstration results on a 64D interacting particle SDE where traditional grid approaches are infeasible.

Abstract

Flow-Based Generative Modeling

Explore similar work

Deep operator learning for efficient sampling from invariant measures of stochastic differential equations

SGNO: Spectral Generator Neural Operator for Long-Horizon PDE Prediction

Guided Flow Matching for Forward and Inverse PDE Problems with Sparse Observations: Algorithm and Theory