Technical Summary: Imposing Boundary Conditions on Neural Operators via Learned Function Extensions

1. Problem Statement

Neural operators have emerged as powerful data-driven surrogates for solving partial differential equations (PDEs), learning mappings from problem conditions to solutions. However, existing frameworks often struggle with general, highly variable boundary conditions (BCs). While many applications treat BCs as the primary drivers of system behavior, standard operator learning (OL) architectures are typically limited to problems with constant or weakly varying BCs across training and inference samples.

The challenge is particularly acute for elliptic PDEs (e.g., Poisson, linear elasticity, hyperelasticity). Unlike parabolic or hyperbolic PDEs where boundary influence propagates gradually or with finite speed, elliptic solutions depend globally on boundary data. Consequently, even minor variations in BCs can induce global variations in the solution. Existing methods to handle BCs are often restricted to specific geometries (e.g., Cartesian grids), fail to support complex mixed-type BCs, or exhibit poor performance when BCs vary significantly between samples.

2. Methodology

The authors propose a general framework that conditions neural operators on complex, non-homogeneous BCs through function extensions. The core idea is to map boundary data to latent "pseudo-extensions" defined over the entire spatial domain, allowing standard domain-to-domain neural operator architectures to consume boundary information naturally.

2.1. Merging and Normalization

To handle mixed BC types (Dirichlet, Neumann, Robin) and component-wise variations without complex branching, the authors propose a unified formulation:
$$B(q, u) = \alpha \odot B_D(u) + \beta \odot B_N(u) - \gamma$$
where $\alpha, \beta, \gamma$ are boundary coefficient functions defined over the entire boundary $\partial\Omega$.

• Normalization: Recognizing that raw BCs often have disparate physical units and magnitudes (e.g., displacement vs. force), the authors introduce a principled shift-and-scale transformation. This ensures well-balanced input functions suitable for stable training while preserving physical consistency.

2.2. The Extender Module

The framework introduces an extender module $\Psi_\theta$ that maps boundary functions $q$ to a domain function $\psi$ (the pseudo-extension). The overall neural operator is constructed as:
$$G_\theta(a, q) = \Phi_\theta \left( \begin{bmatrix} a \\ \Psi_\theta(q) \end{bmatrix} \right)$$
where $a$ represents domain coefficients and $\Phi_\theta$ is the core domain-to-domain operator. Three strategies for $\Psi$ are investigated:

1. Zero Extensions (0X): Padding boundary values with zeros outside the boundary. This is computationally free but less accurate for complex problems.
2. Harmonic Extensions (HX): Solving a Laplace equation to find the smoothest extension. While physically motivated, this requires solving an auxiliary PDE (e.g., via FEM) at inference, creating a computational bottleneck.
3. Learned Pseudo-Extensions (LX): A learnable neural module trained end-to-end with the core operator.
   • Architecture: The proposed LX uses an attention-based mechanism. It processes downsampled boundary features and domain features through a series of Cross-Attention (CA) layers followed by Feed-Forward (FF) blocks.
   • Integration: The extender is integrated into the standard encode-process-decode paradigm. Boundary features are injected into the domain features via CA layers, allowing the model to learn how to propagate boundary information throughout the domain efficiently.

3. Key Contributions

• General Framework: A flexible OL framework that extends existing architectures to explicitly incorporate complex, varying BCs without requiring specialized boundary treatments at the domain edges.
• Unified BC Representation: A procedure that unifies Dirichlet, Neumann, and Robin BCs into a physically consistent, normalized representation, handling component-wise variations and mixed-type segments.
• Benchmark Datasets: The construction of 18 challenging datasets covering Poisson, linear elasticity, and hyperelasticity problems. These datasets feature:
  • Highly variable, mixed-type, and component-wise BCs.
  • Up to eight boundary segments with random sizes, locations, and BC types.
  • Diverse geometries, including non-convex shapes with holes.
• Extender Strategies: The proposal and evaluation of Zero, Harmonic, and Learned extensions, demonstrating that learned pseudo-extensions offer the best trade-off between accuracy and computational efficiency.

4. Results

The framework was evaluated using two state-of-the-art neural operator backbones: RIGNO (graph-based) and GAOT (transformer-based).

• Accuracy: The proposed approach achieves state-of-the-art accuracy, significantly outperforming baselines (including BENO and standard zero-padding) across all 18 datasets.
  • For complex elasticity and hyperelasticity problems, Learned Extensions (LX) reduced average test errors from ~30% (with zero extensions) to ~11% when using RIGNO.
  • For the Poisson problem, harmonic extensions performed well, but learned extensions provided comparable accuracy with lower computational overhead.
• Efficiency: The learned extender allows for accurate BC imposition without the inference-time cost of solving auxiliary PDEs (required for harmonic extensions).
• Scaling:
  • Data Scaling: Model accuracy improves consistently with larger training datasets, with the generalization gap closing as data increases.
  • Model Scaling: Increasing the size of the extender module improves performance, particularly for complex problems where basic extensions fail.
• Robustness and Generalization:
  • Noise: The framework is robust to input noise, maintaining low error rates even with 10% Gaussian noise in boundary functions.
  • Transfer Learning: Pre-training extenders on one geometry or BC type significantly improves performance on unseen geometries and BC types (e.g., pre-training on Square geometry reduced error on Circle geometry by ~85% relative to random initialization).
  • Physics Transfer: Pre-training on linear elasticity reduced errors for hyperelasticity problems from 67% to 23% with minimal fine-tuning.

5. Significance and Claims

The paper claims that learning boundary-to-domain extensions is an effective and practical strategy for imposing complex BCs in existing neural operator frameworks.

• Practicality: The method requires no hyperparameter tuning across different datasets and integrates seamlessly with standard architectures (encode-process-decode).
• Generality: The approach is agnostic to the choice of neural operator backbone (graph-based or transformer-based) and applicable to arbitrary domains and mixed BC types.
• Scientific Impact: By enabling accurate and robust modeling of PDEs with highly variable boundary conditions, this work facilitates the application of scientific machine learning to a broader range of engineering and physics problems, particularly those where boundary conditions are the primary source of variability (e.g., design optimization, inverse problems, and uncertainty quantification).

The authors note limitations, including the current restriction to 2D problems (due to a lack of 3D benchmarks with similar BC variability) and the potential for further optimization of the attention mechanism's efficiency. However, the coordinate-based nature of the extender suggests direct applicability to 3D settings.