# Zero-shot generalization of transformer neural operators to larger domains

arXiv: https://arxiv.org/abs/2606.14597 · HTML: https://mlsanity.com/paper/2606.14597/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 31 references that are in the map.
- First posted on arXiv 12 Jun 2026.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [Extending Context Window of Large Language Models via Positional Interpolation](https://mlsanity.com/paper/2306.15595/) (2023)
- [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://mlsanity.com/paper/2003.03485/) (2020)
- [YaRN: Efficient Context Window Extension of Large Language Models](https://mlsanity.com/paper/2309.00071/) (2023)
- [Symbolic Discovery of Optimization Algorithms](https://mlsanity.com/paper/2302.06675/) (2023)
- [Fourier Neural Operator with Learned Deformations for PDEs on General Geometries](https://mlsanity.com/paper/2207.05209/) (2022)
- [Choose a Transformer: Fourier or Galerkin](https://mlsanity.com/paper/2105.14995/) (2021)
- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://mlsanity.com/paper/2402.02366/) (2024)
- [Rotary Position Embedding for Vision Transformer](https://mlsanity.com/paper/2403.13298/) (2024)
- [Neural Operators with Localized Integral and Differential Kernels](https://mlsanity.com/paper/2402.16845/) (2024)
- [Geometry Aware Operator Transformer as an Efficient and Accurate Neural Surrogate for PDEs on Arbitrary Domains](https://mlsanity.com/paper/2505.18781/) (2025)
- [AB-UPT: Scaling Neural CFD Surrogates for High-Fidelity Automotive Aerodynamics Simulations via Anchored-Branched Universal Physics Transformers](https://mlsanity.com/paper/2502.09692/) (2025)
- [Nonlocality and Nonlinearity Implies Universality in Operator Learning](https://mlsanity.com/paper/2304.13221/) (2023)
- [Continuum Attention for Neural Operators](https://mlsanity.com/paper/2406.06486/) (2024)
- [AROMA: Preserving Spatial Structure for Latent PDE Modeling with Local Neural Fields](https://mlsanity.com/paper/2406.02176/) (2024)
- [Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators](https://mlsanity.com/paper/2402.12365/) (2024)
- [FlashBias: Fast Computation of Attention with Bias](https://mlsanity.com/paper/2505.12044/) (2025)
- [Operator Learning with Domain Decomposition for Geometry Generalization in PDE Solving](https://mlsanity.com/paper/2504.00510/) (2025)
- [Anchored-Branched Steady-state WInd Flow Transformer (AB-SWIFT): a metamodel for 3D atmospheric flow in urban environments](https://mlsanity.com/paper/2603.25635/) (2026)

## Built on by (most cited first, 1 of 1)

- [Learnable composition for neural operators](https://mlsanity.com/paper/2609.03069/) (2026)

## Abstract (arXiv)

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.