Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Junhong Shen, Tanya Marwah, Ameet Talwalkar

TL;DR

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.

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.

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

TL;DR

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.

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.
Paper Structure (36 sections, 13 equations, 6 figures, 9 tables)

This paper contains 36 sections, 13 equations, 6 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. Unified Data Representation
  5. Unified Architecture
  6. Full Workflow and Training
  7. Experiments
  8. Achieving State-of-the-Art Results on PDEBench with Compute Efficiency
  9. Generalizing to Unseen PDEs with Data Efficiency
  10. Ablation Studies
  11. Conclusion and Future Work
  12. Appendix
  13. Datasets
  14. Burgers Equation (1D)
  15. Diffusion-Sorption Equation (1D)
  16. ...and 21 more sections

Figures (6)

  • Figure 1: To adapt pretrained LLMs for PDE solving, UPS first transforms PDE of different dimensions, channels, and resolutions into a unified representation (left panel). Then, the data is processed with a unified architecture that integrates FNO layers, PDE metadata, and LLMs (right panel). The architecture is trained in two stages. In stage 1, we pretrain the embedding network using a joint loss that simultaneously optimizes (i) the distribution similarity between PDE features and text embeddings to align the modalities, and (ii) the prediction performance of extracted PDE features. In stage 2, we fine-tune the entire model on a dataset that combines multiple families of spatiotemporal PDEs with varying domain dimensions, initial/boundary conditions, and coefficients. UPS achieves competitive results with significantly better sample-efficiency than existing methods.
  • Figure :
  • Figure :
  • Figure :
  • Figure :
  • ...and 1 more figures