Factorized kernel attention for scalable PDE learning

The Transformer architecture has demonstrated state-of-the-art performance across various applications and has emerged as a promising tool for data-driven surrogate modeling of partial differential equations (PDEs). However, the application of attention mechanisms to a large number of grid points can lead to instability and remains computationally expensive, despite attempts to introduce linear-complexity variants. In this study, we propose a novel approach called Factorized Transformer (FactFormer), which is based on an axial factorized kernel integral. Specifically, we introduce a learnable projection operator that decomposes the input function into multiple sub-functions, each defined over a one-dimensional domain. These sub-functions are then utilized to compute the instance-based kernel using an axial factorized scheme. We validate the effectiveness of the proposed model on several challenging fluid dynamic systems, including 2D Kolmogorov flow, 3D isotropic turbulence and 3D smoke buoyancy.