🤖 machine learning

Distillation of Foundation Models for Time-dependent PDEs

The paper proposes TREX, a knowledge distillation framework that generates synthetic long-horizon trajectories from fine-tuned foundation models to train compact, efficient student networks for time-dependent PDEs, achieving significant reductions in parameters and inference latency while maintaining or surpassing the original model's accuracy.

Original authors: Daniel Musekamp, Boshra Ariguib, Andrei Manolache, Mathias Niepert

Published 2026-08-13

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Daniel Musekamp, Boshra Ariguib, Andrei Manolache, Mathias Niepert

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the future of a chaotic system, like a swirling storm, a flowing river, or the way heat spreads through a metal plate. In the world of physics, these are described by equations called Partial Differential Equations (PDEs). Solving these equations is like trying to navigate a maze where the walls keep moving; it's incredibly hard and usually requires supercomputers to crunch the numbers step-by-step. Recently, scientists have built "Foundation Models"—massive, all-knowing AI brains trained on thousands of different physics problems. These giants can learn new tasks quickly, but they are so huge and slow that they can't be used for real-time things like controlling a robot arm or predicting weather in the next second. They are like a brilliant but sluggish librarian who needs to read every book in the library before answering a single question.

The big question is: Can we teach a tiny, fast, and nimble student AI to think just like the giant librarian, without needing the giant to do all the work every time? This is the challenge of "knowledge distillation." Usually, to teach a student, you need a lot of examples. But in the real world, we often only have a few snapshots of a system (like a few seconds of video from a sensor) and no idea what the starting conditions were for other scenarios. The paper you are about to read tackles this exact problem: how to take a slow, powerful AI teacher and compress its wisdom into a fast, small student, even when we don't have enough data to train the student from scratch.

The Story of TREX: Teaching a Tiny Robot to Think Like a Giant

Meet TREX (Teacher Rollout Extension). It's a clever new method invented by researchers Daniel Musekamp and his team to solve the "slow giant" problem. Imagine you have a master chef (the Foundation Model) who has cooked millions of dishes and knows exactly how flavors change over time. You want to teach a young apprentice (the Student Model) to cook the same way, but you only have three recipes and a very short time to learn. If you just let the apprentice practice on those three recipes, they might get stuck or make mistakes when the ingredients change slightly.

The problem is that the master chef is too slow to cook every single meal for the apprentice to watch. And you can't just ask the chef to "imagine" new starting ingredients because you don't know what ingredients the apprentice will actually need to handle in the real world.

Here is where TREX gets creative. Instead of waiting for the master chef to cook new meals from scratch, TREX lets the master chef take the few recipes you do have and "roll them out" into the future. Think of it like this: The master chef takes your one recipe for a cake and says, "Okay, let's bake it, then let's bake it again, and again, and again, for a hundred steps." This creates a long, imaginary timeline of how that cake evolves.

But here's the twist: sometimes, the real world gets messy. A wind blows, a spoon drops, or the temperature spikes. To prepare the apprentice for this, TREX occasionally throws a little bit of "noise" (like a gentle shake or a sprinkle of random spice) onto the cake while the master chef is rolling it out. The master chef then has to figure out how to fix the cake and continue cooking from that messy state. This teaches the apprentice not just how to cook a perfect cake, but how to recover when things go wrong.

The Magic Result:
By using this method, the researchers found that they could train a tiny student model that is 3,604 times smaller than the giant teacher. Yet, this tiny student can predict the future of the physics system just as accurately as the giant. In fact, in some tests, the student was even more accurate at long-term predictions because it learned to be more stable and didn't get as confused by small errors as the giant did.

Why does this matter?
The giant models are like a Ferrari that gets 2 miles per gallon—it's fast in a straight line but too expensive to run for everyday tasks. The student models trained with TREX are like a fuel-efficient electric scooter. They are:

Fast: They can make predictions more than 10 times faster than the giant.
Light: They use a fraction of the computer memory (about 3.2 times less).
Smart: They can be built with special "rules of physics" baked into their brains (like knowing that if you rotate a fluid, the physics should rotate with it), which the giant model didn't strictly follow.

What the researchers ruled out:
The team made sure to test if they really needed to know the "initial conditions" (the exact starting state of the universe) to make this work. They found that you do not need to know the full distribution of starting states. You can just start with the few data points you have and let the teacher "roll out" the rest. They also showed that simply copying the teacher's answers isn't enough; the "noisy rollouts" (the messy, shaken-up practice runs) are crucial for the student to learn how to handle real-world chaos.

How sure are they?
The researchers didn't just guess; they ran the numbers on several different physics problems, including fluid dynamics (like water flowing) and compressible gas (like air in a jet engine). They tested this on two different giant AI teachers (called Poseidon and Walrus) and found that the TREX method consistently produced students that matched or beat the teachers' accuracy. They even tested what happens if you remove the "noise" or the "ground truth" data, and the method still held up, though it worked best with both.

In short, TREX is a way to take the super-intelligence of a massive AI and distill it into a tiny, super-fast tool that can run on regular computers, making real-time physics simulations possible for things like weather forecasting, engineering design, and controlling robots. It's like taking a library of genius and shrinking it down to a single, super-smart pocket guide.

Technical Summary: Distillation of Foundation Models for Time-dependent PDEs

1. Problem Definition

Partial differential equations (PDEs) are fundamental to modeling spatiotemporal physical processes such as fluid dynamics. While neural networks have emerged as efficient surrogate models for PDEs, they face a critical trade-off:

Foundation Models (FMs): Large-scale models pretrained on diverse physics datasets (e.g., Poseidon, Walrus) demonstrate strong generalization to downstream tasks with limited data. However, they are computationally expensive, often relying on large transformer architectures that lack specific inductive biases (e.g., equivariance) and require significant inference time and memory. This limits their utility in real-time applications like control and planning.
Compact Surrogates: Smaller models (e.g., Fourier Neural Operators) are efficient but typically require large amounts of task-specific training data to achieve high accuracy, which is often unavailable in experimental or partially observed systems.

The core challenge is to transfer the generalization capabilities of large, pretrained PDE foundation models into compact, efficient student models without requiring access to the full initial-condition distribution of the target domain, which is often unknown or expensive to sample.

2. Methodology: Teacher Rollout Extension (TREX)

The authors propose Teacher Rollout Extension (TREX), a knowledge distillation framework designed to bridge the gap between large foundation models and efficient downstream surrogates.

2.1 Core Concept

TREX addresses the limitation of standard distillation, which often requires sampling new initial conditions from a known distribution $p_0$ . In many downstream scenarios (e.g., experimental data), $p_0$ is unknown. Instead, TREX augments the limited available ground-truth trajectories by generating long synthetic trajectories through teacher rollouts.

2.2 The TREX Procedure

Fine-tuning: A pretrained foundation model (Teacher, $T$ ) is first fine-tuned on a small set of available downstream trajectories.
Rollout Generation: Starting from the initial states of the few available ground-truth trajectories, the fine-tuned teacher is rolled out autoregressively to generate long synthetic trajectories.
- Occupancy Measure: This process samples from the teacher-induced rollout distribution (states encountered during autoregressive prediction) rather than just the short observed trajectories.
- Noise Injection: To broaden state-space coverage and teach the student how the system evolves from perturbed states, periodic Gaussian noise is injected into the teacher's state during rollouts. The teacher then predicts the next state based on this perturbed state.
Student Training: A compact student model ( $S$ $S$ ) is trained on a mixture of:
- Original ground-truth transitions.
- Teacher-generated transitions (both noise-free and noisy segments).
- The student loss combines ground-truth supervision ( $L_{GT}$ ) and teacher distillation loss ( $L_{TREX}$ ).

2.3 Architectural Inductive Biases

Unlike the teacher, which is typically a large, flexible transformer, the student architecture is chosen to incorporate task-specific inductive biases. The authors utilize a Tensorized Fourier Neural Operator (TFNO), which enforces equivariance to translations and suppresses high-frequency noise, properties not explicitly enforced by the teacher.

3. Key Contributions

First Systematic Study: The paper presents the first systematic investigation into distilling general pretrained PDE foundation models into compact downstream surrogate models.
TREX Framework: Introduction of a novel distillation method that generates supervision data via noisy teacher rollouts, eliminating the need for an explicit initial-condition sampler.
Inductive Bias Integration: Demonstration that distilling a general foundation model into a small architecture with suitable inductive biases (e.g., equivariance) can respect physical properties while maintaining the teacher's predictive accuracy.

4. Experimental Results

The authors evaluated TREX using Poseidon and Walrus as teachers across various fluid dynamics benchmarks (Navier-Stokes and Compressible Euler equations).

Accuracy: TREX-trained students matched or surpassed the accuracy of the large foundation model teachers across low-data regimes (4 to 32 trajectories). In some cases, the student outperformed the teacher even when trained solely on teacher-generated data (without ground truth).
Efficiency:
- Parameters: The student models reduced the parameter count by several orders of magnitude (e.g., from ~628M parameters in Poseidon-L to ~174K in TFNO, a reduction of ~3600x).
- Inference Speed: Inference time was reduced by more than an order of magnitude (e.g., from ~270ms to ~22ms per step).
- Memory: GPU memory consumption for a forward pass was reduced by a factor of ~3.2.
Stability and Equivariance:
- Students exhibited more stable long-horizon rollouts with slower error growth compared to teachers.
- The TFNO student demonstrated significantly lower translation equivariance errors compared to the non-equivariant teacher, maintaining consistency under spatial shifts during autoregressive rollouts.
Ablations:
- The method is robust to rollout length and noise hyperparameters.
- Noise injection during rollouts was found to be beneficial, particularly at higher amplitudes.
- Distillation was shown to be effective even when the student architecture (UNet, standard FNO) could not outperform the teacher using only ground truth data.

5. Significance and Claims

The paper claims that TREX effectively reconciles the trade-off between the strong generalization of large foundation models and the efficiency requirements of practical surrogate modeling. By extending limited downstream data through teacher rollouts, TREX enables compact models to learn long-horizon dynamics and physical symmetries that would otherwise require massive datasets or complex architectures.

The authors position this work as a step toward making foundation-model-based surrogates deployable in real-time scientific and engineering applications (e.g., manufacturing control, energy systems) where computational cost and inference speed are critical constraints. They note that while the method relies on the quality of the fine-tuned teacher, it successfully transfers knowledge even when the teacher lacks specific inductive biases, allowing the student to enforce physical constraints like equivariance.