WalrusFlow: Probabilistic Fine-tuning of a Physics Foundation Model via Flow Matching
Abstract
Physics foundation models (FMs) have demonstrated strong forecasting capabilities for complex dynamical systems. However, most physics FMs are trained as point estimators of future states and therefore remain susceptible to degraded spectral fidelity, rapid error accumulation during long-horizon rollouts, and an inability to represent predictive uncertainty. Meanwhile, generative models such as flow matching can retain high-frequency features and naturally represent predictive uncertainty through ensembles, but are usually trained separated from pretrained point-estimator FMs. In this work, we augment a pretrained FM for continuum dynamics with a lightweight conditional flow-based refinement head. The flow head learns structured corrections to the frozen base-model forecasts and, at inference time, generates ensembles of refined solutions. Demonstration studies on two fluid dynamics datasets show that these corrections improve pointwise forecast accuracy, spectral fidelity, probabilistic forecast quality, and long-horizon rollout performance. These results demonstrate that a conditional flow head offers a fine-tuning scheme that can effectively learn forecast corrections and add probabilistic forecasting capabilities to a pretrained deterministic physics FM without retraining the base model.