Zebra：用于解决参数化偏微分方程的上下文和生成式预训练

Zebra: In-Context and Generative Pretraining for Solving Parametric PDEs

October 4, 2024

作者: Louis Serrano, Armand Kassaï Koupaï, Thomas X Wang, Pierre Erbacher, Patrick Gallinari

cs.AI

摘要

解决时变参数偏微分方程（PDEs）是具有挑战性的，因为模型必须适应参数变化，如系数、强制项和边界条件。数据驱动的神经求解器要么在从PDE参数分布中采样的数据上进行训练，希望模型推广到新实例，要么依赖基于梯度的调整和元学习来从观察中隐式编码动态。这往往伴随着增加的推断复杂性。受大型语言模型（LLMs）的上下文学习能力启发，我们引入了Zebra，这是一种新颖的生成自回归变压器，旨在解决参数化PDEs，而无需在推断时进行梯度调整。通过在预训练和推断期间利用上下文信息，Zebra通过对输入序列进行条件化，动态适应新任务，这些输入序列融合了上下文轨迹或先前状态。这种方法使Zebra能够灵活处理任意大小的上下文输入，并通过对多个解轨迹进行采样来支持不确定性量化。我们在各种具有挑战性的PDE场景中评估了Zebra，展示了其适应性、稳健性和与现有方法相比的卓越性能。

English

Solving time-dependent parametric partial differential equations (PDEs) is challenging, as models must adapt to variations in parameters such as coefficients, forcing terms, and boundary conditions. Data-driven neural solvers either train on data sampled from the PDE parameters distribution in the hope that the model generalizes to new instances or rely on gradient-based adaptation and meta-learning to implicitly encode the dynamics from observations. This often comes with increased inference complexity. Inspired by the in-context learning capabilities of large language models (LLMs), we introduce Zebra, a novel generative auto-regressive transformer designed to solve parametric PDEs without requiring gradient adaptation at inference. By leveraging in-context information during both pre-training and inference, Zebra dynamically adapts to new tasks by conditioning on input sequences that incorporate context trajectories or preceding states. This approach enables Zebra to flexibly handle arbitrarily sized context inputs and supports uncertainty quantification through the sampling of multiple solution trajectories. We evaluate Zebra across a variety of challenging PDE scenarios, demonstrating its adaptability, robustness, and superior performance compared to existing approaches.