# Integrating Locality-Aware Attention with Transformers for General Geometry PDEs

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 32 references that are in the map.
- First posted on arXiv 18 Apr 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Fourier Neural Operator with Learned Deformations for PDEs on General Geometries](https://mlsanity.com/paper/2207.05209/) (2022)
- [Multipole Graph Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2006.09535/) (2020)
- [Choose a Transformer: Fourier or Galerkin](https://mlsanity.com/paper/2105.14995/) (2021)
- [GNOT: A General Neural Operator Transformer for Operator Learning](https://mlsanity.com/paper/2302.14376/) (2023)
- [Geometry-Informed Neural Operator for Large-Scale 3D PDEs](https://mlsanity.com/paper/2309.00583/) (2023)
- [U-FNO -- An enhanced Fourier neural operator-based deep-learning model for multiphase flow](https://mlsanity.com/paper/2109.03697/) (2021)
- [Transformer for Partial Differential Equations' Operator Learning](https://mlsanity.com/paper/2205.13671/) (2022)
- [Factorized Fourier Neural Operators](https://mlsanity.com/paper/2111.13802/) (2021)
- [Multiwavelet-based Operator Learning for Differential Equations](https://mlsanity.com/paper/2109.13459/) (2021)
- [U-NO: U-shaped Neural Operators](https://mlsanity.com/paper/2204.11127/) (2022)
- [Convolutional Neural Operators for robust and accurate learning of PDEs](https://mlsanity.com/paper/2302.01178/) (2023)
- [Patches Are All You Need?](https://mlsanity.com/paper/2201.09792/) (2022)
- [Scalable Transformer for PDE Surrogate Modeling](https://mlsanity.com/paper/2305.17560/) (2023)
- [Solving High-Dimensional PDEs with Latent Spectral Models](https://mlsanity.com/paper/2301.12664/) (2023)
- [Improved Operator Learning by Orthogonal Attention](https://mlsanity.com/paper/2310.12487/) (2023)
- [Representation Equivalent Neural Operators: a Framework for Alias-free Operator Learning](https://mlsanity.com/paper/2305.19913/) (2023)
- [Mitigating spectral bias for the multiscale operator learning](https://mlsanity.com/paper/2210.10890/) (2022)

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

- [RETO: A Rotary-Enhanced Transformer Operator for High-Fidelity Prediction of Automotive Aerodynamics](https://mlsanity.com/paper/2605.00062/) (2026)
- [Modularity-Free Conflict-Averse Training for Generalized PINNs](https://mlsanity.com/paper/2606.20156/) (2026)

## Abstract (arXiv)

Neural operators have emerged as promising frameworks for learning mappings governed by partial differential equations (PDEs), serving as data-driven alternatives to traditional numerical methods. While methods such as the Fourier neural operator (FNO) have demonstrated notable performance, their reliance on uniform grids restricts their applicability to complex geometries and irregular meshes. Recently, Transformer-based neural operators with linear attention mechanisms have shown potential in overcoming these limitations for large-scale PDE simulations. However, these approaches predominantly emphasize global feature aggregation, often overlooking fine-scale dynamics and localized PDE behaviors essential for accurate solutions. To address these challenges, we propose the Locality-Aware Attention Transformer (LA2Former), which leverages K-nearest neighbors for dynamic patchifying and integrates global-local attention for enhanced PDE modeling. By combining linear attention for efficient global context encoding with pairwise attention for capturing intricate local interactions, LA2Former achieves an optimal balance between computational efficiency and predictive accuracy. Extensive evaluations across six benchmark datasets demonstrate that LA2Former improves predictive accuracy by over 50% relative to existing linear attention methods, while also outperforming full pairwise attention under optimal conditions. This work underscores the critical importance of localized feature learning in advancing Transformer-based neural operators for solving PDEs on complex and irregular domains.