TNT: Vision Transformer for Turbulence Simulations

TL;DR

This paper introduces TNT, a transformer-based neural network that efficiently simulates turbulence on coarser grids, outperforming traditional methods and demonstrating robustness and potential for broader application.

Contribution

The paper presents a novel transformer architecture, TNT, tailored for turbulence simulation, extending positional embeddings and introducing TMSA for improved spatiotemporal modeling.

Findings

TNT outperforms state-of-the-art U-net in turbulence prediction metrics.

TNT generates stable, long-range turbulence predictions.

Model robustness is confirmed across different initial conditions.

Abstract

Turbulence is notoriously difficult to model due to its multi-scale nature and sensitivity to small perturbations. Classical solvers of turbulence simulation generally operate on finer grids and are computationally inefficient. In this paper, we propose the Turbulence Neural Transformer (TNT), which is a learned simulator based on the transformer architecture, to predict turbulent dynamics on coarsened grids. TNT extends the positional embeddings of vanilla transformers to a spatiotemporal setting to learn the representation in the 3D time-series domain, and applies Temporal Mutual Self-Attention (TMSA), which captures adjacent dependencies, to extract deep and dynamic features. TNT is capable of generating comparatively long-range predictions stably and accurately, and we show that TNT outperforms the state-of-the-art U-net simulator on several metrics. We also test the model performance with different components removed and evaluate robustness to different initial conditions. Although more experiments are needed, we conclude that TNT has great potential to outperform existing solvers and generalize to additional simulation datasets.