PaperlayerView exact revision on arXiv

Source-linked AI summary

Operator Learning Using Weak Supervision from Walk-on-Spheres

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

arXiv:2603.01193v2cs.LG

TL;DR

Neural PDE solvers face costly data generation and unstable physics-informed training from higher-order derivatives. WoS-NO trains neural operators with cheap stochastic weak supervision, achieving stronger generalization and efficiency while reducing GPU memory compared with physics-informed losses.

  • Problem

    Physics-informed neural PDE solvers face optimization challenges from higher-order derivatives, especially when training data are scarce.

  • Method

    WoS-NO trains neural operators without precomputed data by regressing against cheap, unbiased Walk-on-Spheres stochastic solution estimates.

  • Results

    The method achieves stronger generalization and efficiency while reducing GPU memory compared with physics-informed losses.

  • Takeaways & Limitations

    The trained operator performs zero-shot prediction for new PDE instances and geometries without retraining or additional simulations.

  • Takeaways & Limitations

    The method is limited to linear elliptic Poisson-family PDEs and struggles with untrained boundary conditions.

Abstract

from arXiv · show

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 \emph{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

1. Introduction

WoS-NO trains neural operators without precomputed PDE solutions by regressing against cheap, unbiased Walk-on-Spheres estimates, avoiding derivative-based physics losses while learning generalized solution maps. The framework amortizes Monte Carlo costs across PDE families and supports zero-shot prediction on new boundary conditions, coefficients, and geometries.

  • Background: WoS estimates elliptic PDE solutions through Brownian-motion random walks and their interactions with domain boundaries, avoiding grid discretization complexities.The method is particularly associated with Poisson-family equations and provides a grid-free alternative for complex geometries.
  • Problem scope: The analysis focuses on linear elliptic, specifically Poisson-family, PDEs because of WoS’s geometric robustness and computational relevance.Poisson equations are also computational primitives in applications including pressure projection for incompressible fluid dynamics.
  • Our Approach: WoS-NO replaces precomputed solver data and derivative-based physics losses with regression against cheap, unbiased stochastic estimates from WoS.This data-free objective avoids higher-order derivatives and associated optimization challenges while reducing computational and memory costs.
  • Our Approach: The framework amortizes Monte Carlo walks across PDE instances, learning to denoise weak signals and converge toward the true solution operator for Poisson PDE families.The stochastic training process shares walk costs across an entire distribution of PDE instances.
  • Zero-Shot Generalization: The trained operator predicts new PDE instances with different boundary values, coefficient functions, or geometries in one forward pass without retraining or additional simulations.The introduction also reports zero-shot generalization to surface Laplace inpainting and fluid simulations.

2. Background

The background motivates WoS-NO by combining grid-free Monte Carlo representations for Poisson equations with discretization-agnostic neural operators. It identifies the need for data-free, computationally stable training and describes WoS-based weak supervision through recursive random walks and variance control.

  • Walk-on-Spheres: Grid-free Monte Carlo methods address meshing bottlenecks in complex geometries, high-dimensional spaces, and irregular domains by using probabilistic PDE representations.WoS connects Poisson-equation solutions to expected values of stochastic Brownian motion.
  • Neural Operators: Neural operators learn mappings between infinite-dimensional function spaces through integral kernel operators, enabling discretization-agnostic inference and zero-shot super-resolution.Unlike traditional solvers, they do not depend on fixed mesh resolutions.
  • Walk-on-Spheres: WoS recursively advances random walks from intermediate points, incorporating source contributions until the boundary or a maximum step count is reached.For Poisson equations on Ω_T ⊂ R^d with boundary condition u|∂Ω_T = g, the source contribution f(ξ_i) is computed before each jump.
  • Walk-on-Spheres: Variance reduction is achieved by controlling the number of walks L, improving the fidelity of the weak solution estimate.The notation ˆG_1,WoS[a](ξ) denotes a one-trajectory WoS estimate.
  • Motivation: Although MLMC decomposes losses across fidelity levels, it still relies on pre-computed ground truths, motivating a data-free and computationally stable training paradigm.The passage identifies this gap as one addressed by WoS-NO.

3. Problem Setting

The paper formulates a parameterized family of elliptic Poisson problems on variable domains and aims to learn the associated solution operator. Walk-on-Spheres provides stochastic solution estimates that enable weakly supervised operator training across PDE instances.

  • Parameterized PDE family: The problem family consists of Poisson equations with source f ∈ F and Dirichlet boundary data g ∈ B on regular domains Ω_T.The domain Ω_T is open, bounded, connected, and sufficiently regular, with T defining its geometry.
  • Parameterized PDE family: The objective is to learn G: A = T × F × B → U, mapping each PDE specification a = (T, f, g) to its solution u.The learned operator should handle arbitrary combinations of geometry, source, and boundary functions within the parameterized family.
  • Walk-on-Spheres formulation: Walk-on-Spheres reformulates the Poisson solution operator as an expectation over stochastic trajectories evolving through spheres centered at interior points.At each step, the sphere radius is the distance to the boundary, and the walk stops upon exiting the sphere.
  • Weak supervision objective: Using L trajectories, the estimator approximates the ground-truth operator while controlling fidelity through L and supplies unbiased, high-variance weak supervision when L ≤ 10.Early termination and caching accelerate generation but introduce non-zero bias into the practical empirical estimate.

4. Scene Setup

The experiments evaluate WoS-NO on linear Poisson problems with constant or varying coefficients. For varying coefficients, the governing equation is converted to screened form using delta tracking, which yields a corresponding WoS loss.

  • Problem classes: The evaluation covers two problem classes: linear Poisson equations with constant coefficients and with spatially varying coefficients.The varying-coefficient problems are posed on irregular domains.
  • Problem classes: WoS is directly suited to linear and screened Poisson PDEs, motivating the experimental problem setup.The varying-coefficient case is adapted to this framework through a screened reformulation.
  • Constant-coefficient setup: The constant-coefficient setup uses linear Poisson equations on parameterized geometries with radial-basis source functions and periodic boundary terms.The boundary function uses constant, first-harmonic, and second-harmonic angular components with coefficients sampled uniformly from (−1, 1).
  • Varying-coefficient setup: The varying-coefficient PDE includes diffusion, drift, and absorption terms over a three-dimensional domain.The coefficients are α, vω, and σ, respectively.
  • Varying-coefficient setup: Assuming vω = 0, the method applies delta tracking to obtain a screened equation and define a delta-tracking-based WoS loss.The expectation is estimated using a recursive delta-tracking algorithm and an L-trajectory empirical mean.

5. Experiments

Experiments benchmark WoS-NO against data-free neural PDE solvers and traditional numerical methods, showing stronger accuracy, speed, resource efficiency, scalability, and zero-shot generalization. Applications and ablations further assess versatility, while revealing sensitivity to untrained boundary conditions.

  • Baseline comparisons: WoS-NO achieves the lowest L2 error while reducing peak GPU memory and power consumption, and remains faster than PINO and DeepRitz.The comparison uses geometry-aware, multiinstance PDE solvers as baselines.
  • Baseline comparisons: 9.0e−3 L2-error: WoS-NO converges fastest and achieves the lowest error under a fixed 200-minute training budget.This equal-time comparison evaluates Poisson equations with spatially varying coefficients against DeepRitz and PINO.
  • Amortization and inference: 7.2× and 1.9× longer inference time: WoS requires more time than WoS-NO to reach the same L2 error on spatially varying Poisson equations.The comparison measures the time required to match a well-trained WoS-NO’s accuracy.
  • Generalization: 2.1× overall improvement than PINO and 1.59× than DeepRitz: WoS-NO delivers stronger zero-shot inference across three ShapeNet geometries.Under the same time constraint, WoS-NO also achieves 3.73× better performance than WoS.
  • Architecture agnosticism: 100× and 10×: WoS-NO yields the lowest L2 error across GINO, Transolver, and GNOT, outperforming DeepRitz and PINO losses regardless of architecture.This supports the framework’s applicability to any operator architecture tested.
  • Scalability and applications: 2.5e−1 L2-relative error: the pressure-field approximation is obtained orders of magnitude faster than traditional WoS but struggles with untrained boundary conditions.The image-inpainting application likewise achieves competitive accuracy while being orders of magnitude faster than WoS.
  • Scalability and applications: 8× processing-time increase: FEM data processing grows from mesh resolution 12 to 40, whereas WoS-NO uses parallelizable, meshless random walks independent of grid resolution.WoS-NO also has almost no increase in inference time at higher resolution and avoids higher-order derivative computations during training.

6. Future Work · A. Background: Operator Learning

The framework is motivated by operator learning as approximation of mappings between function spaces and uses GINO-based kernel-integral architectures to learn the relevant integral. It also offers a zero-shot formulation for Poisson surface reconstruction, addressing methods that require per-scene optimization and retraining for new inputs.

  • 6. Future Work: Poisson surface reconstruction is identified as a challenging test of PDE-solver generalizability.The passage contrasts this challenge with physics-informed neural fields that often require per-scene optimization.
  • 6. Future Work: The framework can apply to Poisson surface reconstruction in a zero-shot manner by reformulating it as a Poisson problem.This is presented as an alternative to retraining-based generalization for new inputs.
  • A. Background: Operator Learning: A neural operator is a data-driven approximation of mappings between function spaces, written as Gθ: A → U with Gθ ≈ G.Here, G denotes the target mapping defined in equation 3.
  • A. Background: Operator Learning: Neural operators combine pointwise and integral operators through discretized kernel-integral transforms with neural-network-parameterized learnable kernels.The paper uses this formulation to learn the integral presented in Equation (2).
  • A. Background: Operator Learning: The proposed operator uses a GINO-based formulation to approximate the unknown Green’s function as a kernel integral.This connects the operator architecture to the integral representation used by the method.
  • A. Background: Operator Learning: In the discrete setting, the kernel-integral representation becomes a summation.The passage introduces the discrete representation without specifying the full summation expression.
  • A. Background: Operator Learning: Following GINO, the model projects point-cloud inputs onto a uniform latent grid before integrating input functions a = (f, g, T) with an FNO.This describes the main encoding and integration stages of the architecture.
  • A. Background: Operator Learning: A second GNO evaluates the learned integral at validation points over the same spatial domain, which may differ from the input point clouds.The validation points are defined over the desired evaluation set.

B. Derivation of Walk-on-Spheres

The derivation applies Itô’s lemma to reformulate the Poisson equation probabilistically, then constructs a recursive Walk-on-Spheres solution using stopping times, Green’s functions, and Monte Carlo integration.

  • Itô’s lemma is applied to u(Xξ_t) to reformulate the Poisson equation.
  • The stopping time τ is defined as the first exit time of the stochastic process from the domain Ω.
  • Writing the reformulation in expectation form removes the stochastic-integral term because its expectation is zero.
  • Restricting to subdomains Ω0 ⊂ Ω enables the expectation to be rewritten around an initial point ξ0 = ξ and stopping time τ0.
  • The method recursively samples ξ1 ∼ Xξ_τ0, selects a subdomain Ω1 containing ξ1, and defines subsequent stopping times to obtain a recursive solution.
  • The remaining term is computed with Green’s functions and solved through Monte Carlo integrations.

C. Derivation of Delta Tracking · C.1. Definition of terms for Poisson with spatially varying coefficients

The appendix derives Delta Tracking for Poisson equations with spatially varying diffusion and absorption by transforming the PDE into a screened-Poisson formulation and recursively representing its solution. It then specifies the boundary, diffusion, absorption, and source terms used for this problem family.

  • C. Derivation of Delta Tracking: The derivation starts from the variable-coefficient Poisson PDE ∇·(α(x)∇u) − σu(x) = −f(x) in Ω with u(x) = g(x) on ∂Ω.
  • C. Derivation of Delta Tracking: Expanding the second-order term and applying ∇ln(α(x)) = ∇α(x)/α(x) prepares the equation for eliminating its first-order operator.
  • C. Derivation of Delta Tracking: A Girsanov transformation removes the first-order operator, while introducing ¯σ > 0 shifts coefficient heterogeneity into a source term.
  • C. Derivation of Delta Tracking: The transformed source is f′(x, U) = f′(x) + (¯σ − σ′(x))U(x), after which the Feynman-Kac formulation yields the stochastic representation.
  • C. Derivation of Delta Tracking: The solution representation is evaluated recursively, with expectations approximated using the Poisson kernel P¯σ and Green’s function G¯σ.P¯σ handles the first term, while G¯σ handles the second term in the iterative reformulation.
  • C. Derivation of Delta Tracking: The resulting integral is recursively represented by a solution operator G∆ for screened Poisson equations.
  • C.1. Definition of terms for Poisson with spatially varying coefficients: For the spatially varying Poisson problem, the boundary term is defined parametrically with diffusion frequency Φα, and u(x) is set equal to g(x).The specified boundary function combines sine and cosine terms in the spatial coordinates.
  • C.1. Definition of terms for Poisson with spatially varying coefficients: The implementation defines spatial diffusion and absorption coefficients as α(x) and σ(x), and computes f(x) = −α∇2u − ∇α · u + σu from the closed-form solution.The absorption parameters Amax and Amin control the coefficient’s minimum and maximum values.

D. Experimental Settings · D.1. Domain Generation

The experiments generate disc-like linear domains by sampling perturbed polar boundaries and use ShapeNet meshes with signed-distance filtering for complex variable-coefficient geometries. WoS-NO training samples PDE instances and interior points, then regresses neural-operator outputs against batched WoS estimates.

  • D.1. Domain Generation: Linear domains use the polar boundary r(θ) = r0[1 + c1 cos(4θ) + c2 cos(8θ)] with r0 = 1.The coefficients c1 and c2 are uniformly sampled from U(−0.2, 0.2).
  • D.1. Domain Generation: Interior points for linear domains are generated by rejection sampling from an oversampled rectangular grid and retaining candidates inside the boundary.The desired number of points is sampled uniformly from the retained set.
  • D.1. Domain Generation: 0 + x2 1. Candidate points are evaluated against the angle-dependent radius, and only points satisfying the domain criterion are retained.The supplied passage states the criterion as r < r0, with r0 calculated for the given θ.
  • D.1. Domain Generation: Variable-coefficient Poisson experiments use ShapeNet meshes to define complex geometries.The meshes provide the geometric representation for sampling both domain and boundary points.
  • D.1. Domain Generation: A fast signed-distance-field estimate computed with a bounding volume hierarchy identifies domain points inside ShapeNet meshes.Candidate points with negative signed distance are retained as interior points.
  • D.1. Domain Generation: WoS-NO samples PDE inputs and domain points, evaluates batched WoS trajectories, and minimizes mean squared error between operator predictions and WoS estimates.The training loop updates neural-operator parameters using the gradient of the weak-supervision loss.

E. Training details · E.1. Caching prior walks · E.2. Hyperparameters for the models

The study trains three neural-operator architectures on 4,000 PDE instances using sampled domain and boundary data, while caching and progressively refreshing WoS walks to refine weak targets. It specifies separate architectural configurations for GINO, GNOT, and Transolver.

  • E. Training details: The training set contains 4,000 PDE instances over unique geometries and uses GINO, GNOT, and Transolver as architecture-independent neural operators.Each instance supplies 1,024 sampled domain points and 1,024 boundary points, together with pointwise source and varying-coefficient values.
  • E. Training details: The model input comprises 1,024 sampled domain points, 1,024 sampled boundary points, and pointwise source, diffusion, and absorption values when coefficients vary.These quantities are computed from each PDE instance’s sampled geometry.
  • E.1. Caching prior walks: The cache stores intermediate WoS walks and refreshes each epoch with fresh Monte Carlo samples instead of using a fixed noisy target.This amortizes variance reduction over training while adding O(mn) space complexity for m instances and n points per instance.
  • E.1. Caching prior walks: At epoch k, the cached target is the running average of all k · L WoS trajectories generated through that epoch, with L new trajectories per point each epoch.The update is recursive and incorporates the estimate from the L fresh walks generated during epoch k.
  • E.1. Caching prior walks: As training progresses, the cached target’s variance approaches zero and converges to the true solution G[a_j](ξ_i).The empirical loss is evaluated against this progressively refined target.
  • E.2. Hyperparameters for the models: GINO uses embedding dimension 16, up to 600 positional embeddings, graph radii 0.1, 0.1, and 0.05, and a four-layer FNO with [20] Fourier modes.Its hidden, lifting, and projection channels are 64, 256, and 64, with Group Normalization, AdaIn features, and tensor-factorization rank 0.8.
  • E.2. Hyperparameters for the models: GNOT has 12 Transformer layers with hidden dimension 128, single-head linear attention, and a two-expert Mixture-of-Experts feed-forward network without dropout or horizontal Fourier features.Each expert uses a two-layer MLP with inner dimension 8.
  • E.2. Hyperparameters for the models: Transolver uses 6 Transformer layers, hidden dimension 32, 24 attention heads, residual connections, slice number 32, and reference value 16.It excludes time as an input feature and produces a single scalar output.

F. Ablations · F.1. Effects of Control Variates

The ablation studies examine whether control variates improve Walk-on-Spheres Neural Operator performance. Although control variates are intended to reduce estimator variance, no significant improvement is observed because neural operators can learn from noisy inputs.

  • F. Ablations: The ablation section evaluates approaches for improving WoS-NO performance.
  • F.1. Effects of Control Variates: Control variates are presented as a principled method for reducing the variance of Walk-on-Spheres estimates.
  • F.1. Effects of Control Variates: Lower-variance Walk-on-Spheres estimates are expected to produce more consistent results.
  • F.1. Effects of Control Variates: Lower-variance Walk-on-Spheres estimates are also expected to produce more accurate results.
  • F.1. Effects of Control Variates: No significant improvements from control variates are observed in this setting.
  • F.1. Effects of Control Variates: The lack of improvement is attributed to neural operators’ expressivity and ability to learn from noisy inputs.
  • F.1. Effects of Control Variates: Table 3 showcases the effects of control variates on the 2D dataset.

F.2. Effects of Number of Trajectories

WoS-NO was evaluated with 1, 10, 100, and 1000 trajectories, revealing a trade-off between stochastic-estimation fidelity and simulation time. Despite differing trajectory counts, neural operators asymptotically achieve comparable error levels.

  • Trajectory-count evaluation: WoS-NO was trained with L = 1, 10, 100, and 1000 trajectories to assess how trajectory count affects performance.
  • Trajectory-count evaluation: More trajectories reduce variance but increase training time because random-walk simulations are more likely to be longer.
  • Trajectory-count evaluation: The trajectory count creates a fidelity–simulation-time trade-off that requires tuning to balance both costs.
  • Trajectory-count evaluation: Asymptotically, neural operators converge to similar error levels across different numbers of trajectories and stochastic-estimation fidelities.
Loading 2603.01193v2…