Zebra: Pré-treinamento Contextual e Generativo 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

Resumo

Resolver equações diferenciais parciais (EDPs) paramétricas dependentes do tempo é desafiador, pois os modelos precisam se adaptar às variações nos parâmetros, como coeficientes, termos de forçamento e condições de contorno. Solucionadores neurais orientados por dados treinam em dados amostrados da distribuição de parâmetros da EDP na esperança de que o modelo generalize para novas instâncias ou dependem de adaptação baseada em gradientes e meta-aprendizado para codificar implicitamente a dinâmica a partir de observações. Isso frequentemente resulta em maior complexidade de inferência. Inspirados pelas capacidades de aprendizado em contexto de grandes modelos de linguagem (LLMs), apresentamos Zebra, um novo transformador auto-regressivo generativo projetado para resolver EDPs paramétricas sem necessidade de adaptação de gradientes na inferência. Ao alavancar informações em contexto durante o pré-treinamento e a inferência, Zebra se adapta dinamicamente a novas tarefas condicionando-se a sequências de entrada que incorporam trajetórias de contexto ou estados anteriores. Essa abordagem permite que Zebra lide flexivelmente com entradas de contexto de tamanho arbitrário e suporte a quantificação de incerteza por meio da amostragem de múltiplas trajetórias de solução. Avaliamos Zebra em uma variedade de cenários desafiadores de EDP, demonstrando sua adaptabilidade, robustez e desempenho superior em comparação com abordagens 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.