Technical Summary: Zero-shot Generalization of Transformer Neural Operators to Larger Domains

Problem Statement

Transformer-based neural operators have demonstrated strong performance in approximating solution operators for partial differential equations (PDEs) on complex geometries. However, existing approaches typically assume a fixed domain size during training, limiting their ability to generalize at inference. This paper addresses domain extension: the challenge of performing zero-shot inference on spatial domains significantly larger than those encountered during training.

Unlike super-resolution or partial observation tasks, which are interpolation problems, domain extension is fundamentally an extrapolation problem. Current transformer-based operators rely on non-local integration kernels implicitly bounded by the finite training domain. When the inference domain size changes, the implicit periodic nature of standard positional embeddings (e.g., Rotary Positional Embeddings, RoPE) can induce "ghost" interactions between spatially distant points that were never observed together during training. This breaks the model's ability to generalize, as the operator fails to respect the physical locality of interactions relative to the new domain scale.

Methodology

The authors propose a method to enforce explicit, controllable spatial locality within the attention mechanism of transformer neural operators without altering the underlying architecture or sacrificing compatibility with optimized attention kernels (e.g., FlashAttention).

Core Mechanism: Decomposable Locality Bias

The approach combines Rotary Positional Embeddings (RoPE) with a decomposable distance bias added to the attention logits.

1. RoPE Backbone: RoPE provides translation equivariance and expressive relative positional encoding. However, it is inherently non-local and quasi-periodic.
2. Locality Constraint: To prevent spurious long-range interactions, a bias term $\phi$ is subtracted from the attention logits. This bias acts as a geometrical potential that forces the attention kernel to decay rapidly beyond a specific distance $r$, effectively restricting the integration support to a neighborhood $B_r(c)$.
3. Decomposability: To maintain compatibility with efficient GPU kernels, the bias must be decomposable into a function of the query position and a function of the key position ($z = q^\top R_\theta k - \phi$).

Proposed Embeddings

The paper introduces two specific formulations for this bias:

• Local Anisotropic Symmetric Positional Encoding (LASPE):
  Uses an anisotropic $L_2$ distance bias:
  $$z_{LASPE} = q^\top(c) R_\theta^\top(c-\xi) k(\xi) - \frac{\sqrt{d}}{2} (c-\xi)^\top A (c-\xi)$$
  Here, $A$ is a diagonal matrix controlling anisotropy. While effective for symmetric domains, this method relies on quadratic terms which can lead to numerical instabilities at large coordinate values.

• Local Anisotropic Asymmetric Positional Encoding (LAAPE):
  To handle asymmetry (e.g., wind flow where interactions differ upstream vs. downstream) and improve numerical stability, the authors propose a bias based on hyperbolic functions (specifically exponentials):
  $$z_{LAAPE} = q^\top(c) R_\theta^\top(c-\xi) k(\xi) - \frac{\sqrt{d}}{2} \sum_i \left[ \exp(A_-^{1/2}(c-\xi)\cdot e_i) + \exp(A_+^{1/2}(\xi-c)\cdot e_i) \right]$$
  This formulation allows for different decay rates in positive and negative directions ($A_+$ and $A_-$), enabling the kernel support to be asymmetric and anisotropic. It is shown to be a natural extension of LASPE for small displacements but offers superior stability for large domains.

Kernel Width Selection

The authors argue that the kernel width (controlled by the decay parameter $\lambda$) should be chosen based on the training domain size. Specifically, $\lambda$ should be set such that the attention weights become statistically insignificant for distances exceeding the training domain, preventing the model from learning spurious correlations from the extended test domain.

Key Contributions

1. Problem Formulation: The paper formally defines domain extension as a zero-shot generalization challenge for neural operators, highlighting the limitations of current non-local assumptions in extrapolation settings.
2. Novel Positional Encodings: The authors propose LAAPE and LASPE, which combine RoPE with decomposable distance biases. These encodings provide:
   • Controllable Locality: Explicitly bounding the interaction range.
   • Anisotropy and Asymmetry: The ability to model direction-dependent interaction ranges, crucial for complex physical systems like fluid dynamics.
   • Efficiency: Full compatibility with optimized attention kernels (FlashAttention) via decomposability.
3. Empirical Validation: The approach is validated on two academic PDE benchmarks (Shallow Water equations in 1D and Gray-Scott equations in 2D) and a challenging 3D industrial application (microscale atmospheric flow in urban environments).

Experimental Results

The models were trained on small domains and tested on progressively larger domains (scaling factors $s$ up to 10x).

• Academic Benchmarks (1D & 2D):

  • Standard RoPE models showed significant error degradation as the domain size increased (e.g., L1 error for Gray-Scott increased from ~30% to >400% at $s=10$).
  • LASPE and LAAPE models maintained nearly constant error rates across scaling factors.
  • LAAPE outperformed LASPE in the Gray-Scott benchmark at high scaling factors ($s \ge 7$), where LASPE exhibited numerical instabilities due to large coordinate values.

• Industrial Application (3D Atmospheric Flow):

  • In the AB-SWIFT model (microscale atmospheric flow), the task involved strong anisotropy and asymmetry (buildings located upwind).
  • RoPE failed catastrophically on scaled domains ($5\times5$), producing incoherent flow predictions.
  • LAAPE preserved the overall flow structure and accurately predicted wake regions, with error increases being significantly lower than RoPE. The authors attribute the remaining error in LAAPE to new flow patterns arising from the larger domain that were not present in the training data, rather than a failure of the operator itself.

Significance and Claims

The paper claims that enforcing structured spatial locality is a critical inductive bias for scaling neural operators to industrial applications. By explicitly constraining the effective support of interactions to match the training domain, models can generalize to larger domains without requiring proportional increases in training data or model capacity.

The authors position this work as a step toward "appropriate scaling laws" in machine learning for physics, suggesting that performance gains can arise from modulating inductive biases to reflect how information propagates across space, rather than relying solely on data and model scale. The method enables robust zero-shot domain extension, a capability previously limited to domain decomposition strategies or graph-based models, but achieved here within a unified transformer architecture compatible with high-performance computing kernels.

The authors acknowledge limitations, noting that while locality improves generalization, it cannot capture physical correlations that inherently exist at scales larger than the training domain. They suggest that for such cases, fine-tuning on a few large-scale cases remains a necessary future direction.