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.
翻译:我们提出了Poseidon,一个用于学习偏微分方程解算子的基础模型。该模型基于多尺度算子Transformer架构,通过时间条件层归一化技术实现连续时间评估。我们还提出了一种新颖的训练策略,利用时间相关偏微分方程的半群特性,实现了训练数据的大规模扩展。Poseidon在涵盖流体动力学控制方程的多样化大规模数据集上进行了预训练。随后在包含15个挑战性下游任务的测试套件上进行了评估,这些任务涵盖了多种偏微分方程类型和算子。实验表明,Poseidon在样本效率和精度方面均显著超越基线模型,展现出全面优异的性能。该模型对预训练阶段未见过的新物理现象也表现出卓越的泛化能力。此外,Poseidon在预训练和下游任务中均展现出模型规模与数据规模的协同扩展特性。综合而言,我们的研究结果揭示了Poseidon能够从预训练阶段极少量偏微分方程数据中学习有效表示,并泛化至未见过的非相关下游偏微分方程,这证明了其作为通用偏微分方程基础模型的巨大潜力。最后,Poseidon模型及其底层预训练与下游数据集均已开源,代码发布于https://github.com/camlab-ethz/poseidon,预训练模型与数据集发布于https://huggingface.co/camlab-ethz。