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

Published 15 May 2026 in cs.LG | (2605.15806v1)

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_φ$ with $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\times$ on $φ^4$ field theory and $68\times$ on stochastic Burgers, evaluating $\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.

Abstract PDF Upgrade to Chat

Authors (1)

Kai Hidajat

Summary

The paper introduces Martingale Neural Operators (MNO) that decompose terminal marginals of stochastic PDEs using the Doob–Meyer factorization.
It parameterizes both the drift and conditional covariance via specialized FNO stacks, achieving significant reductions in Wasserstein-2 error.
The approach enables one-shot uncertainty quantification with grid-invariant performance across benchmarks like stochastic Burgers' and rough volatility models.

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

Abstract and Motivation

This paper introduces the Martingale Neural Operator (MNO), a functional neural operator architecture for learning the terminal marginal laws of stochastic PDEs. The approach leverages the Doob–Meyer decomposition, which splits a semimartingale into a predictable drift and an unpredictable, zero-mean martingale. Standard deterministic neural operators, such as FNO, achieve high performance in deterministic PDE regimes but collapse to conditional mean predictions under $L^2$ loss when confronting stochastic PDEs, ignoring the stochastic structure essential for uncertainty quantification (UQ). The MNO architecture addresses this by parameterizing both conditional mean (drift) and conditional covariance (residual stochasticity), thereby enabling one-shot uncertainty quantification with a computational complexity and grid-invariance on par with contemporary neural operator methods.

Architectural Formulation

Doob–Meyer-Inspired Parameterization

MNO decomposes its prediction of the terminal state $u_T$ given initial condition $u_0$ as:

$u_T \approx m_\theta(u_0, T) + B_\phi(u_0, T)^\top \xi,$

where $m_\theta$ is the drift-like mean operator and $B_\phi$ is a learned low-rank factor such that $\Gamma_\phi = B_\phi^\top B_\phi$ is positive semidefinite by construction; $\xi \sim \mathcal{N}(0, I_r)$ is a standard Gaussian of rank $r$ . Both $m_\theta$ and $u_T$ 0 are parameterized via independent or shared FNO stacks, with a temporal gating mechanism ensuring initial-value consistency.

This parameterization gives closed-form expressions for the conditional mean and variance:

$u_T$ 1

Theoretical Properties

MNO satisfies several key architectural guarantees:

Drift approximation: Under standard FNO universality, the mean operator can approximate continuous solution operators on bounded domains.
Covariance modeling: Any positive semidefinite covariance operator of rank at most $u_T$ 2 admits a factorization of the MNO form; truncation to $u_T$ 3 yields standard spectral error.
Resolution invariance: Since the architecture operates on function spaces rather than discretizations, zero-shot evaluation over arbitrary discretization is naturally supported.

The architecture enforces zero initial residuals, zero-mean sampling, and channel-aware variance modeling via the factor $u_T$ 4. However, the expressivity of the residual law is tied to the rank and chosen residual distribution (Gaussian in this work).

Empirical Results

Stochastic Burgers' Equation

On the classical stochastic Burgers' equation, MNO dramatically improves marginal law prediction as quantified by Wasserstein-2 ( $u_T$ 5) distance:

Figure 1: Stochastic Burgers' Equation results—MNO recovers a nontrivial residual variance profile while matching deterministic baselines for mean predictions.

$u_T$ 6 is reduced from 0.6483 (Neural SPDE baseline) to 0.0095, signifying a 68-fold improvement.
Mean RMSE remains essentially unchanged compared to deterministic FNO.

Rough Volatility Dynamics

MNO demonstrates robustness to non-semimartingale inputs, such as rough volatility models with Hurst parameter $u_T$ 7, where sequential baselines are challenged.

Figure 2: Rough volatility results—MNO targets terminal marginals and remains accurate as $u_T$ 8, outperforming sequential baselines in the deeply rough regime.

For $u_T$ 9, MNO achieves mean $u_0$ 0 vs. $u_0$ 1 for Neural SDE, a 2.6-fold improvement in terminal marginal accuracy.
The architecture targets marginals rather than pathwise roughness.

$u_0$ 2 Field Theory Benchmark

Figure 3: $u_0$ 3 field theory benchmark (SPDEBench)—MNO achieves $u_0$ 4 vs. $u_0$ 5 for Neural SPDE by utilizing low-rank residual factors.

MNO achieves an unprecedented 120-fold reduction in $u_0$ 6 error (0.0055 vs. 0.6572).

Super-Resolution and Resolution Transfer

Figure 4: 1D zero-shot super-resolution—MNO, trained at coarse resolution, maintains mean RMSE on finer grids, whereas Neural SPDE performance degrades.

MNO generalizes zero-shot to higher resolutions, maintaining accuracy while Neural SPDE's error increases substantially.

Generative Efficiency

Figure 5: Generative efficiency—one-shot moment computation via MNO is $u_0$ 7 faster than diffusion models and yields lower $u_0$ 8 for the same compute budget.

Single-pass moment prediction is $u_0$ 9 vs. $u_T \approx m_\theta(u_0, T) + B_\phi(u_0, T)^\top \xi,$ 0 for diffusion (NFE=25), with superior $u_T \approx m_\theta(u_0, T) + B_\phi(u_0, T)^\top \xi,$ 1.
MNO provides practical computational advantages for UQ tasks over sample-based or iterative generative approaches.

2D Turbulent Flow

Figure 6: 2D turbulent flow—MNO's mean and variance predictions are competitive with FNO, though CNNs may outperform both on mean RMSE in some regimes.

MNO matches FNO in mean RMSE on channel/spatial benchmarks; state-of-the-art is not claimed for 2D mean recovery.
Figure 7: 2D resolution transfer—MNO preserves performance across resolutions without degradation, maintaining parity with FNO.
CNN and spectral baselines lose accuracy rapidly at different grid sizes; MNO maintains grid-invariant behavior.

Additional Structural Diagnostics

The residual head empirically centers and scales as desired (Figure 8).
AR reuse (rolled-out MNO) reveals that pathwise non-Markovian structure is not preserved, establishing that MNO is a terminal-marginal estimator and not a rough path generator (Figure 9).
Resolution transfer retains low variance error at high grid resolutions compared to baseline stochastic operator surrogates (Figure 10).
Drift/covariance head separation is achieved in synthetic settings (Figure 11).
Aleatoric uncertainty predicted by MNO is closely aligned with empirical statistics (Figure 12).

Implications and Limitations

Practical implications: MNO enables deterministic, one-shot uncertainty quantification in operator learning for stochastic PDEs. This is operationally transformative for applications prioritizing terminal risk envelopes, safety constraints, or high-frequency UQ at scale (e.g., in scientific computing, finance, and engineering). Grid-invariant performance further facilitates deployment in variable-resolution and multi-fidelity settings.

Theoretical implications: By structuring operator learning via Doob–Meyer, the approach integrates stochastic analysis priors into functional neural modeling, enabling architectures that are closer to the laws of SPDEs and their physical decomposition. The division between mean and residual spans both semimartingale regimes and broader terminal-marginal modeling in non-Markovian or rough dynamics.

Limitations and open directions:

MNO is restricted to marginal prediction and does not model the full trajectory nor enforce pathwise temporal consistency.
The residual distribution is Gaussian and low-rank in all experiments; extensions to heavy-tailed or multimodal noise would require modified objectives and architectures.
Empirical calibration is limited to marginal moments; full covariance structure matching (off-diagonal accuracy) is weaker, as the architecture targets variance consistency but not cross-covariance explicitly.
For highly deterministic, pattern-forming PDEs (e.g., Gray-Scott at low noise), the added variance head can impair mean-field prediction.

Future directions include extending residual expressivity (non-Gaussian/multimodal factors), pursuing path-consistent functional extensions, explicit Chapman-Kolmogorov or filtration-level regularization, and compositional MNO constructions for multi-step UQ or sequential decision processes.

Conclusion

Martingale Neural Operators instantiate a theoretically principled and empirically validated framework for deterministic uncertainty quantification in stochastic operator learning. The architecture is distinguished by its direct mean/covariance factorization grantable via the Doob–Meyer theorem, computational efficiency, and grid-invariant end-to-end learnability. The resultant advances in $u_T \approx m_\theta(u_0, T) + B_\phi(u_0, T)^\top \xi,$ 2 metrics versus baselines validate the efficacy of the approach for terminal moment estimation in a variety of SPDEs, with substantial practical leverage for UQ-centric workflows.