[论文解读] UPS: Efficiently Building Foundation Models for PDE Solving via Cross-Modal Adaptation

Jun-Hong Shen, Tanya Marwah+1 more|arXiv (Cornell University)|Mar 11, 2024

BIM and Construction Integration被引用 5

一句话总结

UPS 将多样化的时空 PDE 数据统一为一个共同表示，并训练一个基于 Transformer 的模型，使预训练的 LLM 可以适应 PDE 求解，在 PDEBench 任务上实现最先进的结果，同时在数据效率方面表现出色。

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.

研究动机与目标

• 为在不同领域、维度和分辨率上解决多样化的时变 PDE 问题，提供一个数据高效的基础模型方法以建立数据基础。
• 提出一个统一的数据表示和一个三部分架构，结合 FNO 基 PDE 嵌入、LLMs 和一个预测头。
• 利用两阶段跨模态适配将 PDE 嵌入与文本嵌入对齐，并在多个 PDE 家族上进行微调。
• 在 PDEBench 上证明性能与样本效率的提升，具备强大的少样本及零样本迁移能力。

提出的方法

• 将多样的 PDE 路径统一到一个跨越多维度、分辨率和量纲的共享表示空间。
• 使用带有 Fourier Neural Operator (FNO) 层的 PDE 特定嵌入网络从 PDE 数据中提取分辨率不变的特征。
• 通过嵌入类文本的 PDE 描述来整合关于 PDE 的元数据，并在通过 LLM 体处理前将其与 PDE 特征拼接。
• 采用三段式结构：PDE 嵌入网络、作为模型主体的预训练 LLM 层，以及用于下一个 PDE 状态的线性预测头。
• 分两阶段训练：（i）带有模态对齐损失（基于 MMD）和任务损失（nRMSE）的嵌入预训练；（ii）在多个 PDE 家族上的多任务微调。
• 展示与各种 LLM 主干（RoBERTa、T5、CLIP）兼容性，并展现神经算子在网格与分辨率上的不变性。

Figure 1 : UPS outperforms leading competitors on a wide range of tasks from PDEBench, including the single-task neural operators U-Net (Ronneberger et al., 2015 ) , FNO (Li et al., 2020a ) , and ORCA (Shen et al., 2023 ) , and the multi-task MPP-Base (McCabe et al., 2023 ) (MPP only applies to 2D t

实验结果

研究问题

• RQ1 统一的表示和架构是否能够实现跨领域、跨维度、跨分辨率的迁移？
• RQ2 从预训练 LLM 的跨模态适配是否能够提升对多样 PDE 求解任务的数据效率和预测准确性？
• RQ3 部分元数据（文本形式的 PDE 描述）对神经算子框架中的学习与泛化有何影响？
• RQ4 在对未见 PDE 家族、系数和分辨率的少样本和零样本迁移下，UPS 的性能如何？

主要发现

• UPS 在 10 个 PDEBench 任务中的前 8 项达到最先进水平，且往往比基线有较大优势。
• 在同分布任务中，UPS（RoBERTa-Base）在 7 个数据集中的 6 个数据集上取得最低的 nRMSE，并表现出一维和二维任务的强劲性能。
• UPS 在对未见 PDE 家族、系数及更高分辨率的少样本和零样本迁移中表现出色，超越在数据有限条件下的多个基线。
• 阶段-1 的嵌入预训练（带模态对齐和任务学习）对性能至关重要，元数据的纳入进一步提升了结果。
• 该方法数据和计算高效，每个 PDE 家族的训练轨迹不到 5K，且能够在单 GPU 上完成训练。

Figure 2 : To adapt pretrained LLMs for PDE solving, UPS first expresses PDE trajectories of different families, dimensions, channels, and resolutions using a unified data representation (left panel). It then processes the data with a unified transformer architecture that integrates domain-specific