Technical Summary: Neural Operator Processes for Probabilistic Operator Learning under Partial Observations

1. Problem Statement

Neural Operators (NOs), such as the Fourier Neural Operator (FNO), have established themselves as effective paradigms for learning mappings between function spaces, particularly for solving partial differential equations (PDEs). However, standard NOs are typically developed and evaluated under the assumption of dense, fully paired input–output fields. This contrasts with many realistic scientific scenarios, such as data assimilation, where the target physical response is known only through sparse, irregular, or partial observations, even if the underlying domain properties are known densely.

Furthermore, uncertainty quantification (UQ) in operator learning remains underdeveloped, especially under sparse observations. Existing approaches often rely on model-external ensembles, post-hoc calibration, or remain tied to dense input-output formulations. While Neural Processes (NPs) offer a complementary perspective by learning distributions over functions from sparse context sets via amortized latent-variable inference, they are primarily designed for pointwise function regression and have not been widely adapted to structured function-to-function mappings that require dense PDE-scale decoding.

The core challenge addressed in this work is to unify the set-conditioned inference of NPs with the structured operator decoding of NOs to predict full output fields from limited, sparse context while providing coherent uncertainty quantification.

2. Methodology: Neural Operator Processes (NOPs)

The authors introduce Neural Operator Processes (NOPs), a framework that conditions on a sparse context set of joint input–output observations to infer the mapping from a given input field to its dense target solution. The framework supports both deterministic and probabilistic variants within a shared encoder–decoder architecture.

2.1 Architecture Overview

The NOP pipeline consists of three main stages:

1. Context Encoding: The sparse context set $C = \{(x_i^{(c)}, a_i^{(c)}, u_i^{(c)})\}$ is encoded into a context-conditioned representation $R_C(x)$ evaluated at query locations.
2. Initial Hidden Field Construction: The context representation is combined with query-side input information (the input field $a(x)$ on a dense grid) to form an initial hidden field $v_0(x)$.
3. Operator Decoding: A Neural Operator decoder (specifically an FNO-style decoder) processes $v_0(x)$ to produce the final prediction $\hat{u}(x)$.

2.2 Conditioning Strategies

The paper investigates two complementary conditioning mechanisms to bridge the set-based view of NPs with the continuous-domain view of NOs:

• SetConv Conditioning: A global pooling pathway that lifts the irregular sparse context set into a continuous field using kernelized set convolutions (SetConv). This separately encodes the density of context support and the observed local field values.
• Query-Aligned Attention Conditioning: An attention pathway that models which context points are most relevant to each query location, preserving local geometric structure. This allows the model to attend directly to the context set rather than relying solely on kernel proximity.
• Fused Representation: These pathways can be combined into a fused representation before decoding.

2.3 Probabilistic Extension

To enable uncertainty quantification, the deterministic architecture is extended with a global latent variable $z$ that captures task-level uncertainty.

• Latent Inference: A global summary of the context-induced representation is used to infer a context-conditioned prior $p_\theta(z | C)$. During training, an approximate posterior $q_\phi(z | C, U_Q)$ is learned using target outputs.
• Latent Conditioning: The latent variable $z$ is injected at the input stage (broadcast concatenation with $v_0(x)$) rather than modulating the decoder layers, which was found to yield more stable prior-conditioned inference.
• Predictive Distribution: The final predictive distribution is obtained by marginalizing the conditionally factorized likelihood over the latent variable, decomposing uncertainty into a global task-level component (from $z$) and a local predictive component (from the output likelihood).

3. Key Contributions

1. Formulation: The authors formulate operator learning under partial observations as context-conditioned prediction in function space and propose NOPs as the corresponding framework.
2. Architecture: They introduce deterministic and probabilistic NOP architectures sharing an encoder–decoder pipeline, featuring distinct mechanisms for context aggregation (SetConv vs. Attention) and latent uncertainty.
3. Empirical Distinction: The work identifies an empirical distinction between global and local conditioning, demonstrating that geometry-preserving attention is particularly valuable in geometry-sensitive settings, while the preferred pathway depends on the inference objective and PDE regime.
4. Validation: Through systematic experiments and ablations, the authors show that NOPs match or surpass dense-grid neural operators while observing only a small fraction of the response field (under 25% in 1D and under 7% in 2D), provide meaningful uncertainty estimates, and remain computationally competitive at inference.

4. Experimental Results

The framework was evaluated on four benchmark families: Gaussian Process (GP) regression (as a diagnostic), 1D Burgers equation (smooth, periodic), 2D Darcy flow (non-periodic, geometry-sensitive), and 2D Navier–Stokes (periodic but dynamically complex).

4.1 Deterministic Performance

• Sparse Viability: Sparse conditional operator prediction is viable, with NOP variants achieving errors comparable to or slightly lower than standard dense-grid FNO references.
• Regime Dependence:
  • Burgers (Periodic): The attention pathway improves performance even in standard settings, aligning well with the spectral decoder.
  • Darcy (Non-periodic): Preserving local context–query geometry via attention is essential. The SetConv-only variant degrades under uniform sampling but improves with clustered observations, whereas the attention-based variant (DANOP) flattens performance variations across sampling strategies, showing greater robustness to observation geometry shifts.
  • Navier–Stokes: Attention provides greater robustness under shifts in context sampling strategy and budget, though SetConv remains competitive under standard uniform evaluation.

4.2 Probabilistic Performance

• Uncertainty Quantification: The framework provides meaningful uncertainty estimates that widen in weakly constrained regions.
• Latent Conditioning: Probabilistic extensions preserve predictive capacity. The gap between prior-mean and posterior-mean performance is modest, indicating the learned prior is well-aligned with the posterior.
• Variance Parameterization:
  • Burgers: Heteroscedastic (HE) heads perform best, with negligible posterior–prior gaps.
  • Darcy: Homoscedastic (HO) heads often perform better, acting as a regularizer that prevents the model from absorbing difficulty through variance inflation.
  • Navier–Stokes: HE heads are superior; HO heads are too restrictive for this heterogeneous dynamical regime, leading to degraded performance.
• Robustness: The architecture with fused conditioning (PANOP) generally outperforms the SetConv-only variant (PNOP), particularly in zero-latent evaluations, suggesting that geometry-aware conditioning better supports latent stochastic inference.

4.3 Ablations and Limitations

• Cross-Resolution Transfer: While 1D Burgers exhibits successful zero-shot transfer to finer grids, 2D non-periodic settings (Darcy, Navier–Stokes) show significant degradation. This suggests that the learned conditional representation does not transfer cleanly across discretizations in complex regimes, a limitation potentially shared by the interaction between sparse representations and spectral decoders.
• Training vs. Inference Cost: Probabilistic extensions increase training time and memory (due to latent inference graphs) but maintain modest inference latency (milliseconds per batch).
• Failure Modes: A central failure mode is the loss of spatial structure when local context is overly compressed into global summaries. This is benign in smooth periodic settings but damaging in geometry-sensitive regimes like Darcy.

5. Significance and Claims

The paper claims that NOPs provide a practical, explicit, and extensible framework for combining sparse conditional inference, structured operator decoding, and uncertainty quantification. The significance of the work lies in:

• Unification: Bridging operator learning and probabilistic meta-learning in function space.
• Design Principles: Clarifying that explicitly preserving local geometry is a key design requirement for geometry-sensitive settings, and that reliable probabilistic inference requires latent conditioning that complements rather than overwrites the spatial inductive biases of the deterministic trunk.
• Practicality: Demonstrating that full-field prediction with coherent uncertainty is feasible using only a small fraction of observations, offering potential savings in sensing or storage for applications where dense measurements are costly.

The authors remain modest regarding the scope, noting that the benchmarks do not cover strongly chaotic dynamics, highly multiscale systems, or irregular geometries, and that the fixed FNO backbone may limit performance in non-periodic regimes. They identify future work in spatially structured latent representations and decoder families better suited to irregular geometries.