[论文解读] Poseidon: Efficient Foundation Models for PDEs

Maximilian Herde, Bogdan Raonić+5 more|arXiv (Cornell University)|May 29, 2024

Numerical methods for differential equations被引用 9

一句话总结

Poseidon 是一种用于学习偏微分方程解算子的基础模型，在流体动力学数据上进行预训练，具有高样本效率和对未知物理规律的良好泛化，获得强劲的下游性能。

ABSTRACT

We introduce Poseidon, a foundation model for learning the solution operators of PDEs. It is based on a multiscale operator transformer, with time-conditioned layer norms that enable continuous-in-time evaluations. A novel training strategy leveraging the semi-group property of time-dependent PDEs to allow for significant scaling-up of the training data is also proposed. Poseidon is pretrained on a diverse, large scale dataset for the governing equations of fluid dynamics. It is then evaluated on a suite of 15 challenging downstream tasks that include a wide variety of PDE types and operators. We show that Poseidon exhibits excellent performance across the board by outperforming baselines significantly, both in terms of sample efficiency and accuracy. Poseidon also generalizes very well to new physics that is not seen during pretraining. Moreover, Poseidon scales with respect to model and data size, both for pretraining and for downstream tasks. Taken together, our results showcase the surprising ability of Poseidon to learn effective representations from a very small set of PDEs during pretraining in order to generalize well to unseen and unrelated PDEs downstream, demonstrating its potential as an effective, general purpose PDE foundation model. Finally, the Poseidon model as well as underlying pretraining and downstream datasets are open sourced, with code being available at https://github.com/camlab-ethz/poseidon and pretrained models and datasets at https://huggingface.co/camlab-ethz.

研究动机与目标

• 动机：在PDE中引入基础模型，以提升相对于特定任务的神经算子在样本效率上的表现。
• 介绍 Poseidon，一个为PDE解算子量身定制的可扩展基础模型架构。
• 证明在多样化的PDE数据上进行预训练能够在未见PDE和物理现象上实现强泛化。
• 展示 Poseidon 如何随模型规模和数据规模扩展，并提供开源数据集与代码。

提出的方法

• 使用 scOT，一种带前瞻条件的分层多尺度视觉变换器来近似 PDE 解算子 S(t,a)。
• 引入时间条件化层归一化，以实现时间上的连续评估。
• 应用 all2all 训练策略，利用时变PDE的半群性质从轨迹生成更多训练对。
• 在一个大规模、多样化的欧拉/纳维-斯托克斯算子数据集上对 Poseidon 进行预训练，然后在下游任务上进行微调。
• 在最终时刻使用相对L1误差对15个多样化PDE任务进行评估，包括分布外情形。

实验结果

研究问题

• RQ1 一个基于PDE的基础模型能否在少量PDE上进行预训练，从而学习到能泛化到未见PDE和物理现象的表示？
• RQ2 架构、数据规模和模型规模如何影响下游性能和样本效率？
• RQ3 通过 all2all 训练利用半群性质是否能提高学习PDE算子的数据效率？
• RQ4 通过将时间独立的PDE解释为长期极限，Poseidon 在多大程度上能够迁移？
• RQ5 在多样的下游任务中，Poseidon 与特定任务的神经算子及其他PDE基础模型相比如何？

主要发现

• Poseidon 在全部15个下游任务上，在准确性和样本效率方面均优于基线。
• 平均而言，Poseidon 需要大约20个特定任务样本即可达到时间相关PDE用1024个样本的FNO误差（对于时间独立PDE为4096个样本）的水平。
• Poseidon 能很好地泛化到未见PDE和物理现象，包括在预训练中未出现的任务，仅需要少量下游样本。
• 模型规模和数据集规模对下游任务的性能和样本效率均有正向影响。
• 预训练的多样性（数据质量和多样性）对大多数任务的下游准确性有显著影响。