PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Transformer for PDE Surrogate Modeling

Zijie Li, Dule Shu, Amir Barati Farimani

arXiv:2305.17560v2cs.LG

TL;DR

Large-grid PDE surrogate modeling remains challenging because attention over many grid points can be numerically unstable and computationally expensive. FactFormer replaces the multidimensional kernel integral with an axial factorized scheme based on learnable one-dimensional projections, and demonstrates accurate, efficient simulation on 2D Kolmogorov flow and 3D smoke buoyancy.

  • Problem

    Transformer attention for PDEs can become numerically unstable and computationally expensive when applied to problems with many grid points.

  • Method

    FactFormer uses a learnable projection operator to create one-dimensional sub-functions and computes an axial factorized kernel integral over them.

  • Results

    The model simulates 2D Kolmogorov flow on a 256 × 256 grid and 3D smoke buoyancy on a 64 × 64 × 64 grid with good accuracy and efficiency.

  • Takeaways & Limitations

    The factorized scheme provides a computationally efficient low-rank surrogate for full attention in multidimensional PDE problems.

Abstract

from arXiv · show

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.

1 Introduction

PDE solvers discretize continuous domains into grids, while Transformer-based PDE modeling faces scalability challenges as grid-point interactions grow with dimensionality. The paper proposes modifying attention to improve Transformer scalability for high-resolution PDE grids.

  • Motivation: Numerical PDE solvers discretize continuous domains into grids and reduce differential equations to algebraic equations.They use methods including finite difference, finite element, finite volume, and spectral methods.
  • Motivation: Attention-based Transformers have become state-of-the-art across applications and are increasingly applied to PDE modeling.Attention can be interpreted as a learnable kernel integral or Galerkin projection when inputs are functions sampled on grids.
  • Evaluation setting: Figure 1 visualizes predictions and reference ground truth for 2D Kolmogorov flow on a 256 × 256 grid and 3D smoke buoyancy on a 64 × 64 × 64 grid.The smoke-buoyancy visualization shows a zOy cross-section.
  • Challenge: As grid points grow exponentially with dimension, attention over all points creates large interactions that can make deep stacks unstable and costly on high-resolution grids.This motivates a more scalable attention mechanism for PDE modeling.

2 Related works

Related work spans mesh-specific neural PDE solvers, neural operators, solver-enhancing networks, Transformer-based physics simulation, and efficient attention mechanisms. FactFormer is most closely related to neural operators and linear-complexity attention through its learnable projection and kernel-integral components.

  • Neural PDE solver: Mesh-specific neural PDE solvers use architectures such as convolutions for uniform meshes or graph layers for irregular meshes.These models can train from input-target data without access to the governing PDEs.
  • Neural operators: Neural operators learn mappings using function bases such as Fourier, wavelet, learned, or Green’s-function-based representations.Their training can also combine with physics-informed neural networks for more physically consistent predictions.
  • Solver-enhanced methods: Neural networks can augment numerical solvers by correcting solver errors during simulation or reconstructing high-fidelity solutions offline.These approaches target improved solver accuracy and efficiency rather than direct surrogate modeling alone.
  • Transformer for physics simulation: Transformers are used for physics simulation both as neural PDE solvers and as data-driven models when governing PDEs are unavailable.Dot-product attention can be viewed as an approximation of an integral transform.
  • Efficient Transformer: Efficient Transformer research reduces attention cost through softmax-free linear attention or low-rank approximations of query-key products.FactFormer is related to the softmax-free linear-attention direction.

3 Method

FactFormer interprets attention as a learnable kernel integral and factorizes multidimensional interactions across axes. Learnable projections create one-dimensional sub-functions whose kernels update the value tensor sequentially, with latent marching extending predictions across time steps.

  • 3.1 Attention mechanism: Attention maps queries, keys, and values from inputs and computes weighted value averages using a weight function.In self-attention, queries, keys, and values are derived from the same source through learnable projections.
  • 3.1 Attention mechanism: Viewing grid-sampled features as functions, softmax-free attention becomes numerical quadrature of a learnable kernel integral.The kernel is formed from query and key functions, with quadrature weights applied over the discretization grid.
  • 3.2 Multidimensional factorized attention: Standard softmax-free attention has linear sequence-length complexity through matrix associativity but can remain unstable for many grid points and quadratic in channel dimension.Its kernel matrix is low-rank because it is the product of tall, thin query and key matrices.
  • 3.2 Multidimensional factorized attention: Factorized kernel attention computes an axis-specific kernel from each projected sub-function and iteratively applies tensor-matrix products to update the value before a pointwise feedforward network.The axial grid size makes each axis-wise kernel computation affordable for many problems, while each tensor-matrix product costs O(NS_md).
  • 3.2 Multidimensional factorized attention: The learnable projection operator decomposes an n-dimensional input function into multiple one-dimensional sub-functions, using pointwise transformations and pooling over other axes on uniform grids.Each projected function is evaluated along its corresponding spatial axis.
  • 3.3 Training techniques: Latent marching propagates the latent state with a pointwise learnable function, allowing one neural-solver call to produce multiple future steps.The method is chosen because forward Euler prediction in physical space is less stable for fluid problems with relatively large time steps.

4 Experiment

Experiments evaluate FactFormer on fluid-like and steady-state PDE problems under latent-marching and autoregressive protocols, comparing accuracy, efficiency, and attention structure. The results include 2D Kolmogorov flow at 256 × 256 resolution, and factorized attention shows improved stability, accuracy, and computational efficiency over linear attention in the studied comparison.

  • Benchmark problems: The benchmark covers 2D Kolmogorov flow, 3D isotropic turbulence, 3D smoke buoyancy, and 2D Darcy flow.Kolmogorov flow uses a 256 × 256 grid with Re = 1000 and predicts 16 frames over 1 second.
  • Training protocols: All models are evaluated with Latent Marching and Autoregressive training protocols, using relative L2 error, rollout cost, and parameter count as reported measures.Latent Marching predicts multiple future steps per model call, with k = 4 for Kolmogorov flow and smoke buoyancy and k = 2 for isotropic turbulence.
  • Comparison against full attention: Factorized attention produces more stable rollouts and better accuracy than full attention in the studied Kolmogorov-flow comparison.The authors attribute the degraded full-attention result to instability when repeatedly calculating a large attention matrix during rollout.
  • Comparison against full attention: Factorized attention is computationally more efficient than linear attention, with consistent efficiency improvements across the reported grid and model sizes.The benchmark compares forward and backward computational cost, peak memory, encoding time, and propagation time.
  • Comparison against full attention: Less than 5% of linear-attention singular values capture over 90% of total spectral energy, indicating a highly low-rank attention structure.The full-attention spectrum is analyzed with truncated SVD because its matrix size is 16384 × 16384.

5 Conclusion

FactFormer combines a learnable projection operator with a factorized kernel integral for efficient PDE surrogate modeling, while retaining accuracy. The method remains subject to dimensionality and rollout-stability limitations.

  • FactFormer uses a learnable projection operator and factorized kernel integral to balance efficiency and accuracy in PDE surrogate modeling.
  • The factorized kernel integral still evaluates the function on all S1×S2×...×Sm grid points, so it is not free from the curse of dimensionality.
  • Applying the model to non-uniform grids requires non-uniform quadrature weights because the current projection uses mean pooling and exploits uniform grid structure.
  • FactFormer and other neural PDE solvers can become unstable on time-dependent systems because errors accumulate during solving.

A Model implementation details

The implementation specifies FactFormer's hyperparameters, positional encoding, kernel representation, optimization, and convolutional components for selected boundary-value problems.

  • FactFormer uses random Fourier positional encoding with a learnable linear transformation before each attention layer.The encoding is added to each latent representation using its Cartesian coordinate.
  • Kernel dimension is the number of function bases used by each attention head to compute the kernel.The kernel is represented as a sum over dk basis functions.
  • FactFormer is trained with AdamW and a cyclic learning-rate scheduler whose maximum learning rate is 3e-4.
  • For 3D smoke buoyancy and 2D Darcy flow, a four-layer U-shaped CNN block follows attention layers to better account for boundary values.All convolutional layers use kernel size 3 and padding 1; the first layer uses stride 2.
  • Every model compresses temporal inputs with a 2D convolution and uses a three-layer MLP to project the latent encoding.The temporal convolution uses filters of size (1, Tin).

B Baseline implementation details

The baselines comprise FNO, F-FNO, Dil-ResNet, and a linear-attention Transformer, implemented with specified repositories, architectures, normalization, and training schedules.

  • The baseline set includes FNO, F-FNO, Dil-ResNet, and a linear-attention Transformer.
  • FNO and F-FNO use hidden size 96, four layers, and task-dependent Fourier mode counts, with group normalization before the final fully connected layer.The mode numbers are 12 for 3D problems, 24 for 2D turbulence, and 20 for Darcy flow.
  • Dil-ResNet uses three residual blocks with five CNN layers of width 128 and truncated dilation patterns.Without truncation, 3D training would take over a week and cannot fit on a single A6000 GPU for two-step rollout training.
  • The linear-attention Transformer follows OFormer's attention implementation with Galerkin-style normalization.
  • Experiments use PyTorch 1.8 except for FNO and F-FNO, which use PyTorch 1.13, with 100k iterations for language-modeling setups and 64k for autoregressive setups.

C Visualization of error trend

This section visualizes average frame-wise error trends for time-dependent systems, distinguishing autoregressive and language-modeling results with dashed and solid lines.

  • The section presents average frame-wise error trends for the investigated time-dependent systems.
  • The plots report error trends for vorticity in 2D Kolmogorov flow, pressure and velocity in 3D isotropic turbulence, and marker field in 3D smoke buoyancy.
  • Dashed lines denote autoregressive results, while solid lines denote language-modeling results.

D Further ablation study

The further experiments examine runtime scaling, hyperparameter effects, learned kernels, and training stability. Factorized attention is more computationally efficient than linear attention in the reported 2D and 3D benchmarks, while learned kernels show structured patterns and runs converge similarly across seeds.

  • Model scaling performance: The number of attention heads has a crucial impact on final prediction loss in the 2D Kolmogorov-flow ablation.The experiments use data sampled from a 128 × 128 grid and vary one hyperparameter direction at a time.
  • Visualization of learned kernels: Rotary positional encoding produces stationary learned kernels whose values depend on relative distance between spatial points.The kernels are visualized in Figure 14.
  • Visualization of learned kernels: The learned kernel matrices exhibit symmetric patterns and are diagonal dominated despite the non-symmetric nature of QK^T.
  • Influence of random seed: Three random seeds converge to similar loss levels with only marginal differences.The comparison is shown through averaged frame-wise loss trends.

E Dataset details

The dataset details describe four PDE settings: 2D Kolmogorov flow, 3D isotropic turbulence, 3D smoke buoyancy, and 2D Darcy flow. They specify governing dynamics, boundary or forcing conditions, numerical solvers, resolutions, and train/test splits.

  • 2D Kolmogorov flow: The 2D Kolmogorov-flow dataset contains 100 training and 20 testing trajectories, each spanning 10 seconds and 160 frames.The flow is periodic in all spatial directions and uses forcing factor n = 8 with a dragging term 0.1ω(x).
  • 3D isotropic turbulence: The 3D isotropic-turbulence dataset uses periodic boundaries, Taylor Reynolds number 84, and 1000 training and 100 testing trajectories.Its reference spectral Galerkin simulation uses a 60 × 60 × 60 spatial resolution and 0.005s temporal resolution.
  • 2D Darcy flow: The 2D Darcy-flow dataset uses constant forcing, a Gaussian-random-field coefficient, and 1000 training and 100 testing samples.The data come from a second-order finite-difference solver on a 421 × 421 resolution grid.

F Results visualization

The visualization section presents predicted fields for the evaluated PDE problems, using a middle first-axis cross-section for 3D cases. The figures cover 2D flow fields and multiple physical variables in 3D turbulence and smoke buoyancy.

  • Visualization protocol: For 3D problems, visualizations show the cross-section at the middle of the first axis.
  • 2D Kolmogorov flow: The 2D Kolmogorov-flow results include two example samples.
  • 3D isotropic turbulence: The 3D isotropic-turbulence visualizations include pressure and the x-, y-, and z-components of velocity.
  • 3D smoke buoyancy: The 3D smoke-buoyancy visualizations include the smoke marker field and all three velocity components.
  • 2D Darcy flow: The 2D Darcy-flow visualization shows the flow field.

G Broader impact

The paper positions FactFormer as an efficient attention-based surrogate for numerical PDE solvers while acknowledging broad applicability and long-term-stability boundaries. It also identifies possible combinations with other solver architectures and illustrates the distinction from Axial Transformer.

  • Broader impact: The authors report that FactFormer balances accuracy and efficiency as a surrogate for numerical solvers, but no single data-driven model is guaranteed to suit every PDE.
  • Limitations: Long-term stability of the proposed model cannot be guaranteed, as with most concurrent neural PDE solvers.
  • Future combinations: The architecture could potentially be combined with output-variable relationships or common neural architectures such as U-Net.
  • Architecture comparison: The illustrative comparison distinguishes FactFormer from Axial Transformer while omitting positional encoding, multi-head, and softmax details for simplicity.The Axial Transformer example shows a column-wise attention block.
Loading 2305.17560v2…