Zebra : Pré-entraînement en contexte et génératif pour résoudre des EDP paramétriques

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

October 4, 2024

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

cs.AI

Résumé

La résolution d'équations aux dérivées partielles (EDP) paramétriques dépendantes du temps est complexe, car les modèles doivent s'adapter aux variations des paramètres tels que les coefficients, les termes de forçage et les conditions aux limites. Les solveurs neuronaux basés sur les données s'entraînent soit sur des données échantillonnées à partir de la distribution des paramètres des EDP dans l'espoir que le modèle généralise à de nouvelles instances, soit s'appuient sur une adaptation basée sur les gradients et l'apprentissage méta pour encoder implicitement les dynamiques à partir des observations. Cela s'accompagne souvent d'une complexité accrue de l'inférence. Inspiré par les capacités d'apprentissage en contexte des grands modèles de langage (LLM), nous présentons Zebra, un nouveau transformateur auto-régressif génératif conçu pour résoudre des EDP paramétriques sans nécessiter d'adaptation de gradient lors de l'inférence. En exploitant les informations en contexte à la fois lors de la pré-formation et de l'inférence, Zebra s'adapte dynamiquement à de nouvelles tâches en se basant sur des séquences d'entrée qui intègrent des trajectoires de contexte ou des états précédents. Cette approche permet à Zebra de gérer de manière flexible des entrées de contexte de taille arbitraire et de prendre en charge la quantification de l'incertitude en échantillonnant plusieurs trajectoires de solution. Nous évaluons Zebra dans divers scénarios d'EDP complexes, démontrant sa capacité d'adaptation, sa robustesse et ses performances supérieures par rapport aux approches existantes.

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.