Martingale Neural Operators

Code for Martingale Neural Operators: Learning Stochastic Marginals via Doob-Meyer Factorization (arXiv:2605.15806).

Neural operators are strong deterministic surrogates for PDEs, but they collapse to the conditional mean on stochastic problems and lose variance and tail structure. Martingale Neural Operators (MNO) encode the Doob-Meyer decomposition—predictable drift plus zero-mean martingale innovation—as an architectural prior: a drift head predicts the conditional mean, and a low-rank basis (B_\phi) with (\Sigma = B_\phi^\top B_\phi) (positive semi-definite by construction) parameterizes uncertainty. This repository implements MNO in 1D and 2D, the training objective, simulators, baselines, and the full experiment suite from the paper.

Installation

Requires Python 3.13+ and uv. From the repository root:

uv sync

Modes are lite (smoke test, ~minutes), standard (paper-scale, ~hours on GPU), and full (denser ablations, ~6h on a Blackwell 6000). Results in experiments/results/; figures in experiments/figures/.

Run the full suite

# Core paper benchmarks (Burgers, rough vol, generative, reaction-diffusion, phi^4)
uv run python experiments/run_all.py standard --group repro_core

# All experiments + plots
uv run python experiments/run_all.py standard

# Theory propositions only
uv run python experiments/run_all.py standard --group theory

Use --skip-plots to skip figure regeneration, or --experiments burgers,phi4 to run a subset. Logs are written to experiments/results/logs/.

Regenerate figures from saved JSON

uv run python -m experiments.plot_results

# Or export a LaTeX summary table from saved results:
uv run python -m experiments.exp_summary

Precomputed metrics under experiments/results/ and figures under experiments/figures/ are included so plots can be reproduced without retraining.

Using MNO in code

import torch
from src.mno import MartingaleNeuralOperator
from src.training import MNOLoss, MNOTrainer

model = MartingaleNeuralOperator(
    in_channels=1,
    out_channels=1,
    width=48,
    modes=16,
    n_layers=4,
    rank=16,
    noise_type="gaussian",
)
u0 = torch.randn(8, 1, 64)  # batch, channels, grid

mean = model(u0, t=1.0)                    # conditional mean
samples = model.sample(u0, t=1.0, n_samples=32)
mean, var = model.get_moments(u0, t=1.0)   # mean and diagonal variance

2D problems use src.mno_2d.MartingaleNeuralOperator2d with the same decomposition on spatial grids.

Repository layout

src/
  mno.py              # 1D MartingaleNeuralOperator
  mno_2d.py           # 2D variant
  training.py         # MNOLoss, MNOTrainer
  baselines.py        # FNO and deterministic baselines
  probabilistic_baselines.py
  simulators.py       # SPDE / SDE simulators
  evaluation.py       # Wasserstein, coverage, metrics

experiments/
  case_*.py           # Application benchmarks (Burgers, phi^4, Gray-Scott, …)
  prop_*.py           # Theory checks (Props 0-4)
  config.py           # LITE / STANDARD / FULL presets
  suite_registry.py   # Experiment metadata and groups
  run_all.py          # Batch runner
  plot_results.py     # Paper figures from JSON results

scripts/
  setup_blackwell.sh
  setup_l40.sh

Checkpoints are saved under experiments/models/ (gitignored). Set MNO_RUNTIME_PROFILE=blackwell on supported GPUs for TF32 and memory tweaks (see experiments/runtime.py).

Experiment registry

Key	Script	Description
burgers	case_burgers	Stochastic Burgers
phi4	case_phi4	phi^4 field theory
rough_vol	case_rough_vol	Rough volatility SDEs
generative	case_generative_sde	Generative efficiency vs diffusion
reaction_diffusion	case_reaction_diffusion	Reaction-diffusion blow-up
gray_scott	case_gray_scott	2D Gray-Scott
turbulent_flow	case_turbulent_flow	2D turbulent flow
resolution_2d	case_resolution_2d	2D zero-shot resolution
prop_0 … prop_4	prop_*	Martingale verification, mirror, resolution, identifiability, uncertainty
burgers_ablations	case_burgers_ablations	Rank / backbone / loss ablations

See experiments/suite_registry.py for baseline lists and EXPERIMENT_GROUPS.

Citation

@article{hidajat2026mno,
  title   = {Martingale Neural Operators: Learning Stochastic Marginals via Doob-Meyer Factorization},
  author  = {Kai Hidajat},
  journal = {arXiv preprint arXiv:2605.15806},
  year    = {2026},
  url     = {https://arxiv.org/abs/2605.15806}
}