Multi-Level Monte Carlo Training of Neural Operators

TL;DR

The paper tackles the high computational cost of training neural operators for PDE-related tasks by introducing a multi-level Monte Carlo (MLMC) gradient estimator that leverages a hierarchy of function discretizations. The method expresses the gradient on the finest level as a telescopic sum of coarse-level gradients and level-specific corrections, coupled with a suite of mini-batch strategies and a cache-friendly, multi-resolution data batcher. The authors provide a rigorous mathematical framework, define a function-space hierarchy, and establish optimal sample allocations to achieve a target mean-squared error with reduced cost, complemented by practical MLMC-SGD optimization. Empirically, MLMC training yields substantial speedups (up to about 60%) across Fourier neural operators, MP-PDE, GINOT, and mesh-free CNN architectures, while maintaining or improving accuracy on Darcy flow, Navier–Stokes, cylinder flow, and DeepJEB benchmarks, thereby offering a Pareto-frontier for accuracy versus training time in neural-operator learning.

Abstract

Operator learning is a rapidly growing field that aims to approximate nonlinear operators related to partial differential equations (PDEs) using neural operators. These rely on discretization of input and output functions and are, usually, expensive to train for large-scale problems at high-resolution. Motivated by this, we present a Multi-Level Monte Carlo (MLMC) approach to train neural operators by leveraging a hierarchy of resolutions of function dicretization. Our framework relies on using gradient corrections from fewer samples of fine-resolution data to decrease the computational cost of training while maintaining a high level accuracy. The proposed MLMC training procedure can be applied to any architecture accepting multi-resolution data. Our numerical experiments on a range of state-of-the-art models and test-cases demonstrate improved computational efficiency compared to traditional single-resolution training approaches, and highlight the existence of a Pareto curve between accuracy and computational time, related to the number of samples per resolution.

Figures (8)

• Figure 1: Darcy test problem. Final test loss against the average computational training time per epoch. The baseline plotted as a dotted black line reports the standard FNO performance when trained on increasing data resolutions, while the colored lines highlight a "Pareto curve" of MLMC accuracy/training time trade-off when varying the number of levels $m$ and sampling factor $\delta$.
• Figure 2: Navier-Stokes problem. Same as \ref{['fig:darcy_FNO_pareto_bound']} with a FNO-3d trained on the time-dependent Navier-Stokes equations at increasing mesh resolutions $8\times 8$, $16\times 16$, $32\times 32$, $64\times 64$ (baseline) or with MLMC. The trained models are evaluated on a distinct test set at the finest resolution $64\times 64$. The magenta curve corresponds to MLMC at finest level $32\times 32$ and coarsest level $16\times 16$ with decreasing $\delta$, and yields a better cost/accuracy trade-off compared to the baseline.
• Figure 3: Navier-Stokes problem. Left: Comparison between the ground truth data (1st row), and the MLMC predicted solutions at $t = 40$ with $(m,\delta) = (4,4)$ (2nd row) and $(2,2)$ (3rd row) for three different initial conditions. Right: Evolution of the coarse loss at the lowest resolution during training (top-left) and the pair losses in the telescopic sum expansion of the fine loss with $\delta=2$.
• Figure 4: DeepJEB test problem. (a) Three sample geometries from the DeepJEB dataset at resolution $R_{\text{test}}$, $R_2$, and $R_1$. (b) Comparison between the ground truth data (first column), and the MLMC solutions (second column) for different shapes.
• Figure 5: DeepJEB test problem. Left: Comparison between the norm of the fine gradients, the norm of the MLMC coarse gradients and MLMC gradients approximation over the same batch. Right: Absolute error between the fine gradients, and the gradients of the coarse and telescopic sum given by \ref{['eq:MLMC_batch']} over the same batch.

Theorems & Definitions (3)

• Proposition 2.1
• Remark 2.2
• proof : Proof of \ref{['prop:MLMC']}