Research spotlight · NeurIPS 2026

Distillation of Foundation Models for Time-dependent PDEs

Let a large model teach a small one to predict evolving physics.

Daniel Musekamp · Boshra Ariguib · Andrei Manolache · Mathias Niepert

A large, slower teacher neural network extends a few available states into a trajectory of predicted states. One short amber-and-blue segment of a faint attractor is connected by dashed zoom guides to the enlarged trajectory below. Each dot represents a complete system state. A smaller, faster student network learns from these examples and runs independently after training.

The method in a nutshell

A large foundation model can learn to predict evolving physical systems from just a few simulations. But running that model can still be expensive.

Teacher Rollout Extension (TREX) turns it into a teacher. After fine-tuning, the teacher repeatedly predicts later states, extending the available simulations into longer training examples. Optional noise exposes the student to perturbed states and the teacher’s recovery.

A compact student learns from this richer data and can incorporate physical symmetries. TREX needs no sampler for new initial conditions: it starts with the examples already available.

The payoff: in the paper’s PDE benchmarks, students can match or surpass the teacher’s accuracy with over 10× faster inference and orders of magnitude fewer parameters.

Read the paper

Distillation of Foundation Models for Time-dependent PDEs. NeurIPS 2026. Publication record