Zebra: Preentrenamiento Generativo y en Contexto para Resolver EDPs Paramétricas

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

October 4, 2024

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

cs.AI

Resumen

Resolver ecuaciones en derivadas parciales (EDP) paramétricas dependientes del tiempo es desafiante, ya que los modelos deben adaptarse a variaciones en parámetros como coeficientes, términos forzantes y condiciones de contorno. Los solucionadores neuronales basados en datos se entrenan con datos muestreados de la distribución de parámetros de las EDP con la esperanza de que el modelo se generalice a nuevas instancias, o bien se basan en adaptación basada en gradientes y meta-aprendizaje para codificar implícitamente la dinámica a partir de observaciones. Esto a menudo conlleva una mayor complejidad de inferencia. Inspirados en las capacidades de aprendizaje en contexto de los grandes modelos de lenguaje (LLMs), presentamos Zebra, un novedoso transformador auto-regresivo generativo diseñado para resolver EDP paramétricas sin necesidad de adaptación de gradientes en inferencia. Al aprovechar la información en contexto durante el preentrenamiento y la inferencia, Zebra se adapta dinámicamente a nuevas tareas condicionándose a secuencias de entrada que incorporan trayectorias de contexto o estados previos. Este enfoque permite a Zebra manejar de manera flexible entradas de contexto de tamaño arbitrario y soporta la cuantificación de incertidumbre mediante el muestreo de múltiples trayectorias de solución. Evaluamos Zebra en una variedad de escenarios desafiantes de EDP, demostrando su adaptabilidad, robustez y rendimiento superior en comparación con enfoques existentes.

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.