Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AMR-Transformer: Enabling Efficient Long-range Interaction for Complex Neural Fluid Simulation

Zeyi Xu, Jinfan Liu, Kuangxu Chen, Ye Chen, Zhangli Hu, Bingbing Ni

TL;DR

AMR-Transformer tackles the challenge of modeling long-range interactions in high-resolution CFD by coupling an adaptive mesh refinement (AMR) tokenizer with a constraint-aware pruning heuristic and an encoder-only Transformer solver. The AMR tokenizer creates multi-scale patches guided by Navier-Stokes constraints, enabling focused computation on dynamically rich regions, while the Transformer captures global dependencies across patches. Extensive experiments on CFDBench, PDEBench, and a new 128×128 shockwave dataset show state-of-the-art accuracy with dramatically reduced token counts (2–10×) and substantial FLOP reductions (up to 60× vs ViT), demonstrating strong efficiency-accuracy tradeoffs for complex fluid phenomena. This approach enables high-fidelity, long-range CFD simulations at scale, with practical impact for turbulence, shocks, and other multi-scale flow regimes.

Abstract

Accurately and efficiently simulating complex fluid dynamics is a challenging task that has traditionally relied on computationally intensive methods. Neural network-based approaches, such as convolutional and graph neural networks, have partially alleviated this burden by enabling efficient local feature extraction. However, they struggle to capture long-range dependencies due to limited receptive fields, and Transformer-based models, while providing global context, incur prohibitive computational costs. To tackle these challenges, we propose AMR-Transformer, an efficient and accurate neural CFD-solving pipeline that integrates a novel adaptive mesh refinement scheme with a Navier-Stokes constraint-aware fast pruning module. This design encourages long-range interactions between simulation cells and facilitates the modeling of global fluid wave patterns, such as turbulence and shockwaves. Experiments show that our approach achieves significant gains in efficiency while preserving critical details, making it suitable for high-resolution physical simulations with long-range dependencies. On CFDBench, PDEBench and a new shockwave dataset, our pipeline demonstrates up to an order-of-magnitude improvement in accuracy over baseline models. Additionally, compared to ViT, our approach achieves a reduction in FLOPs of up to 60 times.

AMR-Transformer: Enabling Efficient Long-range Interaction for Complex Neural Fluid Simulation

TL;DR

AMR-Transformer tackles the challenge of modeling long-range interactions in high-resolution CFD by coupling an adaptive mesh refinement (AMR) tokenizer with a constraint-aware pruning heuristic and an encoder-only Transformer solver. The AMR tokenizer creates multi-scale patches guided by Navier-Stokes constraints, enabling focused computation on dynamically rich regions, while the Transformer captures global dependencies across patches. Extensive experiments on CFDBench, PDEBench, and a new 128×128 shockwave dataset show state-of-the-art accuracy with dramatically reduced token counts (2–10×) and substantial FLOP reductions (up to 60× vs ViT), demonstrating strong efficiency-accuracy tradeoffs for complex fluid phenomena. This approach enables high-fidelity, long-range CFD simulations at scale, with practical impact for turbulence, shocks, and other multi-scale flow regimes.

Abstract

Accurately and efficiently simulating complex fluid dynamics is a challenging task that has traditionally relied on computationally intensive methods. Neural network-based approaches, such as convolutional and graph neural networks, have partially alleviated this burden by enabling efficient local feature extraction. However, they struggle to capture long-range dependencies due to limited receptive fields, and Transformer-based models, while providing global context, incur prohibitive computational costs. To tackle these challenges, we propose AMR-Transformer, an efficient and accurate neural CFD-solving pipeline that integrates a novel adaptive mesh refinement scheme with a Navier-Stokes constraint-aware fast pruning module. This design encourages long-range interactions between simulation cells and facilitates the modeling of global fluid wave patterns, such as turbulence and shockwaves. Experiments show that our approach achieves significant gains in efficiency while preserving critical details, making it suitable for high-resolution physical simulations with long-range dependencies. On CFDBench, PDEBench and a new shockwave dataset, our pipeline demonstrates up to an order-of-magnitude improvement in accuracy over baseline models. Additionally, compared to ViT, our approach achieves a reduction in FLOPs of up to 60 times.

Paper Structure

This paper contains 17 sections, 13 equations, 3 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. AI for CFD
  4. Accumulation Approaches
  5. Methodology
  6. AMR Tokenizer
  7. N-S Constraint-Aware Pruning Module
  8. Transformer Neural Solver
  9. Experiment
  10. Implementation
  11. CFDBench AND PDEBench Benchmark
  12. Shockwave Dataset
  13. Comparative Analysis with SOTA
  14. Analysis of AMR tokenizer
  15. Ablation Study
  16. ...and 2 more sections

Figures (3)

  • Figure 1: AMR tokenizer results in high-resolution simulation. Visualization of the AMR tokenizer applied to 1024 × 1024 shock wave and explosion simulations. The AMR tokenizer captures fine-scale structures while reducing token count. Right panels compare the regular 512 × 512 grid with the AMR partitioning, each right panel displays the total cell count and partitioning scheme (bottom-right), along with the mean squared error (MSE) relative to the 1024 × 1024 ground truth (top-left).
  • Figure 2: AMR Tokenizer: The AMR tokenizer adaptively partitions the input domain $I$, where $H \times W \times c$ represents the full spatial resolution of the domain with $c$ channels of features. AMR tokenizer refines the mesh progressively, from low resolution to high resolution, depth by depth. At each depth, cells undergo subdivision and storage based on customized Navier-Stokes constraints, including the velocity gradient, vorticity, momentum, and Kelvin-Helmholtz instability. Grayscale cells marked with a colored × fail the corresponding physical constraint, meaning they are discarded and neither stored nor subdivided. Cells that pass the corresponding physical constraints are marked with the color of the most sensitive constraint, and stored and subdivided in the quadtree structure. The stored cells are aggregated by averaging (shown as "Storage (AVG)"), producing a compact multi-resolution representation of the domain. The second row shows the mesh results of progressive subdivision at each depth.
  • Figure 3: Visualization of five problems: Each column represents one of the five benchmark problems (Cavity, Cylinder, Dam, Shockwave, and Tube). These visualizations highlight the diversity and complexity of the flow scenarios, from boundary-driven flows in Cavity, vortex shedding around the Cylinder, rapid changes in the Dam break, steep gradients in Shockwave, to confined flows in the Tube.