Zebra: контекстное и генеративное предварительное обучение для решения параметрических уравнений в частных производных

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

October 4, 2024

Авторы: Louis Serrano, Armand Kassaï Koupaï, Thomas X Wang, Pierre Erbacher, Patrick Gallinari

cs.AI

Аннотация

Решение временно-зависимых параметрических уравнений в частных производных (УЧП) представляет собой сложную задачу, поскольку модели должны адаптироваться к изменениям параметров, таким как коэффициенты, воздействующие члены и граничные условия. Нейросетевые решатели, основанные на данных, лиелятся на обучение на данных, отобранных из распределения параметров УЧП в надежде, что модель обобщит на новые случаи, или полагаются на градиентную адаптацию и метаобучение для неявного кодирования динамики из наблюдений. Это часто сопровождается увеличением сложности вывода. Вдохновленные возможностями обучения в контексте больших языковых моделей (БЯМ), мы представляем Zebra, новый генеративный авторегрессионный трансформер, разработанный для решения параметрических УЧП без необходимости градиентной адаптации при выводе. Используя информацию в контексте как на этапе предварительного обучения, так и при выводе, Zebra динамически адаптируется к новым задачам, условиям на входных последовательностях, которые включают траектории контекста или предшествующие состояния. Этот подход позволяет Zebra гибко обрабатывать входы произвольного размера и поддерживать оценку неопределенности путем выборки нескольких траекторий решения. Мы оцениваем Zebra в различных сложных сценариях УЧП, демонстрируя его адаптивность, надежность и превосходную производительность по сравнению с существующими подходами.

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.