Distillation of Foundation Models for Time-dependent PDEs

Affiliation: University of Stuttgart, University of Cologne, Bitdefender

arXiv: 2608.11937v1 [cs.LG] 12 Aug 2026

Foundation models for time-dependent partial differential equations (PDEs) are trained on large and diverse collections of physical systems and can generalize effectively to new downstream tasks. After fine-tuning on only a few trajectories from a target domain, they can achieve strong accuracy in low-data regimes. However, these models are typically large and computationally intensive, limiting their usefulness as fast surrogates for numerical solvers.

The paper states: "Foundation models for PDEs have been introduced... By pretraining a surrogate on a large set of diverse physics, foundation models can transfer knowledge to downstream tasks with fewer training samples. On the downside, their inference is computationally expensive, as they require more parameters to generalize well. Slow inference prohibits their use in applications that have real-time requirements, such as planning and control tasks in manufacturing or energy systems."

Additionally, the authors note that "the used architectures are mostly large transformers that lack inductive biases, since these may vary from problem to problem. For example, many PDEs are locally equivariant to rotations and translations, and globally as well if the boundary conditions are periodic."

The paper proposes TREX, a knowledge distillation framework that transfers the predictive capability of a pretrained foundation model into a compact and efficient student. The approach works as follows:

1. Fine-tuning: The foundation model (teacher) is first fine-tuned on a small set of trajectories in the target domain.

2. Teacher Rollout Extension: Starting from the few available downstream trajectories, TREX generates long teacher rollouts and trains the student on the resulting teacher-labeled trajectories.

3. Noise Injection: Periodic noise injection can further broaden local state-space coverage and teach the student how the teacher evolves from perturbed states.

4. Student Training: The student is trained on a mixture of ground-truth and teacher-labeled transitions.

The key insight is that TREX does not require access to the true initial-condition distribution. Instead, it reuses the few ground-truth trajectories as anchors and expands them through the dynamics learned by the foundation model.

The paper formalizes this through an occupancy measure: TREX constructs an extended teacher-induced state occupancy measure ρKT = (1/K) Σ t=0 K-1 (T t)# ν0, where (T t)# ν0 denotes the pushforward of ν0 under t autoregressive teacher steps and K is the rollout horizon. This describes the state distribution induced by rolling out the teacher from the selected set of ground-truth samples.

The noise injection procedure is defined as: z̃t = zt + mt ξt, zt+1 = T(z̃t), where mt = 1 t > 0 ∧ t ≡ 0 mod T Noise, and (ξt)c(x) N(0, σ2s2c). Here, sc denotes the global standard deviation of channel c computed over the dataset, and σ is a fixed scaling factor.

The final student objective combines ground-truth supervision with the teacher rollout distillation: LS = LGT + λTREX LTREX.

The paper provides the following contributions:

1. First systematic study of distilling general pretrained PDE foundation models into compact downstream surrogate models.

2. TREX, a novel distillation method that generates supervision data using noisy teacher rollouts, in some cases surpassing the teacher's accuracy.

3. Demonstration that by distilling the foundation model into a small, lightweight student architecture with suitable inductive biases, inference time can be reduced by more than an order of magnitude while respecting physical properties such as equivariances.

• Poseidon-L (Herde et al., 2024): A SwinV2 vision transformer with 628.6M parameters, pretrained on six different datasets containing Navier-Stokes and compressible Euler equations.

• Walrus (McCabe et al., 2025): A larger vision transformer with 1.3B parameters that uses the last six states as context for next state prediction.

• TFNO (Tucker Factorized Fourier Neural Operator): Uses Tucker factorization to reduce parameters. The network uses 32 modes per dimension, 32 hidden channels, three layers, and a GELU activation function, resulting in only 174.4K parameters.

• U-Net (Gupta & Brandstetter, 2023): 9,161,332 parameters.

• FNO (Li et al., 2021): 5,042,836 parameters.

• CE-RPUI: Compressible Euler with Random Piecewise Uniform Initial conditions

• NS-SVS: Incompressible Navier-Stokes with Sinusoidal Vortex Sheet initial conditions

• NS-PwC: Navier-Stokes with Piecewise Constant initial conditions

• NS-BB: Navier-Stokes with Brownian Bridge initial conditions

• Kolmogorov Flow (for Walrus experiments): Reynolds number 5000, 501 time steps

The paper reports: TREX matches or surpasses its teacher network across the selected low-data settings. The results show median relative L1 errors for various numbers of training trajectories (N = 4, 8, 16, 32).

Key findings:

• TREX matches or surpasses the teacher's accuracy across all selected downstream settings

• IC-KD (which requires full knowledge of the initial condition distribution) also matches the teacher but requires access to the IC distribution, which is not the case for experimental settings

• Relational KD slightly improves the student but cannot match the teacher in settings where the teacher outperformed ground truth learning

The paper reports: While surpassing the teacher's accuracy, the student needs 3604 times fewer parameters. Specifically:

Model Parameters 1-step Time Full Rollout Time

Poseidon-L (lead-time) 628.6M 270.9 ms 1908.2 ms

Poseidon-L (autoreg.) 628.6M 269.4 ms 1887.7 ms

TFNO (student) 174.4K 22.1 ms 152.1 ms

The inference time is reduced by over an order of magnitude, and the memory consumption for a forward pass by a factor of 3.2.

The paper notes: While the student has a larger error at the beginning, it has a more stable long-term rollout and has a slower error growth. This behavior is observed across multiple datasets.

The paper conducts an ablation on NS-PwC (four trajectories):

Variant Median rel. L1

TREX 0.056 ± 0.003

w/o GT 0.060 ± 0.002

w/o GT & w/o Noise 0.064 ± 0.006

Teacher 0.070 ± 0.011

The paper states: Even without the ground truth data, TREX still outperforms its teacher. Removing noise injections during the rollout reduces the gap, but not fully.

• Rollout Length: The algorithm is quite robust to the choice of this hyperparameter. The performance only decreases slowly for very long rollouts.

• Noise Strength: Choosing the noise quite high at a level of one time the channel-wise standard deviation sc produces the best results.

The paper tests other student architectures (UNet and FNO): "TREX is also able to let other student architectures match the teacher's accuracy for all considered numbers of trajectories. While the TFNO was able to improve over the teacher using only ground truth for a higher number of trajectories in the main experiments, this was not the case for the UNet or even a standard FNO, and was only possible through distillation."

For the Walrus experiments on Kolmogorov Flow: The student trained with TREX can also successfully distill the multi-step teacher foundation model. TREX improves long-rollout accuracy while reducing inference time by more than an order of magnitude.

The paper performs a self-distillation experiment using the TFNO trained only on ground-truth samples as the teacher: While self-distillation has a small positive regularizing effect, the main advantage comes from distilling the pretrained, higher-capacity foundation model.

The paper investigates the effect of explicit equivariance built into the student architecture. The TFNO student is equivariant to translations on the periodic grid, while the teacher (Poseidon) is not.

The results show:

• The student has substantially lower equivariance error than the teacher across all datasets

• The shifted rollout error of the student remains essentially unchanged, while the teacher can degrade substantially, most clearly on NS-SVS

The paper concludes: "These results suggest that equivariance can be a useful property of the student, especially in autoregressive simulations where small symmetry errors may accumulate over long rollouts. More broadly, it is encouraging that TREX can use a non-equivariant teacher to train a compact student that preserves the teacher's predictive behavior while enforcing the symmetry of the downstream task."

The paper acknowledges several limitations:

1. "TREX relies on the quality of the fine-tuned teacher. Since the method samples from the teacher-induced rollout distribution, any systematic bias or instability in the teacher's long-horizon dynamics may be transferred to the student."

2. "Rollout extension is most useful when long trajectories continue to explore informative regions of state space. Its benefit may therefore be limited for systems that rapidly collapse to a trivial fixed point, exhibit strongly transient dynamics, or contain multiple disconnected long-time regimes that are not represented by the available trajectories."

3. "The Gaussian perturbations used by TREX are not guaranteed to satisfy all physical constraints of the underlying PDE and should therefore be viewed as a state-space augmentation strategy rather than as physically valid perturbations."

The paper concludes: "We introduce TREX, a simple and effective knowledge distillation framework for PDE foundation models. By extending limited downstream data through long and noisy teacher rollouts, TREX enables the student to match the teacher's long-horizon dynamics, thereby providing valuable additional information. Our results demonstrate that distillation in the foundation model regime is not only feasible but highly effective: compact student models can match the accuracy of large teachers while achieving substantial improvements in inference speed and memory efficiency. Additionally, the student can incorporate inductive biases to faithfully reproduce prior knowledge about the system, such as equivariances. Overall, TREX bridges the gap between the strong generalization capabilities of large foundation models and the efficiency requirements of practical surrogate modeling."

Improvements to AI systems:

1. Distillation framework for large pretrained models: Implement TREX to compress any large pretrained autoregressive model (not just PDE foundation models) into a compact student network. The student can match or exceed teacher accuracy while using 3,600x fewer parameters and achieving >10x faster inference.

2. Noise-injected rollout data augmentation: Use the teacher's own rollouts with periodic Gaussian perturbations as synthetic training data. This expands limited ground-truth trajectories into dense coverage of the state space, enabling the student to learn stable long-horizon dynamics without access to the true initial-condition distribution.

3. Inductive bias injection via student architecture: Choose a student architecture with built-in equivariances (e.g., translation equivariance for periodic domains) that the teacher lacks. The distillation process transfers predictive behavior while enforcing physical symmetries, reducing accumulated symmetry errors in autoregressive rollouts.

4. Hybrid ground-truth + teacher-labeled training objective: Train the student on a weighted mixture of real data and teacher-generated transitions. This balances fidelity to actual physics with the teacher's learned dynamics, improving robustness in low-data regimes (4–32 trajectories).

5. Self-distillation as regularization: Apply the same rollout-extension technique using a student trained only on ground truth as the teacher. This provides a mild regularizing effect, useful when no large pretrained model is available.

What the improved AI system can do:

• Real-time PDE surrogate modeling: Deploy compact models (e.g., 174K parameters) for applications like manufacturing control or energy system planning, where large foundation models (628M–1.3B parameters) are too slow for real-time inference.

• Stable long-horizon prediction: Generate accurate multi-step rollouts with slower error growth than the teacher, critical for weather forecasting, fluid dynamics simulation, or climate modeling.

• Symmetry-preserving simulation: Enforce known physical invariances (rotations, translations, periodic boundaries) even when the original teacher is not equivariant, reducing drift in long simulations.

• Data-efficient adaptation: Fine-tune on as few as 4–8 observed trajectories and still achieve teacher-level accuracy, enabling deployment in experimental settings where collecting data is expensive or impossible.

• Multi-teacher distillation: Distill from different foundation models (e.g., Poseidon-L or Walrus) into a single efficient student, combining strengths of various architectures.

• Transferable framework: Apply the same rollout-extension + noise injection + inductive-bias student approach to other sequence-prediction domains (e.g., video prediction, molecular dynamics, or traffic flow) where large pretrained models exist but real-time deployment is needed.

Foundation models for time-dependent partial differential equations (PDEs) are trained on large and diverse collections of physical systems and can generalize effectively to new downstream tasks. After fine-tuning on only a few trajectories from a target domain, they can achieve strong accuracy in low-data regimes. However, these models are typically large and computationally intensive, limiting their usefulness as fast surrogates for numerical solvers. We propose Teacher Rollout Extension (TREX), a knowledge distillation framework that transfers the predictive capability of a pretrained foundation model into a compact and efficient student. Starting from a fine-tuned teacher, TREX augments limited downstream data by generating long synthetic trajectories through teacher rollouts, optionally with periodic noise injection. This procedure samples from the teacher-induced rollout distribution without requiring explicit knowledge of the initial-condition distribution, while exposing the student to long-horizon states and local recovery behavior around states encountered during autoregressive prediction. The student can further incorporate task-specific inductive biases, such as equivariance, that the teacher does not necessarily enforce. We evaluate TREX on multiple PDE benchmarks. The resulting students can match or surpass the teacher's accuracy while reducing the number of parameters by several orders of magnitude and achieving more than an order-of-magnitude speedup in inference.