Zebra: In-Context and Generative Pretraining for Solving Parametric PDEs
Louis Serrano, Armand Kassaï Koupaï, Thomas X Wang, Pierre Erbacher, Patrick Gallinari
2024-10-13
Summary
This paper introduces Zebra, a new method for solving time-dependent parametric partial differential equations (PDEs) using in-context learning and generative pretraining.
What's the problem?
Solving parametric PDEs can be very challenging because these equations depend on various factors like coefficients and boundary conditions, which can change over time. Traditional methods often require complex adjustments and can struggle to adapt to new situations, making them less efficient and harder to use.
What's the solution?
Zebra addresses these challenges by using a generative auto-regressive transformer that leverages in-context information during both training and when solving new problems. This means that Zebra can learn from previous examples and apply that knowledge without needing to adjust its parameters every time it encounters a new situation. It processes sequences of input data to adapt dynamically, allowing it to handle a wide range of scenarios effectively. The results show that Zebra performs better than existing methods in various PDE tasks.
Why it matters?
This research is significant because it offers a more flexible and efficient way to solve complex equations that are important in many fields, such as physics and engineering. By improving how models learn from data and adapt to new conditions, Zebra could help scientists and engineers tackle real-world problems more effectively, ultimately leading to advancements in technology and understanding of dynamic systems.
Abstract
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.