Scalable Transformer for PDE Surrogate Modeling

TL;DR

This paper introduces FactFormer, a scalable Transformer model with axial factorized kernels for efficient and accurate PDE surrogate modeling on high-dimensional grids, outperforming traditional methods in stability and computational cost.

Contribution

The paper proposes a novel Factorized Transformer with axial kernel decomposition, enabling efficient PDE surrogate modeling on large multi-dimensional grids.

Findings

Successfully simulates 2D Kolmogorov flow on 256x256 grid

Accurately models 3D smoke buoyancy on 64x64x64 grid

Offers a computationally efficient low-rank surrogate for full attention

Abstract

Transformer has shown state-of-the-art performance on various applications and has recently emerged as a promising tool for surrogate modeling of partial differential equations (PDEs). Despite the introduction of linear-complexity attention, applying Transformer to problems with a large number of grid points can be numerically unstable and computationally expensive. In this work, we propose Factorized Transformer (FactFormer), which is based on an axial factorized kernel integral. Concretely, we introduce a learnable projection operator that decomposes the input function into multiple sub-functions with one-dimensional domain. These sub-functions are then evaluated and used to compute the instance-based kernel with an axial factorized scheme. We showcase that the proposed model is able to simulate 2D Kolmogorov flow on a $256\times 256$ grid and 3D smoke buoyancy on a $64\times64\times64$ grid with good accuracy and efficiency. The proposed factorized scheme can serve as a computationally efficient low-rank surrogate for the full attention scheme when dealing with multi-dimensional problems.