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

초록

시간 의존적 매개변수 편미분 방정식(PDEs)을 해결하는 것은 도전적입니다. 모델은 계수, 강제항, 경계 조건과 같은 매개변수의 변화에 적응해야 합니다. 데이터 주도형 신경 해법은 PDE 매개변수 분포에서 샘플링된 데이터로 학습하거나 그래디언트 기반 적응 및 메타-러닝을 활용하여 관측에서 역동학을 암시적으로 인코딩할 수 있습니다. 이는 종종 추론 복잡성의 증가와 함께 발생합니다. 대형 언어 모델(LLMs)의 맥락 학습 능력에서 영감을 받아, 우리는 그래디언트 적응이 추론에서 필요하지 않은 매개변수 PDE를 해결하기 위해 설계된 새로운 생성적 자기 회귀 트랜스포머인 Zebra를 소개합니다. 사전 훈련 및 추론 중에 맥락 정보를 활용함으로써, Zebra는 입력 시퀀스에 의존하여 새로운 작업에 동적으로 적응하며, 맥락 궤적 또는 이전 상태를 통합합니다. 이 접근 방식은 Zebra가 임의 크기의 맥락 입력을 유연하게 처리하고 여러 해결 궤적을 샘플링하여 불확실성 양자화를 지원합니다. 우리는 다양한 어려운 PDE 시나리오를 통해 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.