UPS: Efficiently Building Foundation Models for PDE Solving via Cross-Modal Adaptation

Junhong Shen, Tanya Marwah, Ameet Talwalkar

Keywords: PDE Solving, Neural Operators, Large Language Models, Cross-Modal Adaptation, Fine-Tuning

TL;DR: We introduce UPS (Unified PDE Solvers), an effective and data/compute-efficient approach to learning unified neural operators for diverse spatiotemporal PDE families using pretrained LLMs.

Abstract: We present Unified PDE Solvers (UPS), a data- and compute-efficient approach to developing unified neural operators for diverse families of spatiotemporal PDEs from various domains, dimensions, and resolutions. UPS embeds different PDEs into a shared representation space and processes them using a FNO-transformer architecture. Rather than training the network from scratch, which is data-demanding and computationally expensive, we warm-start the transformer from pretrained LLMs and perform explicit alignment to reduce the modality gap while improving data and compute efficiency. The cross-modal UPS achieves state-of-the-art results on a wide range of 1D and 2D PDE families from PDEBench, outperforming existing unified models using 4 times less data and 26 times less compute. Meanwhile, it is capable of few-shot transfer to unseen PDE families and coefficients.

Submission Number: 99