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

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.

• The paper introduces UPS, a novel unified PDE solver that leverages cross-modal LLM adaptation to significantly reduce data and compute requirements.
• The methodology standardizes diverse PDE representations and employs a two-stage training process integrating pretrained LLMs with domain-specific neural operator layers.
• Empirical results demonstrate UPS achieves state-of-the-art performance on PDEBench with remarkable few-shot learning capability across various 1D and 2D tasks.

Unified Neural Operators for Diverse Spatiotemporal PDEs via LLM Adaptation

Introduction to Unified PDE Solvers

The quest for solving partial differential equations (PDEs) is central to numerous scientific and engineering disciplines. Traditionally approached through analytical techniques or numerical methods, the rise of data-driven solutions, notably Deep Learning (DL), has opened up new avenues for addressing these complex problems. While DL-based methods like neural operators have shown promise in approximating solution maps to PDE families, they typically require training a unique model for each PDE family, involving substantial data and computational resources. Recent efforts towards developing foundation models seek to train a unified model capable of transferring across PDE families, albeit with significant data and computational demands.

This work introduces Unified PDE Solver (\ourmodel), a novel approach leveraging the prowess of LLMs and generative AI for solving a wide spectrum of spatiotemporal PDEs. By unifying PDEs into a consistent representation and ingeniously incorporating LLMs into the operator learning, \ourmodel demonstrates effective and data-efficient learning across various PDE families.

\ourmodel Methodology

\ourmodel tackles the challenge of processing diverse PDE data by proposing a standardized data representation and an innovative LLM-based network architecture. The unified data representation bridges PDEs of different dimensions, ensuring a homogenized input format. The ensuing unified network architecture synergistically combines pretrained LLMs with domain-specific FNO layers, effectively processing the unified PDE data.

A significant contribution of this work is the two-stage cross-modal adaptation process. Initial training aims at aligning the LLMs’ understanding of PDEs, while subsequent fine-tuning harnesses multitask learning across various PDE tasks. This methodology not only facilitates leveraging the vast knowledge base of LLMs but also enhances performance with considerably fewer training samples.

Empirical Validation and Implications

\ourmodel's efficacy is rigorously benchmarked on the PDEBench framework, encompassing a wide range of 1D and 2D PDE tasks. The model exhibits superior performance, achieving state-of-the-art results on multiple benchmarks and demonstrating remarkable few-shot learning capabilities on unseen PDE families and conditions. These results underscore the potential of \ourmodel to serve as a generalized solver for complex physical systems.

Discussion and Outlook

This work represents a pivotal step towards realizing generalized foundation models for solving PDEs efficiently. By effectively adapting pretrained LLMs to the domain of PDEs, we not only achieve impressive empirical results but also dramatically reduce the data and compute requirements typically associated with training unified neural PDE solvers from scratch.

The successful adaptation of LLMs to PDE solving, as demonstrated in this study, opens up promising avenues for future research. Extending this approach to higher-dimensional PDEs and other types of physical systems, as well as further exploring the potential of LLMs in solving inverse problems, represents exciting directions for advancing the field.

\ourmodel strikes a balance between leveraging existing AI advancements and tailoring solutions to the computational physics domain. As the field of AI continues to evolve, particularly in the field of LLMs, \ourmodel presents a scalable and efficient framework for benefiting from these advancements, pushing the boundaries of what is achievable in computational physics and beyond.