Technical Summary: Kastor – An Efficient Fine-Tuning Strategy for Generative Emulation of PDE Simulations

Problem Statement

Physical simulations based on Partial Differential Equations (PDEs) are fundamental to modern science and engineering but suffer from massive computational costs due to the need for fine spatial and temporal meshes. While Machine Learning (ML) offers a path to accelerate these simulations by learning continuous mappings as surrogate models, existing approaches face significant hurdles. Standard auto-regressive ML emulators often accumulate errors over long time horizons, leading to divergence. Furthermore, they struggle to capture the inherent stochasticity of complex physical systems and often fail to reproduce correct statistical distributions, such as power spectrum densities. The challenge lies in developing surrogate models that are not only fast and accurate but also stable, physically plausible, and capable of generating probabilistic trajectories that reflect system uncertainty.

Methodology

The authors propose Kastor, a comprehensive methodology to adapt the Walrus physics foundation model—a deterministic, pre-trained transformer architecture—into a highly efficient and accurate generative surrogate. The approach integrates three primary technical innovations:

1. Two-Stage Inference Scheme

To mitigate error accumulation while maintaining computational efficiency, Kastor employs a two-stage inference strategy:

• Causal Strided Model: A primary causal model predicts states at a large time stride ($\tau$, typically $\tau=4$). This reduces the number of auto-regressive steps required to reach a specific lead time, thereby limiting the compounding of errors.
• Non-Causal Temporal Super-Resolution: A lightweight, non-causal up-sampling network fills in the missing intermediate states between the strided predictions. Unlike block-wise auto-regressive approaches that split model capacity across multiple future states, this up-sampler uses all available context (past and future strided states) to resolve missing frames. This design significantly reduces computational cost (by a factor of $\sim37.5\%$ at long lead times) while preserving trajectory density.

2. Mean Prediction Regularization (MPR)

Recognizing that PDE emulation often involves a dominant deterministic mode with smaller stochastic components, the authors introduce Mean Prediction Regularization (MPR). This objective constrains the generative model to predict the mean of the target distribution when conditioned on null noise.

• Implementation: During training, the model performs a deterministic pass (with noise set to zero) alongside the stochastic pass. The loss for this deterministic pass is a standard regression objective (e.g., MAE) against the ground truth.
• Effect: MPR stabilizes the training of Functional Generative Networks (FGN) and diffusion-based emulators, ensuring the model captures the primary deterministic dynamics before learning the stochastic residuals. This prevents the generative component from overwhelming the physical signal.

3. Gradient Difference Loss (GDL)

Standard training losses like Mean Absolute Error (MAE) or Continuous Ranked Probability Score (CRPS) may fail to preserve critical spatial derivatives essential for governing physical dynamics. Kastor incorporates Gradient Difference Loss (GDL), which explicitly matches the spatial gradients between the forecast and ground truth states. This ensures the preservation of sharp features (e.g., shock fronts) and improves the spectral consistency of the simulations.

Training and Architecture Details

• Backbone: The method leverages the Walrus architecture (convolutional encoder, causal transformer with spatial/time attention, convolutional decoder).
• Stochasticity Injection: Stochasticity is introduced via Adaptive Layer-Normalization (AdaLN) conditioned on low-dimensional Gaussian noise, combined with patch jittering (random spatial shifts of input data) to prevent grid-imprinting artifacts and enhance ensemble diversity.
• Loss Function: The training objective combines CRPS (for generative fidelity), MPR (for deterministic stability), and GDL (for physical fidelity).

Key Results

The authors evaluated Kastor on 10 diverse 2D datasets from The Well benchmark, covering domains such as fluid dynamics, acoustic scattering, and reaction-diffusion systems.

• Forecasting Accuracy: Kastor achieves an average 42.9% reduction in forecasting error (measured by fCRPS skill score) compared to a reference model fine-tuned using the standard Walrus methodology.
• Variance-Normalized RMSE (VRMSE): The retrained Kastor model outperforms the reported numbers from the original Walrus paper on 8 out of 10 datasets when evaluated using variance-normalized RMSE.
• Spectral Consistency: The inclusion of GDL significantly improves the Log Spectral Distance (LSD), ensuring the power spectrum density of the generated trajectories matches the ground truth more closely than deterministic baselines.
• Efficiency: The two-stage inference scheme reduces computational costs by approximately 37.5% compared to standard auto-regressive rollouts while maintaining or improving accuracy, particularly at longer lead times where error accumulation is most severe.
• Ablation Studies:
  • Time Stride: A stride of $\tau=4$ was found optimal for most datasets, though $\tau=1$ remained superior for highly sub-sampled simulations (e.g., Active Matter, Rayleigh-Bénard).
  • Regularization: MPR was critical for the performance of FGN models, which otherwise underperformed compared to diffusion baselines.
  • Gradient Loss: GDL provided marginal to significant improvements across datasets, particularly for those with sharp discontinuities like the Acoustic Scattering Maze.

Significance and Claims

The paper claims that Kastor represents a robust framework for deriving strong generative surrogate models from pre-trained deterministic foundation models. Its significance lies in demonstrating that:

1. Deterministic-to-Generative Adaptation: Converting a deterministic foundation model into a generative one yields superior predictive capabilities, provided the training is appropriately constrained (via MPR).
2. Efficiency-Accuracy Trade-off: Shifting from standard single-step auto-regressive rollouts to a two-stage inference scheme (strided prediction + non-causal super-resolution) successfully balances computational efficiency with high-fidelity trajectory generation.
3. Physical Fidelity: By combining generative modeling with spatial gradient matching and mean regularization, the approach successfully preserves the spectral distribution and physical stability of complex systems, outperforming both standard deterministic baselines and more computationally heavy diffusion models.

The authors conclude that their methodology offers a practical path toward scalable, probabilistic PDE emulation that meets the rigorous demands of engineering and scientific simulation.