Efficient High-Accuracy PDEs Solver with the Linear Attention Neural Operator

TL;DR

Neural operators learning PDE mappings face a scalability-accuracy trade-off between softmax attention ($O(N^2 d)$) and linear attention ($O(N d^2)$). LAN0 introduces a compact set of agent tokens with complexity $O(M N d)$ to mediate global interactions, preserving expressive power while reducing cost. A universal approximation theorem for LANO is established, and empirical results across solid and fluid mechanics benchmarks show substantial accuracy gains over slice-based baselines like Transolver, along with discretization-invariant performance. This work provides a scalable, high-accuracy operator backbone applicable to irregular domains and multi-physics PDEs, with potential for uncertainty quantification and inverse problems.

Abstract

Neural operators offer a powerful data-driven framework for learning mappings between function spaces, in which the transformer-based neural operator architecture faces a fundamental scalability-accuracy trade-off: softmax attention provides excellent fidelity but incurs quadratic complexity $\mathcal{O}(N^2 d)$ in the number of mesh points $N$ and hidden dimension $d$, while linear attention variants reduce cost to $\mathcal{O}(N d^2)$ but often suffer significant accuracy degradation. To address the aforementioned challenge, in this paper, we present a novel type of neural operators, Linear Attention Neural Operator (LANO), which achieves both scalability and high accuracy by reformulating attention through an agent-based mechanism. LANO resolves this dilemma by introducing a compact set of $M$ agent tokens $(M \ll N)$ that mediate global interactions among $N$ tokens. This agent attention mechanism yields an operator layer with linear complexity $\mathcal{O}(MN d)$ while preserving the expressive power of softmax attention. Theoretically, we demonstrate the universal approximation property, thereby demonstrating improved conditioning and stability properties. Empirically, LANO surpasses current state-of-the-art neural PDE solvers, including Transolver with slice-based softmax attention, achieving average $19.5\%$ accuracy improvement across standard benchmarks. By bridging the gap between linear complexity and softmax-level performance, LANO establishes a scalable, high-accuracy foundation for scientific machine learning applications.

Figures (6)

• Figure 1: The architecture of LANO incorporates an agent attention block, as illustrated on the right-hand side. In this block, the query, key, and value matrices $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{N \times d}$ are first obtained. A pooling operation is then applied to $\mathbf{Q}$ to derive the agent token $\mathbf{A} \in \mathbb{R}^{M \times d}$. In the aggregation flow, $\mathbf{A}$ serves as the query, $\mathbf{K}$ as the key, and $\mathbf{V}$ as the value. In the mediated flow, $\mathbf{Q}$ acts as the query, $\mathbf{A}$ as the key, and the output from the aggregation flow as the value. Both attention flows follow the standard scaled dot-product softmax attention mechanism.
• Figure 2: Qualitative comparison of model performance on the Elasticity benchmark. Ground truth (GT), Transolver and LANO predictions, along with the corresponding input geometry and absolute error fields.
• Figure 3: Qualitative comparison of model performance on the Plasticity benchmark. Ground truth (GT), Transolver and LANO predictions at $t=20$, along with the corresponding input die geometry and absolute error fields.
• Figure 4: Qualitative comparison of model performance on the Airfoil benchmark. Ground truth (GT), Transolver and LANO predictions, along with the corresponding input airfoil geometry and absolute error fields.
• Figure 5: Qualitative comparison of model performance on the Pipe benchmark. Ground truth (GT), Transolver and LANO predictions, along with the corresponding input pipe geometry and absolute error fields.

Theorems & Definitions (14)

• Theorem 1
• proof
• Definition 1
• Lemma 1
• proof
• Remark 1
• Remark 2
• Theorem 2: Universal Approximation of LANO
• proof
• Lemma 2