일반화 가능한 장기 물리 시뮬레이션을 위한 잠재 생성 기반 솔버

Latent Generative Solvers for Generalizable Long-Term Physics Simulation

본 연구에서는 이질적인 PDE 시스템에 대한 장기 예측 시뮬레이션을 다룹니다. 우리는 Latent Generative Solvers (LGS)라는 두 단계 프레임워크를 소개합니다. LGS는 (i) 사전 학습된 VAE를 사용하여 다양한 PDE 상태를 공유된 잠재 물리 공간으로 매핑하고, (ii) 플로우 매칭을 통해 학습된 트랜스포머를 사용하여 확률적 잠재 역학을 학습합니다. 핵심 메커니즘은 불확실성 조절 장치(uncertainty knob)이며, 이는 학습 및 추론 과정에서 잠재 입력을 변경하여 솔버가 예측 과정에서 발생하는 오차를 수정하고 자기 회귀 예측을 안정화하도록 합니다. 또한, 모델이 생성한 경로를 사용하여 시스템 설명자(context)를 업데이트하는 플로우 강제(flow forcing) 방법을 사용하여 학습/테스트 조건의 일관성을 높이고 장기적인 안정성을 향상시킵니다. 우리는 12가지 PDE 패밀리에 걸쳐 약 250만 개의 경로를 포함하는 데이터셋을 사용하여 LGS를 사전 학습했습니다. LGS는 짧은 시간 범위에서는 기존의 강력한 결정론적 신경 연산자 모델과 유사한 성능을 보이지만, 장기적인 시간 범위에서는 예측 오차를 크게 줄입니다. 잠재 공간에서의 학습과 효율적인 아키텍처 설계는 기존의 생성 모델 기반 방법보다 최대 70배 낮은 FLOPs를 달성하여 확장 가능한 사전 학습을 가능하게 합니다. 또한, 제한된 미세 조정(finetuning) 예산 하에서 256x256 해상도의 Kolmogorov flow 데이터셋에 효율적으로 적용할 수 있음을 보여줍니다. 전반적으로, LGS는 일반화 가능하고 불확실성을 고려하는 신경 PDE 솔버를 개발하는 실용적인 방법을 제공하며, 이는 장기 예측 및 과학적 워크플로우에서 더 안정적인 성능을 제공합니다.

We study long-horizon surrogate simulation across heterogeneous PDE systems. We introduce Latent Generative Solvers (LGS), a two-stage framework that (i) maps diverse PDE states into a shared latent physics space with a pretrained VAE, and (ii) learns probabilistic latent dynamics with a Transformer trained by flow matching. Our key mechanism is an uncertainty knob that perturbs latent inputs during training and inference, teaching the solver to correct off-manifold rollout drift and stabilizing autoregressive prediction. We further use flow forcing to update a system descriptor (context) from model-generated trajectories, aligning train/test conditioning and improving long-term stability. We pretrain on a curated corpus of $\sim$2.5M trajectories at $128^2$ resolution spanning 12 PDE families. LGS matches strong deterministic neural-operator baselines on short horizons while substantially reducing rollout drift on long horizons. Learning in latent space plus efficient architectural choices yields up to \textbf{70$\times$} lower FLOPs than non-generative baselines, enabling scalable pretraining. We also show efficient adaptation to an out-of-distribution $256^2$ Kolmogorov flow dataset under limited finetuning budgets. Overall, LGS provides a practical route toward generalizable, uncertainty-aware neural PDE solvers that are more reliable for long-term forecasting and downstream scientific workflows.