# Distillation of Foundation Models for Time-dependent PDEs

arXiv: https://arxiv.org/abs/2608.11937 · HTML: https://mlsanity.com/paper/2608.11937/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 35 of its 48 references that are in the map.
- First posted on arXiv 12 Aug 2026.

## Builds on (most cited first, 30 of 35)

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [A ConvNet for the 2020s](https://mlsanity.com/paper/2201.03545/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [FitNets: Hints for Thin Deep Nets](https://mlsanity.com/paper/1412.6550/) (2014)
- [Knowledge Distillation: A Survey](https://mlsanity.com/paper/2006.05525/) (2020)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [Relational Knowledge Distillation](https://mlsanity.com/paper/1904.05068/) (2019)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [Physics-Informed Neural Operator for Learning Partial Differential Equations](https://mlsanity.com/paper/2111.03794/) (2021)
- [Be Your Own Teacher: Improve the Performance of Convolutional Neural Networks via Self Distillation](https://mlsanity.com/paper/1905.08094/) (2019)
- [Do Deep Nets Really Need to be Deep?](https://mlsanity.com/paper/1312.6184/) (2013)
- [Graph networks as learnable physics engines for inference and control](https://mlsanity.com/paper/1806.01242/) (2018)
- [Poseidon: Efficient Foundation Models for PDEs](https://mlsanity.com/paper/2405.19101/) (2024)
- [Towards Multi-spatiotemporal-scale Generalized PDE Modeling](https://mlsanity.com/paper/2209.15616/) (2022)
- [Self-Distillation Amplifies Regularization in Hilbert Space](https://mlsanity.com/paper/2002.05715/) (2020)
- [DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training](https://mlsanity.com/paper/2403.03542/) (2024)
- [The Well: a Large-Scale Collection of Diverse Physics Simulations for Machine Learning](https://mlsanity.com/paper/2412.00568/) (2024)
- [Data-Free Knowledge Distillation for Deep Neural Networks](https://mlsanity.com/paper/1710.07535/) (2017)
- [Knowledge Transfer with Jacobian Matching](https://mlsanity.com/paper/1803.00443/) (2018)
- [Multi-Grid Tensorized Fourier Neural Operator for High-Resolution PDEs](https://mlsanity.com/paper/2310.00120/) (2023)
- [Frequency-Aligned Knowledge Distillation for Lightweight Spatiotemporal Forecasting](https://mlsanity.com/paper/2507.02939/) (2025)
- [A Framework and Benchmark for Deep Batch Active Learning for Regression](https://mlsanity.com/paper/2203.09410/) (2022)
- [Group Equivariant Fourier Neural Operators for Partial Differential Equations](https://mlsanity.com/paper/2306.05697/) (2023)
- [Learning Neural PDE Solvers with Parameter-Guided Channel Attention](https://mlsanity.com/paper/2304.14118/) (2023)
- [A Library for Learning Neural Operators](https://mlsanity.com/paper/2412.10354/) (2024)
- [Walrus: A Cross-Domain Foundation Model for Continuum Dynamics](https://mlsanity.com/paper/2511.15684/) (2025)
- [Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers](https://mlsanity.com/paper/2405.17527/) (2024)
- [PROSE-FD: A Multimodal PDE Foundation Model for Learning Multiple Operators for Forecasting Fluid Dynamics](https://mlsanity.com/paper/2409.09811/) (2024)

## Abstract (arXiv)

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.