REVIEW 3 major objections 5 minor 21 references
Distillation of Foundation Models for Time-dependent PDEs
T0 review · 3 major / 5 minor · reviewed 2026-08-16 · deepseek-v4-flash
Pith's one-line read Large pretrained PDE foundation models can be distilled into compact surrogates that match or beat the teacher's accuracy with about 3,604 times fewer parameters and an order-of-magnitude faster inference, using only long teacher rollouts…
desk verdict TREX is a novel distillation method for PDE foundation models, but the 'student beats teacher' claim needs a budget-matched teacher baseline before it can be trusted. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the teacher-induced rollout occupancy measure $\rho_T^K = \frac{1}{K}\sum_{t=0}^{K-1} (T^t)_\# \nu_0$, where $T$ is the fine-tuned teacher, $(T^t)_\#$ is the pushforward of the empirical distribution $\nu_0$ of available downstream initial states after $t$ autoregressive steps, and $K$ is the rollout horizon. TREX samples this measure by rolling the teacher out for long horizons and, at fixed intervals, perturbing the state with Gaussian noise $\tilde{z}_t = z_t + m_t \xi_t$ so the student sees states near the rollout path and the teacher's recovery from them. The student is trained with a $k$-step autoregressive rollout loss on a mixture of ground-truth and teacher-labeled sub-trajectories, and can be chosen as a tensorized Fourier neural operator, a spectral architecture whose Tucker factorization cuts parameter count and enforces translation equivariance.
What would settle it
Take a fine-tuned teacher whose long-horizon rollouts are known to diverge from the true solution (or to converge to a wrong attractor), run TREX with and without ground-truth mixing, and compare the student's error-over-time curve with the teacher's: the claim predicts the student inherits the teacher's long-horizon behavior, so a student that instead tracks the ground truth beyond the training horizon, or that fails to improve in the low-data regime, would count against the central mechanism.
Extended reading notes
Core claim
On its own terms, the paper claims that a PDE foundation model's knowledge can be transferred to a compact surrogate by training the student on the state distribution the teacher visits when rolled out autoregressively from the few available ground-truth trajectories, rather than by sampling new initial conditions. TREX fine-tunes the teacher on the target data, generates long (optionally noise-injected) rollouts, and trains the student on a mix of real and teacher-labeled transitions with a $k$-step rollout loss. The authors report that students trained this way match or surpass the teacher's median relative $L_1$ error at $N=4$ to $N=32$ trajectories on CE-RPUI, NS-SVS, NS-PwC, and NS-BB; the TFNO student uses 3,604 times fewer parameters than Poseidon-L and reduces per-step inference from about 270 ms to about 22 ms, with memory reduced by a factor of about 3.2. They additionally show that the framework extends to the multi-frame Walrus teacher on Kolmogorov flow, and that an equivariant student architecture makes translation-equivariance error essentially zero, a symmetry the transformer-based teacher does not enforce.
Load-bearing premise
TREX's load-bearing premise is that long autoregressive rollouts of the fine-tuned teacher, started from a few ground-truth states and optionally perturbed by Gaussian noise, provide useful supervision for the downstream task; the authors themselves note that if the teacher drifts to an incorrect or distorted attractor, TREX reinforces that bias, and that the noise injections need not satisfy the PDE's physical constraints.
Editorial extensions
If this is right
- A TREX-trained student can serve as a real-time surrogate: on the Poseidon benchmarks, per-step inference falls from about 270 ms to about 22 ms while median relative $L_1$ error stays at or below the teacher's.
- Distillation no longer requires sampling the downstream initial-condition distribution, opening the method to experimental or partially observed data where only a few trajectories exist.
- Enforcing translation equivariance in the student gives it roughly zero equivariance error on periodic domains, and its shifted autoregressive rollouts stay accurate while the teacher's degrade (most visibly on NS-SVS).
- The method transfers across student architectures (TFNO, FNO, U-Net) and across teacher models (Poseidon, Walrus), including multi-frame teachers that cannot be used with IC-KD.
- When the foundation model does not generalize well on the target task relative to a ground-truth-trained student, TREX can hurt; the authors recommend applying it only after a validation check confirms the teacher helps.
Reading between the lines
- The occupancy-measure view suggests a testable scaling law: gains should grow with how much state space the teacher's rollouts cover relative to the handful of observed trajectories, so the method should show smaller improvements on transient or non-recurrent dynamics.
- TREX's noise injection is a generic state-space augmentation, so combining it with physics-aware perturbations (for example, divergence-free or otherwise constraint-preserving noise) could preserve physical validity while keeping the same distillation objective.
- The demonstrated translation equivariance likely extends to other known symmetries of the target PDE, such as rotations on isotropic domains or reflections, but the paper only tests spatial shifts.
- Because TREX adds data rather than changing the loss on existing points, it should compose with feature-based distillation, so an even smaller student may be possible by matching teacher representations as well as rollouts.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes TREX, a knowledge-distillation method for compressing PDE foundation models into compact student surrogates. Starting from a fine-tuned teacher, TREX expands a small set of ground-truth trajectories by autoregressive teacher rollouts with optional periodic Gaussian noise injection, and trains the student on a mixture of ground-truth and teacher-labeled k-step rollout losses. Experiments use Poseidon-L as teacher and TFNO as student on four fluid-dynamics tasks, plus Walrus on Kolmogorov flow, with ablations on rollout length, noise amplitude, student architecture, equivariance, and self-distillation. The main claim is that TREX-trained students match or surpass the teacher's low-data accuracy with far fewer parameters and faster inference, without requiring the initial-condition distribution.
Significance. If the results hold, TREX is a practically useful contribution: it addresses a real obstacle to deploying PDE foundation models as fast surrogates and removes the need for an initial-condition sampler, which is often unavailable for experimental or partially observed systems. The evaluation has genuine strengths: the student is scored against held-out ground-truth trajectories rather than against the teacher, the same student optimizer budget is used across TREX, IC-KD, relational KD, and ground-truth baselines, and the ablations isolate the contributions of rollout extension, noise, student architecture, and self-distillation. The reported parameter reduction (about 3600x), inference speedup (more than 12x), and memory reduction are clear from Table 1. The principal weakness is that the headline comparison between student and teacher does not appear to be optimization-budget matched, so the 'surpass the teacher' part of the claim needs additional experimental support.
major comments (3)
- [§5 and §F] The headline 'match or surpass the teacher' comparison is not budget-matched. The student receives 25k batch updates (Sec. 5), while the Poseidon teacher fine-tuning protocol reported in Sec. F is 200 epochs at batch size 16; on the small downstream training sets this amounts to at most a few thousand optimizer steps, and at N=4 possibly only a few hundred, i.e., far fewer than the student's budget. The teacher is also fine-tuned with the original n-to-n Poseidon loss while the student is trained with a 2-step autoregressive rollout loss. Because no budget-matched teacher baseline or teacher training curves are reported, the abstract and Sec. 5.1 claim that the distilled student can surpass the teacher may be explained by additional optimization on the same few trajectories rather than by TREX's knowledge transfer. Please report the teacher fine-tuning budget, add a teacher baseline trained with the same optimizer budget and ideally the same autoregressive objective, and re-state the headline claim accordingly.
- [Table 5 and Fig. 3] The claim that the student 'surpasses' the teacher is not statistically supported by the reported confidence intervals. For several headline cells the intervals overlap substantially, for example CE-RPUI N=4 (TREX 0.421±0.015 vs. Teacher 0.446±0.074) and NS-SVS N=4 (0.043±0.037 vs. 0.043±0.030), and with only three seeds no paired-significance statement is possible. Please either report paired per-seed differences or soften the 'surpass' wording to 'match or improve in aggregate' where the intervals do not support a stronger statement.
- [Sec. 5, task selection] The four Poseidon tasks were selected on the validation set as cases where the fine-tuned teacher outperforms a TFNO trained only on downstream data. This selection is disclosed, but it means the experiments demonstrate TREX in favorable transfer settings. The abstract and conclusion should explicitly state this restricted scope, and ideally the paper should include at least one task where the teacher does not beat the scratch student, in order to delimit when TREX should be applied.
minor comments (5)
- [§5.1] The text says Poseidon datasets contain 8 time steps (Sec. E), but Fig. 4c and the surrounding discussion report errors at timesteps 8-10 as temporal extrapolation; please clarify how the evaluation extends beyond the stated 8-frame horizon.
- [Table 3] The Walrus experiment is reported as a single run without confidence intervals or repeated seeds; this is a noticeably lower evidentiary standard than the Poseidon experiments and should either be augmented with seeds or explicitly labeled as a pilot study.
- [Eq. (7)] The channel standard deviation s_c is introduced only in prose; please define it directly next to the equation to avoid ambiguity with the noise scaling factor sigma.
- [§5.1] There is a missing space in 'Time stepst>7' and a few other minor typographical errors; a careful proofreading pass would improve readability.
- [§G] The hyperparameters T_Noise, lambda, and the student rollout length k are fixed without a sensitivity study; reporting at least a small ablation for lambda would strengthen the claim that the method is robust to its free parameters.
Circularity Check
No significant circularity: the student is evaluated on held-out ground truth, not on teacher outputs, and teacher rollouts are the method's data-augmentation mechanism rather than a hidden fit.
full rationale
The paper's central claim is that a compact TFNO student distilled from Poseidon-L via TREX matches or surpasses the teacher's accuracy on held-out test trajectories (Fig. 3, Tab. 5). The evaluation metric is the median relative L1 error against ground-truth states, so the student's error is not equal to the teacher's error by construction. The teacher is used as a label generator for additional training trajectories (Sec. 4.1, LS = L_GT + lambda_TREX L_TREX), which is the intended mechanism of distillation, not a fitted input renamed as a prediction. The paper includes multiple controls that break any would-be circularity: an Only Ground Truth baseline, Relational KD, IC-KD, a self-distillation experiment (Sec. H.1) showing that distilling a TFNO teacher into the same architecture gives only a small regularization effect, and an ablation without ground-truth data (Tab. 2). The acknowledged limitation that a biased teacher can bias TREX (Sec. A) is a correctness risk, not a circularity. The only self-citation (Takamoto et al., 2023, for noise injection in autoregressive solvers) is non-load-bearing and is accompanied by independent citations. No load-bearing derivation reduces to its own input, and no 'prediction' is obtained by fitting a parameter to the quantity it is said to predict.
Assumptions & free parameters
free parameters (6)
- Noise scale sigma =
sigma = 1
- Noise injection period T_Noise =
T_Noise = 10
- Teacher rollout length K =
K = 100
- TREX loss weight lambda =
lambda = 1
- Student rollout length k =
k = 2 (main), 5 (U-Net)
- TFNO compression factor =
0.1
assumptions (3)
- domain assumption Teacher long-horizon rollouts from a few ground-truth initial states provide useful supervision for the downstream task.
- domain assumption Gaussian noise injection is an effective state-space augmentation even though perturbed states may violate PDE constraints.
- domain assumption The downstream PDEs are translation-equivariant on periodic domains, and enforcing this in the student is beneficial.
Cite this review
Pith. "Pith review of Distillation of Foundation Models for Time-dependent PDEs." pith.science (2026). https://pith.science/paper/THBYSEWL
@misc{pith2026260811937,
author = {Pith},
title = {Pith review of: Distillation of Foundation Models for Time-dependent PDEs},
year = {2026},
howpublished = {\url{https://pith.science/paper/THBYSEWL}},
note = {Machine review of arXiv:2608.11937}
}
read the original 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.
Figures
Figures from the paper (6 more)
Reference graph
Works this paper leans on
-
[4]
A library for learning neural operators.arXiv preprint arXiv:2412.10354,
Jean Kossaifi, Nikola Kovachki, Zongyi Li, David Pitt, Miguel Liu-Schiaffini, Robert Joseph George, Boris Bonev, Kamyar Azizzadenesheli, Julius Berner, Valentin Duruisseaux, and Anima Anandku- mar. A library for learning neural operators.arXiv preprint arXiv:2412.10354,
-
[6]
Yuxuan Liu, Jingmin Sun, Xinjie He, Griffin Pinney, Zecheng Zhang, and Hayden Schaeffer. Prose-fd: A multimodal pde foundation model for learning multiple operators for forecasting fluid dynamics. arXiv preprint arXiv:2409.09811,
-
[7]
Swin transformer v2: Scaling up capacity and resolution
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. InProceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12009–12019, 2022a. Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell,...
work page 2022
-
[10]
Wal- rus: A cross-domain foundation model for continuum dynamics.arXiv preprint arXiv:2511.15684,
Michael McCabe, Payel Mukhopadhyay, Tanya Marwah, Bruno Regaldo-Saint Blancard, Francois Rozet, Cristiana Diaconu, Lucas Meyer, Kaze WK Wong, Hadi Sotoudeh, Alberto Bietti, et al. Wal- rus: A cross-domain foundation model for continuum dynamics.arXiv preprint arXiv:2511.15684,
-
[12]
Spectral-inspired Operator Learning with Limited Data and Unknown Physics
Han Wan, Rui Zhang, and Hao Sun. Spectral-inspired neural operator for data-efficient pde simulation in physics-agnostic regimes.arXiv preprint arXiv:2505.21573,
-
[13]
Zhanhong Ye, Zining Liu, Bingyang Wu, Hongjie Jiang, Leheng Chen, Minyan Zhang, Xiang Huang, Qinghe Meng Zou, Hongsheng Liu, Bin Dong, et al. Pdeformer-2: A versatile foundation model for two-dimensional partial differential equations.arXiv preprint arXiv:2507.15409,
-
[15]
Unisolver: Pde- conditional transformers towards universal neural PDE solvers
Hang Zhou, Yuezhou Ma, Haixu Wu, Haowen Wang, and Mingsheng Long. Unisolver: Pde- conditional transformers towards universal neural PDE solvers. In Aarti Singh, Maryam Fazel, Daniel Hsu, Simon Lacoste-Julien, Felix Berkenkamp, Tegan Maharaj, Kiri Wagstaff, and Jerry Zhu (eds.),Forty-second International Conference on Machine Learning, ICML 2025, Vancouver...
work page 2025
-
[16]
First, TREX relies on the quality of the fine-tuned teacher
13 A Limitations While TREX provides an effective framework for distilling PDE foundation models, it comes with several limitations. First, 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 transferre...
work page 2022
Show all 21 references
-
[17]
Cheng et al
enforces that the learned model follows the invariant measure by minimizing the Maximum Mean Discrepancy between the data distribution and the learned attractor. Cheng et al. (2025) proposes a Poincaré Flow Neural Network, which is trained by dividing the data into the contrac...
2025
-
[18]
version with the implementation provided by Kossaifi et al. (2025). The factorization reduces the number of parameters with a compression factor of 0.1. The network uses 32 modes per dimension, 32 hidden channels, three layers, and a GELU (Hendrycks & Gimpel,
2025
-
[19]
as provided by Gupta & Brandstetter (2023). The U-Net has 4 levels of resolution, increasing the number of channels per level by 1, 2, 2, and 4, starting at 16 hidden channels after the embedding, resulting in 9,161,332 parameters. For U-Net, we found it advantageous to train ...
2023
-
[21]
aims to make the distance between an input pair i,jfrom a mini-batch in the student feature space close to the distance in the feature space: LRKD = X i<j ℓδ d(ti,tj),d(si,sj) ,(31) wheretandsare the teacher and student features, respectively. The loss is the Huber distance ℓδ...
2023
-
[1963]
Decoupled weight decay regularization
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019,
2019
-
[1984]
Kingma and Jimmy Ba
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.),3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings,
2015
-
[2015]
U-net: Convolutional networks for biomedical image segmentation
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. InMedical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, pp. 234–241. Springer,
2015
-
[2019]
Omnifluids: Unified physics pre-trained modeling of fluid dynamics.arXiv preprint arXiv:2506.10862,
Rui Zhang, Qi Meng, Han Wan, Yang Liu, Zhi-Ming Ma, and Hao Sun. Omnifluids: Unified physics pre-trained modeling of fluid dynamics.arXiv preprint arXiv:2506.10862,
-
[2021]
The model uses the same number of modes and channels as the TFNO, resulting in 5,042,836 parameters
without factorization for the model ablations. The model uses the same number of modes and channels as the TFNO, resulting in 5,042,836 parameters. The FNO uses 4 layers. Poseidon.As the teacher, we use the Poseidon-L model (Herde et al., 2024). Poseidon-L was pretrained on si...
2024
-
[2022]
Raphael Gontijo Lopes, Stefano Fenu, and Thad Starner
IEEE, 2022b. Raphael Gontijo Lopes, Stefano Fenu, and Thad Starner. Data-free knowledge distillation for deep neural networks.arXiv preprint arXiv:1710.07535,
-
[2023]
Gaussian error linear units (gelus).arXiv preprint arXiv:1606.08415,
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus).arXiv preprint arXiv:1606.08415,
-
[2024]
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean
URL https://arxiv.org/abs/2405.19101. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network.arXiv preprint arXiv:1503.02531,
-
[2025]
Stuart, and Anima Anandkumar
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhat- tacharya, Andrew M. Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In9th International Conference on Learning Representations, ICLR ...
2021
Reviewed August 16, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.