Distillation of Foundation Models for Time-dependent PDEs

2026-08-12

AI Summary

While large foundation models demonstrate strong performance in solving time-dependent partial differential equations, their high computational cost limits their practicality as replacements for efficient numerical solvers. This work proposes the Teacher Rollout Extension (TREX) framework, which leverages knowledge distillation to transfer capabilities from a pretrained teacher model to a lightweight student model. By using long-horizon synthetic trajectories generated by the teacher to augment limited downstream data, TREX enables sampling of rollout trajectories without requiring prior knowledge of the initial condition distribution. This exposes the student model to both long-term dynamics and local recovery behaviors, while allowing integration of task-specific inductive biases—such as equivariance. Combined with noise injection and an equivariant network architecture, the resulting student model achieves several orders of magnitude fewer parameters, over tenfold faster inference, and accuracy comparable to or exceeding that of the teacher.

Abstract

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.

Problem

• foundation models
• time-dependent PDEs
• knowledge distillation
• computational efficiency
• surrogate modeling

Innovation

• knowledge distillation
• foundation models
• time-dependent PDEs
• teacher-student framework
• trajectory synthesis