Operator Learning Using Weak Supervision from Walk-on-Spheres

Hrishikesh Viswanath; Hong Chul Nam; Xi Deng; Julius Berner; Anima Anandkumar; Aniket Bera

cs.LGMar 1, 2026

Operator Learning Using Weak Supervision from Walk-on-Spheres

Authors: Hrishikesh Viswanath, Hong Chul Nam, Xi Deng, Julius Berner, Anima Anandkumar, Aniket Bera

Organizations: Purdue University, USA · ETH Zurich, Switzerland · California Institute of Technology, USA · NVIDIA, USA

Abstract

Training neural PDE solvers is often bottlenecked by expensive data generation or unstable physics-informed neural network (PINN) involving challenging optimization landscapes due to higher-order derivatives. To tackle this issue, we propose an alternative approach using Monte Carlo approaches to estimate the solution to the PDE as a stochastic process for weak supervision during training. Leveraging the Walk-on-Spheres method, we introduce a learning scheme called Walk-on-Spheres Neural Operator (WoS-NO) which uses weak supervision from WoS to train any given neural operator. We propose to amortize the cost of Monte Carlo walks across the distribution of PDE instances using stochastic representations from the WoS algorithm to generate cheap, noisy, estimates of the PDE solution during training. This is formulated into a data-free physics-informed objective where a neural operator is trained to regress against these weak supervisions, allowing the operator to learn a generalized solution map for an entire family of PDEs. This strategy does not require expensive pre-computed datasets, avoids computing higher-order derivatives for loss functions that are memory-intensive and unstable, and demonstrates zero-shot generalization to novel PDE parameters and domains. Experiments show that for the same number of training steps, our method exhibits up to 8.75 $\times$ improvement in $L_2$ -error compared to standard physics-informed training schemes, up to 6.31 $\times$ improvement in training speed, and reductions of up to 2.97 $\times$ in GPU memory consumption. We present the code at https://github.com/neuraloperator/WoS-NO

Figures & tables

Panel of Figure 1 — Figure 1: Given arbitrary, unseen input geometry and boundary conditions, we compare our method (WoS-NO) with WoS at equal execution time, and with DeepRitz and PINO at equal training time. We visualize the relative absolute error against the analytic solution of the PDE (the ground truth). We show that WoS-NO achieves the strongest performance (lowest relative error) in comparison with other baselines. During equal-time training, WoS-NO achieves 2.1 $\times$ overall improvement than PINO and 1.59 $\times$ than DeepRitz. During inference, WoS-NO achieves 3.73 $\times$ better performance than WoS under the same time constraint.

Figure 2: Weak Supervision Loss with WoS: Our algorithm learns the given family of parametrized Poisson equations $\Delta u=f$ on $\Omega_{T}\subset\mathbb{R}^{d}$ and $u|_{\partial\Omega_{T}}=g$ and is agnostic to underlying neural operator architecture. The WoS method defines the recursive process of the random walk, stopping once the boundary or the maximum number of steps is reached. The source contribution $f(\xi_{i})$ is computed for each intermediate point $\xi_{i}$ before jumping to the next point $\xi_{i+1}$ . We achieve variance reduction by controlling the number of walks $L$ to improve the fidelity of the weak solution. $\hat{\mathcal{G}}_{1,\text{WoS}}[a](\xi)$ denotes the estimation of 1-trajectory WoS estimation. $\xi_{K}$ denotes the termination condition where the boundary value $g(\xi_{K})$ is added if the point is within the tolerance region. The operator estimate is denoted by $\mathcal{G}_{\theta}[a](\xi)$ . We illustrate the overall learning process, with WoS integral serving as the weak supervision for the neural operator.

APPROACH	TRAINING TIME (Min)	AVERAGE MSE ERROR $\pm\text{Std}$	PEAK GPU MEMORY (MB)	PEAK GPU POWER (W)
PINO	$85.25$	$2.5e^{-3}\pm 4.7e^{-3}$	$1523.5$	$78.43$
Deepritz Operator	$35.401$	$1.0e^{-2}\pm 1.6e^{-2}$	$859.43$	$49.52$
WoS-NO	$\mathbf{13.5}$	$\mathbf{8.2e^{-4}\pm 5.4e^{-4}}$	$\mathbf{627}$	$\mathbf{48.81}$

Table 1: Performance and compute comparison of our WoS-NO against data-free baselines on the linear Poisson family. All methods were trained for 20,000 steps, with the evaluation metrics averaged over 1000 unseen linear parameterizations. Our approach achieves the highest accuracy with significantly lower GPU resource consumption.

APPROACH	TRAINING TIME (Min)	AVERAGE MSE ERROR $\pm\text{Std}$	PEAK GPU MEMORY (MB)	PEAK GPU POWER (W)
PINO	$411$	$9.4e^{-3}\pm 1.1e^{-2}$	8587.9	122.85
Deepritz Operator	$198$	$1.1e^{-2}\pm 2.4e^{-2}$	$3903.1$	$116.59$
WoS-NO	$\mathbf{188}$	$\mathbf{9.0e^{-3}\pm 7.9e^{-3}}$	$\mathbf{2886.3}$	$\mathbf{103.90}$

Table 2: Performance and compute comparison of our WoS-NO against data-free baselines on Poisson equations with spatially varying coefficients. All methods were trained for 50,000 steps, with the evaluation metrics averaged over 1000 unseen spatially varying parameterizations. Our approach achieves the highest accuracy with significantly lower GPU resource consumption.

Panel of Figure 3 — Figure 3: Left: Training Poisson equations with spatially varying coefficients with equal time for 200 minutes, WoS-NO demonstrates the lowest $L_{2}$ error while converging the fastest. In contrast, PINO requires a much longer time to converge. Middle: After WoS-NO training is finished, we compare the amount of time needed to achieve the same level of accuracy as $L_{2}$ -error to achieve the same $L_{2}$ error as a well-trained WoS-NO for 4096 pointwise estimations. Right: GINO, Transolver and GNOT are trained on DeepRitz, PINO and WoS-NO losses, and across all three operator architectures, WoS-NO is the strongest with the lowest $L_{2}$ -error.

Figure 4: Left: This graph showcases the inference time of a single instance, as a function of domain size. In the case of FEM, it denotes the resolution of the mesh, while for the WoS-based methods, it’s the number of query points Middle: This figure highlights the amortization achieved with our approach. Meshing time includes time needed to create meshes in the domain. Data preprocessing time includes the simulation time to create the regression target (FEM solution and WoS simulation). Training time represents the time required for the neural operator to learn the data-driven objective from the points sampled from the FEM solution. Our approach does not require meshing, and the total time for WoS-NO is lower than PINO while approximately equal to DeepRitz total time, making it an efficient operator for training. Right: This figure highlights how the data preprocessing time for FEM scales with the resolution of the inputs. Increasing the resolution results in a direct increase in computation time (hours). This is avoided by WoS-NO since WoS is highly parallelizable on GPU and does not depend on the input resolution.

Figure 5: Quantitative comparison for biharmonic inpainting. We evaluate our pre-trained operator against a traditional WoS solver and the scikit-image baseline. The plot shows the final error versus total runtime (log scale, in seconds) required to inpaint 20 masks. The average MSE compared to the ground truth for our method is $2.8e^{-3}\pm 7.2e^{-3}$ , WoS ( $5.7e^{-3}\pm 5.2e^{-3}$ ) and scikit-image ( $5.4e^{-3}\pm 5.0e^{-3}$ ). The total wall clock time for WoS solver is 557.2s, Scikit-Image 6.8s and WoS-NO 2.8s.

Figure 6: Qualitative and quantitative results showing the pressure field magnitude in the von Kármán vortex. ( Left ) Visual comparison of the pressure fields computed by the baseline WoS solver [ 12 ] (top) and our pre-trained operator (bottom). Our method achieves a relative error of $2.5\times 10^{-1}$ with respect to the baseline. ( Right ) Log-scale plot of the total wall-clock time required to compute the pressure solution over 70 simulation time-steps, comparing the computational cost of both methods. The total wall clock time for WoS solver is 1698.2s and WoS-NO 251.1s

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

CONTROL VARIATES	TRAIN L2 Err	TEST L2 Err
No	$4.05e^{-4}$	$8.2e^{-4}$
Yes	$4.33e^{-4}$	$8.35e^{-4}$

Appendix

Table 3: Showcasing the effects of control variates on the 2D dataset.

Figure 7: Left: This figure showcases the overall training time as the number of trajectories changes for each instance. The higher number of trajectories increases the training time due to a higher likelihood of longer random walks. Right: This plot shows comparable performance for different numbers of walks for each instance. Asymptotically, neural operators converge to a similar level of errors regardless of the fidelity of the stochastic estimation.

Explore similar work

On the training of physics-informed neural operators for solving parametric partial differential equations
Jun 4, 2026Nanxi Chen, Chuanjie Cui, Airong Chen +2Neural Network Optimization Gradient Interference
College of Civil Engineering, Tongji University, Shanghai 200092 · Department of Engineering Science, University of Oxford, Oxford OX1 3PJ · Institute for Foundations of Data Science, Yale University, New Haven, CT 06520
Preconditioned Physics-Informed Neural Operator Training
Sep 28, 2026Shizheng Wen, Siddhartha Mishra, Marius ZeinhoferPhysics-Informed Neural Operators PDE Operator Learning
ETH Zürich
Learning Physical Operators using Neural Operators
Feb 26, 2026Vignesh Gopakumar, Ander Gray, Dan Giles +5PDE Surrogate Modeling Physics-Informed Neural Operators
UCL Centre for Artificial Intelligence · UK Atomic Energy Authority · LIX, CNRS, ´Ecole Polytechnique +2

Jun 4, 2026J/K move · Enter open · S save

On the training of physics-informed neural operators for solving parametric partial differential equations

Nanxi Chen, Chuanjie Cui, Airong Chen, Sifan Wang+1

College of Civil Engineering, Tongji University, Shanghai 200092 · Department of Engineering Science, University of Oxford, Oxford OX1 3PJ · Institute for Foundations of Data Science, Yale University, New Haven, CT 06520

Neural Network Optimization Gradient Interference PDE Operator Learning Physics-Informed ML Physics-Informed Neural Operators

Open paper AskPDF arXiv

Physics-informed neural operators (PINOs) aim to learn solution operators for partial differential equations by using the governing physics as supervision, rather than relying solely on paired input-output simulation data. By incorporating physical constraints into the training objective, PINOs combine the cross-instance generalization of neural operators with the data efficiency of physics-informed learning. Despite this promise, how to train PINOs efficiently and robustly remains less well-understood than the training of either data-driven neural operators or physics-informed neural networks (PINNs). To bridge this gap, we examine key components of the PINO training pipeline, including architecture design, optimizer choice, loss balancing, and collocation-point sampling strategy. We study three representative operator backbones, Deep Operator Network (DeepONet), Fourier Neural Operator (FNO), and Continuous Vision Transformer (CViT), across five diverse parametric PDE systems. Our results show that CViT provides consistently strong and stable performance across the considered benchmarks. Beyond architecture, we find that several optimization pathologies previously identified in PINN training naturally arise in PINOs, including gradient conflicts and causal violation. We also find that mitigation algorithms developed for PINNs remain effective in the PINO setting. We further compare physics-informed and data-driven training under different data regimes, revealing that a carefully designed physics-informed training pipeline can match, and in some cases, outperform purely data-driven neural operators. Taken together, these findings provide a systematic empirical understanding of the optimization challenges in PINO training and inform a practical pipeline for efficient and robust physics-informed operator learning. Code and data are available at https://github.com/NanxiiChen/PI-CViT.

Abstract

Explore similar work

On the training of physics-informed neural operators for solving parametric partial differential equations

Preconditioned Physics-Informed Neural Operator Training

Learning Physical Operators using Neural Operators