CViT: Continuous Vision Transformer for Operator Learning

TL;DR

CViT introduces a novel neural operator architecture combining vision transformers with grid-based embeddings and cross-attention, enabling flexible, high-performance learning of complex physical systems across various PDEs.

Contribution

The paper presents CViT, a new neural operator architecture that adapts vision transformer techniques for operator learning in physical sciences, demonstrating superior performance without extensive pretraining.

Findings

Achieves state-of-the-art results on multiple PDE benchmarks.

Effectively handles discontinuities and multi-scale features.

Outperforms larger models with less pretraining.

Abstract

Operator learning, which aims to approximate maps between infinite-dimensional function spaces, is an important area in scientific machine learning with applications across various physical domains. Here we introduce the Continuous Vision Transformer (CViT), a novel neural operator architecture that leverages advances in computer vision to address challenges in learning complex physical systems. CViT combines a vision transformer encoder, a novel grid-based coordinate embedding, and a query-wise cross-attention mechanism to effectively capture multi-scale dependencies. This design allows for flexible output representations and consistent evaluation at arbitrary resolutions. We demonstrate CViT's effectiveness across a diverse range of partial differential equation (PDE) systems, including fluid dynamics, climate modeling, and reaction-diffusion processes. Our comprehensive experiments show that CViT achieves state-of-the-art performance on multiple benchmarks, often surpassing larger foundation models, even without extensive pretraining and roll-out fine-tuning. Taken together, CViT exhibits robust handling of discontinuous solutions, multi-scale features, and intricate spatio-temporal dynamics. Our contributions can be viewed as a significant step towards adapting advanced computer vision architectures for building more flexible and accurate machine learning models in the physical sciences.