Kastor: An efficient fine-tuning strategy for generative emulation of PDE simulations
Guillaume Couairon, Alexis Jacq, Yu-Han Wu, Renu Singh, Yana Hasson, Quentin Berthet, Romuald Elie
cs.LG, math.NA
Submitted: 2026-08-06
Comments: 34 pages, 32 figures
License: http://creativecommons.org/licenses/by/4.0/
The gist: Machine learning offers a promising avenue to accelerate physical simulations by replacing computationally expensive traditional Partial Differential Equation (PDE) solvers with fast, differentiable
Terminology
Abstract
Machine learning offers a promising avenue to accelerate physical simulations by replacing computationally expensive traditional Partial Differential Equation (PDE) solvers with fast, differentiable surrogate models. However, standard auto-regressive ML emulators often suffer from error accumulation over long horizons and struggle to capture the stochasticity of complex physical systems. In this paper, we propose Kastor, a comprehensive methodology to adapt a deterministic physics foundation model into a highly efficient and accurate generative surrogate. First, we introduce a two-stage inference scheme that combines a large-stride causal auto-regressive model with a non-causal temporal super-resolution network, significantly reducing error accumulation while minimizing computational cost. Second, we present Mean prediction regularization (MPR), a novel training objective that constrains the generative model to predict the deterministic distribution mean under null noise conditioning. This regularization dramatically improves the performance and stability of both Functional Generative Networks (FGN) and diffusion-based emulators. Finally, we demonstrate that incorporating spatial gradient matching improves the accuracy and physical fidelity of the simulations as measured by power spectrum density. Extensive evaluations on diverse simulation datasets of the benchmark The Well show that with these components, our model outperforms competing methods in forecasting accuracy, spectral consistency, and computational efficiency. Our model achieves a 42.9% average reduction in forecasting compared to our reference based on the Walrus finetuning methodology, and outperforms Walrus for 8 out of 10 datasets on variance-normalized RMSE (VRMSE).
Sources
- Skillful joint probabilistic weather forecasting from marginals
- Fluid Intelligence: A Forward Look on AI Foundation Models in Computational Fluid Dynamics
- U-Cast: A Surprisingly Simple and Efficient Frontier Probabilistic AI Weather Forecaster
- Diffusion Fine-tuning with Rewarded Moment Matching Distillation
- Adam: A Method for Stochastic Optimization
- PROSE-FD: A Multimodal PDE Foundation Model for Learning Multiple Operators for Forecasting Fluid Dynamics
- BCAT: A Block Causal Transformer for PDE Foundation Models for Fluid Dynamics
- Walrus: A Cross-Domain Foundation Model for Continuum Dynamics
- PhysiX: A Foundation Model for Physics Simulations
- MORPH: PDE Foundation Models with Arbitrary Data Modality
- Test-time Generalization for Physics through Neural Operator Splitting
- Towards a Foundation Model for Partial Differential Equations Across Physics Domains
- Towards a Physics Foundation Model
- PDEformer-2: A Versatile Foundation Model for Two-Dimensional Partial Differential Equations
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks