cs.LGMay 15, 2026

Martingale Neural Operators: Learning Stochastic Marginals via Doob-Meyer Factorization

Authors: Kai Hidajat

Organizations: Department of Applied Mathematics University of Washington Seattle, WA, USA

Abstract

Neural operators excel as deterministic surrogates, but inevitably collapse to the conditional mean when applied to stochastic PDEs, discarding the variance and tail structure upon which uncertainty quantification depends. Recovering this structure typically requires Monte Carlo rollouts or grafted generative models, both of which surrender the one-shot efficiency and resolution invariance that define the operator paradigm. To resolve this, we draw on the Doob-Meyer theorem, which establishes that any semimartingale fundamentally decomposes into a predictable drift and an unpredictable, zero-mean martingale. Translating this theorem into an architectural prior, we introduce the Martingale Neural Operator (MNO). MNO maps an initial condition directly to the conditional mean and covariance of the terminal law, parameterized by a drift-like mean and a low-rank factor BφB_φ with Bφ⊤BφB_φ^\top B_φ positive semi-definite by construction. For our experiments, we use a Gaussian residual instantiation. Across 1D SPDEs, rough volatility, and 2D operator tasks, MNO reduces Wasserstein distance by up to 120×120\times on φ4φ^4 field theory and 68×68\times on stochastic Burgers, evaluating ∼3×\sim 3\times faster than a conditional diffusion baseline at matched wall-clock training budgets. On 2D tasks, MNO is comparable to FNO on zero-shot resolution transfer and turbulent flow, while quasi-deterministic systems such as Gray-Scott remain a failure mode.

Explore similar work

Oct 18, 2024math.OC

Polynomial Scaling is Possible For Neural Operator Approximations of Structured Families of BSDEs

Neural operator (NO) architectures learn nonlinear maps between infinite-dimensional function spaces and are widely used to accelerate simulation and enable data-driven model discovery. While universality results ensure expressivity, they do not address \emph{complexity}: for broad operator classes described only through regularity (e.g.\ uniform continuity or CrC^r-regularity), information-theoretic lower bounds imply that minimax-optimal NO approximation rates scale \emph{exponentially} in the reciprocal accuracy 1/ε1/\varepsilon. This has shifted the focus of NO theory toward identifying additional problem-specific structure, beyond regularity, under which suitably tailored NO architectures can leverage to unlock polynomial scaling in 1/ε1/\varepsilon. We exhibit the first polynomial-scaling regime for NO approximations of solution operators in stochastic analysis; by identifying structured families of \emph{non-Markovian} BSDEs with randomized terminal condition parameterized by the Sobolev-regular terminal condition and by Sobolev-regular additive nonlinear perturbations of the generator. We prove that their solution operator can be approximated (uniformly over the family) by a tailored NO whose number of trainable parameters grows \emph{polynomially} in 1/ε1/\varepsilon. We unlock this polynomial scaling regime by \emph{informing the NO's inductive bias} by factoring out the singular part of the associated semilinear elliptic PDE Green's function and by incorporating the Doléans--Dade exponential of the BSDE's common non-Markovian factor into the NO's decoding layers. As a byproduct, we extend polynomial-scaling guarantees from families of linear elliptic PDEs on regular domains to the semilinear setting.
Takashi Furuya, Anastasis Kratsios
May 16, 2026stat.ML

Diffusion-Based Stochastic Operator Networks for Uncertainty Quantification in Stochastic Partial Differential Equations

We introduce a novel framework for uncertainty quantification of solution operators associated with stochastic partial differential equations (SPDEs). Although SPDEs play a central role in modeling complex physical systems under uncertainty, their practical use typically requires specifying the magnitude and structure of model uncertainties that are often unknown and difficult to infer from noisy measurements. To address this challenge, we develop a stochastic operator-learning framework that learns directly from noisy data and outputs both a mean solution field and a quantification of uncertainty. The proposed method, namely the Stochastic Operator Network (SON), is constructed by combining the structure of the Deep Operator Network (DeepONet) with Stochastic Neural Networks (SNNs) to model stochasticity and enable probabilistic prediction. The training procedure is carried out by minimizing a Hamiltonian-type loss and optimizing the resulting objective using the Stochastic Maximum Principle. Numerical experiments on benchmark SPDEs under multiple uncertainty sources demonstrate the accuracy and robustness of the proposed method in capturing solution structure and quantifying predictive uncertainty.
Phuoc-Toan Huynh, Richard Archibald, Feng Bao
Jun 22, 2026cs.LG

Neural Operator Processes for Probabilistic Operator Learning under Partial Observations

Neural operators learn mappings between function spaces, but are typically developed with dense input-output training fields and fully observed inputs at inference. Many scientific problems require instead predicting solution fields from sparse, irregular, or partial observations under uncertainty. We introduce Neural Operator Processes (NOPs), a framework that unifies neural-process conditioning with neural-operator decoding to predict full output fields from limited context. NOPs condition on sparse joint input-output observations and support deterministic and probabilistic prediction within a shared encoder-decoder architecture. We study two conditioning strategies, convolutional pooled summaries and query-aligned attention, and analyze how their interaction with latent stochastic variables depends on PDE geometry. Across function regression and three PDE benchmarks, we find that sparse conditional operator learning is viable and can match dense-grid behavior in several regimes, that preserving local context-query geometry is essential in non-periodic settings but less so in spectrally smooth periodic regimes, and that uncertainty-aware operator learning succeeds when latent conditioning complements rather than overwrites the local geometric pathway. These results provide a basis for probabilistic operator learning under partial observations and help bridge operator learning and probabilistic meta-learning in function space.
Jose Miguel Lara-Rangel, Serge Guillas