Zebra: パラメトリックPDEの解決のためのインコンテキストおよび生成プリトレーニング

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

要旨

時間依存パラメトリック偏微分方程式（PDE）を解くことは困難であり、係数、強制項、境界条件などのパラメータの変動に適応する必要があります。データ駆動型ニューラルソルバーは、PDEパラメータの分布からサンプリングされたデータでトレーニングするか、勾配ベースの適応やメタラーニングに依存して観測から動態を暗黙的にエンコードすることで、新しいインスタンスに汎化することを期待します。これには推論の複雑さが伴います。大規模言語モデル（LLM）のインコンテキスト学習能力に触発され、我々はZebraという新しい生成オートレグレッシブトランスフォーマーを紹介します。Zebraは、勾配適応を推論時に必要とせずにパラメトリックPDEを解決するよう設計されています。事前トレーニングと推論の両方でインコンテキスト情報を活用することで、Zebraは、入力シーケンスに基づいて動的に新しいタスクに適応し、コンテキストの軌跡や直前の状態を組み込んだ入力シーケンスに条件付けることで、任意のサイズのコンテキスト入力を柔軟に処理し、複数の解軌跡をサンプリングすることによる不確実性の定量化をサポートします。我々は、Zebraをさまざまな困難なPDEシナリオで評価し、既存の手法と比較してその適応性、堅牢性、および優れた性能を示しています。

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.