Positional Knowledge is All You Need: Position-induced Transformer (PiT) for Operator Learning

Junfeng Chen,Kailiang Wu

2024-05-15

Abstract:Operator learning for Partial Differential Equations (PDEs) is rapidly emerging as a promising approach for surrogate modeling of intricate systems. Transformers with the self-attention mechanism$\unicode{x2013}$a powerful tool originally designed for natural language processing$\unicode{x2013}$have recently been adapted for operator learning. However, they confront challenges, including high computational demands and limited interpretability. This raises a critical question: Is there a more efficient attention mechanism for Transformer-based operator learning? This paper proposes the Position-induced Transformer (PiT), built on an innovative position-attention mechanism, which demonstrates significant advantages over the classical self-attention in operator learning. Position-attention draws inspiration from numerical methods for PDEs. Different from self-attention, position-attention is induced by only the spatial interrelations of sampling positions for input functions of the operators, and does not rely on the input function values themselves, thereby greatly boosting efficiency. PiT exhibits superior performance over current state-of-the-art neural operators in a variety of complex operator learning tasks across diverse PDE benchmarks. Additionally, PiT possesses an enhanced discretization convergence feature, compared to the widely-used Fourier neural operator.

Machine Learning,Numerical Analysis

What problem does this paper attempt to address?

The paper attempts to address the issues of high computational demand and limited interpretability in existing Transformer models based on self-attention mechanisms for operator learning in partial differential equations (PDEs). Specifically, the authors propose a new architecture called the Position-induced Transformer (PiT), which is based on an innovative positional attention mechanism. Compared to classical self-attention, this mechanism shows significant advantages in operator learning tasks. The positional attention mechanism relies solely on the spatial relationships between the sampling points of the input function, rather than the function values themselves. This greatly improves computational efficiency and provides better interpretability in the numerical solution of PDEs.

Moreover, PiT outperforms the current state-of-the-art neural operators in various complex operator learning tasks and exhibits excellent discretization convergence, especially when dealing with high-dimensional and complex geometric problems in different benchmarks. By introducing positional attention and its variants (cross-positional attention and local positional attention), the authors not only address the computational bottleneck issues of existing methods but also enhance the model's generalization ability across different grid resolutions. These improvements enable PiT to achieve outstanding performance in various PDE benchmarks and demonstrate its advantages in discretization convergence.