Zebra: Preallenamento In-Context e Generativo per la Risoluzione di PDE Parametriche

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

October 4, 2024

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

cs.AI

Abstract

Risolvere equazioni differenziali parziali (PDE) parametriche dipendenti dal tempo è una sfida, poiché i modelli devono adattarsi alle variazioni nei parametri come coefficienti, termini forzanti e condizioni al contorno. I risolutori neurali basati sui dati vengono addestrati su dati campionati dalla distribuzione dei parametri delle PDE nella speranza che il modello generalizzi a nuove istanze o si basano sull'adattamento basato sul gradiente e sull'apprendimento meta per codificare implicitamente la dinamica dalle osservazioni. Questo spesso comporta un aumento della complessità inferenziale. Ispirati alle capacità di apprendimento contestuale dei grandi modelli linguistici (LLM), presentiamo Zebra, un nuovo trasformatore auto-regressivo generativo progettato per risolvere PDE parametriche senza richiedere adattamenti del gradiente durante l'inferenza. Sfruttando le informazioni contestuali durante sia la preformazione che l'inferenza, Zebra si adatta dinamicamente a nuovi compiti condizionandosi su sequenze di input che incorporano traiettorie di contesto o stati precedenti. Questo approccio consente a Zebra di gestire in modo flessibile input di contesto di dimensioni arbitrariamente grandi e supporta la quantificazione dell'incertezza attraverso il campionamento di diverse traiettorie di soluzione. Valutiamo Zebra in una varietà di scenari PDE impegnativi, dimostrandone l'adattabilità, la robustezza e le prestazioni superiori rispetto agli approcci esistenti.

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.