BCAT: A Block Causal Transformer for PDE Foundation Models for Fluid Dynamics

TL;DR

BCAT is a novel block causal transformer model that significantly improves the accuracy of fluid dynamics solution predictions across diverse datasets and tasks by capturing spatial dependencies more effectively.

Contribution

The paper introduces BCAT, a block causal transformer architecture tailored for PDE-based fluid dynamics prediction, demonstrating superior accuracy and adaptability over prior methods.

Findings

3.5x accuracy improvement in next frame prediction

Average relative error of 1.18% across tasks

40% better accuracy after fine-tuning on turbulence data

Abstract

We introduce BCAT, a PDE foundation model designed for autoregressive prediction of solutions to two dimensional fluid dynamics problems. Our approach uses a block causal transformer architecture to model next frame predictions, leveraging previous frames as contextual priors rather than relying solely on sub-frames or pixel-based inputs commonly used in image generation methods. This block causal framework more effectively captures the spatial dependencies inherent in nonlinear spatiotemporal dynamics and physical phenomena. In an ablation study, next frame prediction demonstrated a 3.5x accuracy improvement over next token prediction. BCAT is trained on a diverse range of fluid dynamics datasets, including incompressible and compressible Navier-Stokes equations across various geometries and parameter regimes, as well as the shallow-water equations. The model's performance was evaluated on 6 distinct downstream prediction tasks and tested on about 8K trajectories to measure robustness on a variety of fluid dynamics simulations. BCAT achieved an average relative error of 1.18% across all evaluation tasks, outperforming prior approaches on standard benchmarks. With fine-tuning on a turbulence dataset, we show that the method adapts to new settings with more than 40% better accuracy over prior methods.