Title: OPERATOR LEARNING USING WEAK SUPERVISION FROM WALK-ON-SPHERES
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Neural operator, walk on spheres, poisson solver, weak supervision, monte-carlo method

Abstract: Training neural PDE solvers is often bottlenecked by expensive data generation
or unstable physics-informed neural network (PINN) that involves challenging
optimization landscapes due to higher-order derivatives. To tackle this issue, we
propose an alternative approach using weak supervision from stochastic processes
to produce training data of varying quality. Specific to Poisson PDEs, an efficient
Monte-Carlo algorithm called Walk-on-Spheres (WoS) is capable of generating
solutions using efficient random walks. We introduce a learning scheme called
Walk-on-Spheres Neural Operator (WoS-NO) using weak supervision from WoS to
train any given neural operator. The central principle of our method is to amortize
the cost of Monte Carlo walks across the distribution of PDE instances. Our
method leverages stochastic representations using the WoS algorithm to generate
cheap, noisy, yet unbiased 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. Leveraging the unbiased
nature of these estimates, the operator learns a generalized solution map for an
entire family of PDEs. This strategy results in a mesh-free framework that operates
without expensive pre-computed datasets, avoids the need for 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
upto 8.75× improvement in L2-error compared to standard physics-informed
training schemes, upto 6.31× improvement in training-speed, and reductions of up
to 2.97× in GPU memory consumption.