Probabilistic Operator Learning Framework

• Probabilistic Operator Learning Framework is a method that models mappings between infinite-dimensional function spaces with intrinsic stochastic noise.
• It employs stochastic processes via SDE-driven branch networks and leverages the stochastic maximum principle with adjoint BSDEs for efficient training.
• This approach enables accurate surrogate modeling for stochastic PDEs and noisy systems by capturing both mean behavior and pointwise uncertainty.

A probabilistic operator learning framework provides a mathematically rigorous and computational approach for learning mappings between infinite-dimensional function spaces where intrinsic stochasticity, noise, or uncertainty is present. These frameworks are central for scientific computing, surrogate modeling of stochastic or random partial differential equations (PDEs), and anywhere the underlying operator exhibits non-negligible variability in its input-output relations. Recent developments, such as the Stochastic Operator Network (SON) (Bausback et al., 10 Jul 2025), address the direct quantification and propagation of operator uncertainty—eschewing purely deterministic architectures in favor of models that learn and predict both mean behavior and pointwise uncertainty.

1. Motivation and Conceptual Foundations

Classical operator learning methods such as DeepONet and Fourier Neural Operator (FNO) are deterministic; their predictions are entirely characterized by point solutions even when the true solution operator is stochastic (e.g., solutions to SDEs/SPDEs, data corrupted by process noise, stochastic parameter dependencies). Previous uncertainty quantification (UQ) solutions—Bayesian neural operators, Monte Carlo dropout, information bottleneck, or ensembles—tended to focus on epistemic uncertainty (model parameters) or exogenous input noise, rather than capturing intrinsic "operator noise," i.e., stochasticity that is a fundamental property of the mapping itself.

Probabilistic operator learning frameworks systematically lift operator learning from deterministic function approximation to the level of learning function-valued random mappings, providing direct access to the conditional distribution of the output function, not merely its expectation.

2. Stochastic Operator Network (SON): Operator Networks with SDE-driven Branches

The Stochastic Operator Network (SON) (Bausback et al., 10 Jul 2025) exemplifies this shift by making the branch network of the classical DeepONet architecture a stochastic process—specifically, the solution to a stochastic differential equation (SDE):

$$A_{n+1} = A_n + h\,\mu(A_n, \theta_n) + \sqrt{h}\,\sigma(\theta_n) \omega_n$$

where:

• $A_n$ is the neuron activation at layer $n$
• $\mu(\cdot)$: trainable drift (deterministic part)
• $\sigma(\cdot)$: trainable diffusion (models stochasticity, i.e., uncertainty)
• $\omega_n$: i.i.d. standard normal variables

In the continuous limit $h \to 0$:

$$A_T = A_0 + \int_0^T \mu(A_t, \theta_t) dt + \int_0^T \sigma(\theta_t) dW_t$$

where $W_t$ is Brownian motion.

This SDE-driven representation causes the operator's output to be an explicit random variable, with the learned diffusion $\sigma(\theta_t)$ parameterizing pointwise uncertainty.

The SON operator prediction at location $y$ is: $$\hat{G}_{SON}(u)(y) = \langle\,\beta(u; \theta_\beta, \epsilon),\, \tau(y; \theta_\tau)\,\rangle + b_0$$ where $\beta(\cdot)$ is the (random, SDE-driven) branch output, and $\tau(\cdot)$ is the trunk output.

3. Training via the Stochastic Maximum Principle and Adjoint Backward SDEs

SON reframes operator learning as a stochastic optimal control problem, incorporating the Stochastic Maximum Principle (SMP). The expected terminal loss is:

$$J(\theta) = \mathbb{E}\left[ \Phi(A_T, \Gamma) + \int_0^T r(A_t, \theta_t)\,dt \right]$$

Contrary to standard backpropagation, SON uses forward-backward SDEs for optimization:

• Forward SDE (data/branch network):

$$dA_t^* = \mu(A_t^*, \theta_t)dt + \sigma(\theta_t)dW_t$$

• Backward SDE (adjoints):

$$dB_t^* = -\left[\,r_a(A_t^*, \theta_t^*) + \mu_a^\top B_t^* + \sigma_a^\top C_t^*\,\right]dt - C_t^*dW_t$$

where $B_t^*$, $C_t^*$ serve as adjoint sensitivity processes.

The key innovation is to replace standard gradients with Hamiltonian gradients: $$\nabla_{\theta} J(\theta) = \mathbb{E}\left[ \mu_\theta^\top B_t + \sigma_\theta^\top C_t + r_\theta \right] = \mathbb{E} \left[ \nabla_{\theta} H(A_t, B_t, C_t, \theta) \right]$$ where the Hamiltonian $H$ incorporates both mean and diffusion contributions.

Numerically, the forward SDE and backward BSDE are discretized (Euler–Maruyama), with sample-wise efficient adjoint solutions.

4. Architecture and Comparison to Deterministic Operator Networks

Model	Output Type	Uncertainty Type	Training Gradient	Branch Net	Captures Output Variance
DeepONet	Deterministic mean	None / Only epistemic	Loss backprop	Feed-forward	No
SON	Random variable	Intrinsic (diffusion)	Hamiltonian, SMP (BSDE)	SDE-driven	Yes (pointwise)

The deterministic DeepONet architecture cannot model output variance due to operator noise; its predictions are always mean values with nearly zero variance under repeated trials on noisy data. SON, by contrast, empirically matches both mean and variance (uncertainty), recovers noise scaling, and delivers output distributions that match the ground-truth sample covariance profiles of the data.

5. Uncertainty Quantification and Empirical Evaluation

SON quantifies output uncertainty directly through its learned diffusion coefficients, supporting robust, sample-based UQ. In experiments involving noisy antiderivative tasks, stochastic ODEs (1D/2D), double integrals in 2D/3D, and stochastic elliptic PDEs, SON matches both mean and variance of noisy target operators, accurately modeling noise even where amplitude varies spatially.

Comparative analysis with DeepONet (for identical data, architecture, and optimizer) demonstrates:

• SON matches true variance and mean,
• Vanilla DeepONet captures only mean, with almost zero output variance,
• SON requires no significant increase in training time due to efficient backward SDE solvers.

6. Practical Relevance, Computational Aspects, and Outlook

The SON framework provides a principled, scalable approach to probabilistic operator learning that is computationally efficient—forward/backward SDE simulation adds minimal cost compared to deterministic feedforward NNs, due to sample-efficient adjoint methods.

Its technical innovations—SDE-based latent representation, SMP/BSD-driven gradients, and direct modeling of pointwise stochasticity—yield a more expressive operator surrogate model, with guaranteed quantification of uncertainty at every output location.

SON is applicable to a broad class of operator learning problems where modeling, propagating, and calibrating intrinsic stochasticity is essential, including:

• Surrogate modeling for stochastic PDE/ODE solvers,
• Physical systems with process/environment noise,
• Real-world scenarios with noisy measurements intrinsically linked to the true operator.

The SON framework thus marks a substantial step forward in aligning operator learning methodology with the realities of noisy, stochastic scientific and engineering systems, enabling reliable and uncertainty-aware predictions directly from data.

---

Summary Table: Deterministic vs Probabilistic Operator Learning (SON)

Capability	DeepONet	SON
Stochastic/noisy operators	❌	✅ Explicitly modeled
Pointwise uncertainty quantification	❌	✅ Learned diffusion
Empirical output variance matching	❌	✅ Matches true variance
Output: mean & variance	Mean only	Mean & variance (random band)
Training cost	Low	Comparable (with efficient BSD)
Optimization principle	Loss gradient	Hamiltonian (SMP/BSDE)