Zero-shot generalization of transformer neural operators to larger domains

Abstract: Transformer-based neural operators have shown remarkable performance for approximating solution operators of partial differential equations on complex geometries. However, existing approaches implicitly assume a fixed domain size, which limits their ability to generalize at inference. In this work, we investigate domain extension, namely zero-shot inference on spatial domains that are significantly larger than those encountered during training. We argue that this setting fundamentally requires spatial locality and translation equivariance. We propose to implement this locality via a decomposable bias in the attention logits computation, enabling finely controllable locality while remaining fully decomposable into query-key inner products and directly compatible with optimized attention kernels. Combined with rotary positional embeddings, it enables expressive embeddings with controllable spatial support without altering the transformer architecture. We empirically show that our approach substantially improves zero-shot generalization to larger domains across two PDE benchmarks and a 3D industrial atmospheric flow application. Our code and datasets are available at https://github.com/cerea-daml/domain-extension.

• The paper introduces decomposable locality biases (LASPE and LAAPE) to enforce spatial locality in transformer neural operators for stable domain extrapolation.
• It demonstrates that LAAPE consistently outperforms RoPE and LASPE across 1D, 2D, and 3D PDE benchmarks by reducing normalized L1 error and preserving solution structures.
• The approach offers practical insights for industrial PDE simulations by mitigating error spikes and enhancing numerical stability during large-scale domain extensions.

Zero-shot Generalization of Transformer Neural Operators to Larger Domains

Problem Formulation and Locality Principle

The paper "Zero-shot generalization of transformer neural operators to larger domains" (2606.14597) addresses a critical challenge in neural operator learning for PDEs: inference on spatial domains that exceed those observed during training, a scenario termed "domain extension". Unlike super-resolution or partial observation tasks, domain extension is fundamentally an extrapolation problem. The architecture must avoid spurious long-range dependencies due to periodic embedding artifacts when scaling the physical domain.

The authors argue that robust domain extension requires explicit enforcement of spatial locality and translation equivariance in transformer-based neural operators. While transformer attention is inherently non-local, learning on bounded training domains restricts the effective kernel support, leaving models vulnerable to pathological attention patterns—especially when rotary positional embeddings (RoPE) produce periodic "ghost" interactions for out-of-distribution spatial distances. This insight motivates the design of architectural inductive bias so that all interactions during inference remain local to the training region, independent of domain size.

Figure 1: Domain upscaling: RoPE alone creates spurious long-range interactions when scaling the domain, whereas the proposed decomposable bias enforces locality.

Decomposable Local Bias Construction

To solve the locality problem, the paper proposes augmenting RoPE-based attention logits with a decomposable distance-based bias, ensuring compatibility with efficient matrix multiplication-based attention kernels (e.g., FlashAttention). This bias is formulated in two versions:

• LASPE (Local Anisotropic Symmetric Positional Encoding): Extends FlashBias with an anisotropic L2 decay term, restricting kernel support explicitly in the attention computation.
• LAAPE (Local Anisotropic Asymmetric Positional Encoding): Employs direction-dependent exponential decay to enable asymmetric spatial locality, critical for domains with geometric/behavioral asymmetry (e.g., upwind/downwind flow around obstacles).

Both approaches provide tunable spatial support, either isotropic or directionally restricted, without breaking decomposability into query-key products. The asymmetric LAAPE is a natural extension of LASPE, reducing to it in the symmetric case and providing superior stability at large coordinate ranges.

Figure 2: Comparison of kernel shapes enforced by LASPE and LAAPE biases, demonstrating faster decay and enhanced asymmetry with LAAPE.

Empirical Evaluation on Multiple PDE Domains

The authors evaluate locality-enforced transformer neural operators across three benchmarks:

• 1D Shallow Water Equations: Models the propagation of water height and velocity.
• 2D Gray-Scott Reaction-Diffusion: Captures complex interactions between chemical species.
• 3D Microscale Atmospheric Flow: Industrial-scale simulations of wind flow in urban environments, exhibiting strong anisotropy and upwind/downwind asymmetry.

In each case, the model is trained exclusively on simulations over small domains and evaluated zero-shot on domains with up to 10x scaling. The performance is measured through normalized L1 error on relevant field quantities.

Key numerical results are as follows:

• RoPE degrades rapidly as domain size increases, with substantial error spikes and qualitatively incorrect predictions, especially for scaling factors beyond the training range.
• LASPE maintains accuracy for moderate scaling, but exhibits numerical instability at large coordinate values, as seen in error maps and convergence degradation.
• LAAPE remains stable and accurate across all tested scaling factors, consistently outperforming both RoPE and LASPE on large domains in terms of both mean error and qualitative solution fidelity.

  Figure 3: Autoregressive prediction of $V$ in 3-step Gray-Scott simulation with $s_x \times s_y = 3 \times 3$; LAAPE maintains sharp patterns, RoPE prediction collapses.

  

  Figure 4: Comparison at $s_x \times s_y = 7 \times 7$ under BFloat16 precision: LASPE suffers strong spatially-varying numerical errors; LAAPE prediction remains uniform and consistent.

Robustness in Industrial Atmospheric Flow

In the industrial 3D atmospheric flow benchmark, the architecture leverages LAAPE to support asymmetric interactions—critical for encoding the upstream building area and downstream wake, each with their own kernel support. RoPE fails drastically on large domains, producing physically implausible results, while LAAPE preserves coherent flow structures and accurately predicts long-range wake patterns, albeit with error growth due to novel flow configurations absent from training.

Figure 5: Left: Upwind buildings region and corresponding kernel support in encoder/processor blocks. Right: Flow predictions for $s_x \times s_y=5\times 5$, LAAPE preserves wake structure while RoPE prediction collapses.

Visualizations and Quantitative Scaling

The appendix further demonstrates model predictions in ground truth overlays for 1D/2D/3D tasks, underscoring the stability and accuracy as domain scales increase.

Figure 6: Shallow water ground truth and predictions at $s=1$, showing accurate reproduction of velocity and height.

Figure 7: Shallow water ground truth and predictions at $s=3$, demonstrating robust zero-shot generalization.

Figure 8: Shallow water ground truth and predictions at $s=10$, LAAPE maintains solution structure.

Similar visualizations for Gray-Scott and atmospheric flow show consistent qualitative performance for LAAPE, with error maps and flow fields closely matching ground truth across domain scaling.

Numerical and Computational Limitations

The paper performs a careful numerical stability analysis. LASPE is shown to be sensitive to floating-point precision issues as domain coordinates grow, leading to substantial overflow/underflow and spatially heterogeneous predictions. LAAPE, with exponential terms more naturally compatible with hardware limits, supports domain scaling at up to $181\,\lambda$ (where $\lambda$ is the decay width) under BFloat16, covering the vast majority of physical simulation scenarios.

Figure 9: Numerical stability experiment; LASPE error increases with large coordinates in Bfloat16, LAAPE remains stable.

Implications, Scalability, and Future Directions

By enforcing locality explicitly through decomposable bias functions, transformer neural operators achieve robust zero-shot generalization to arbitrarily large domains, without retraining or modification. This advances the practical utility of neural operators in industrial-scale PDE settings, where training data is expensive and domain scaling is unavoidable.

A key implication is that controlling kernel support provides a principled scaling law for deep operator learning: generalization depends not only on data/model size, but on the spatial structure encoded in the attention mechanism. Sparse attention matrices induced by locality can potentially reduce computational complexity from quadratic to linear, though further algorithmic work (e.g., mesh reordering, space-filling curves) is required for optimal exploitation in irregular point clouds.

Further, enforcing physically meaningful locality is essential for situations where extrapolation beyond training correlations is risky, and where global solution behavior does not scale proportionally with domain size. Fine-tuning on larger domains may be necessary to capture emergent long-range correlations beyond those seen during training. The authors suggest integration with mesh partitioning, hierarchical transformers [Zhdanov2025], and more general relative positional encodings for broader applicability.

Conclusion

The study provides a rigorous, empirically validated methodology for robust zero-shot spatial domain extension in transformer neural operators. Through analytically decomposable locality-constrained positional biases, the proposed architectures retain high expressiveness, effective translation equivariance, and numerical stability on large-scale PDE domains. These findings have direct practical impact for deployment in computational engineering and physics, and motivate further research in locality-aware neural operator scaling and efficient attention for unstructured meshes.