NAGO: Noise-Aware Generative Operator via Flow Matching
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 noise measure as a design variable rather than a fixed default. While existing generative approaches mainly shape the sampling process through interpolation schedules or guidance heuristics, NAGO directly controls the marginal velocity field in functional flow matching by choosing a noise measure with suitable regularity. We establish conditions under which the resulting noise-induced flows are well defined in infinite-dimensional settings, and show how noise regularity affects the spectral geometry and stiffness of the marginal velocity field. This analysis leads to a principled design rule: match the regularity of the noise measure to that of the target invariant measure. The resulting method enables stable training and fast sampling without relying on downstream correction mechanisms. Experiments on three representative SPDEs with multiscale invariant measures show that NAGO consistently improves the accuracy--efficiency trade-off over standard flow matching baselines, while achieving up to two orders of magnitude speedup compared with classical numerical solvers.