Physics Informed Token Transformer for Solving Partial Differential Equations

TL;DR

This paper tackles PDE solving by fusing physics directly into a transformer-based neural operator. It introduces PITT, which tokenizes governing equations and uses a latent equation representation to generate an analytically driven update $x_{t+1}=x_t+F_P(x_t)$, thereby enforcing physics-informed updates alongside a neural operator. The authors demonstrate broad improvements over Fourier Neural Operator, DeepONet, and OFormer across 1D and 2D benchmarks, including Heat, Burgers', KdV, Navier–Stokes, and Poisson, with two equation-embedding schemes and measurable gains in next-step, fixed-future, and rollout tasks, as well as interpretable attention maps that reflect parameter sensitivities. The results suggest that equation-embedded transformers can deliver accurate, efficient PDE surrogates with enhanced physical fidelity and generalization, and open avenues for extending to 3D systems and more robust rollout training.

Abstract

Solving Partial Differential Equations (PDEs) is the core of many fields of science and engineering. While classical approaches are often prohibitively slow, machine learning models often fail to incorporate complete system information. Over the past few years, transformers have had a significant impact on the field of Artificial Intelligence and have seen increased usage in PDE applications. However, despite their success, transformers currently lack integration with physics and reasoning. This study aims to address this issue by introducing PITT: Physics Informed Token Transformer. The purpose of PITT is to incorporate the knowledge of physics by embedding partial differential equations (PDEs) into the learning process. PITT uses an equation tokenization method to learn an analytically-driven numerical update operator. By tokenizing PDEs and embedding partial derivatives, the transformer models become aware of the underlying knowledge behind physical processes. To demonstrate this, PITT is tested on challenging 1D and 2D PDE neural operator prediction tasks. The results show that PITT outperforms popular neural operator models and has the ability to extract physically relevant information from governing equations.

Figures (20)

• Figure 1: The Physics Informed Token Transformer (PITT) uses standard multi head self-attention to learn a latent embedding of the governing equations. This latent embedding is then used to perform numerical updates using linear attention blocks. The equation embedding acts as an analytically-driven correction to an underlying data-driven neural operator.
• Figure 2: Example setup for the 2D Poisson equation. a) Input boundary conditions and geometry. b) Target electric field output.
• Figure 3: PITT FNO prediction decomposition for 1D Heat equation. Left: The FNO module of PITT predicts a large change to the final frame of input data. Middle The numerical update block corrects the FNO output. Right The combination of FNO and numerical update block output very accurately predicts the next step.
• Figure 4: Error accumulation for rollout experiments using standard embedding.
• Figure 5: Rollout results for 2D Navier Stokes using our novel embedding method.