{"id":"979cae56-cc6a-4e47-8634-c807552854dd","arxiv_id":"2608.11937","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":6,"one_line_summary":"TREX distills a fine-tuned PDE foundation model into a compact student by training on noisy teacher rollouts, matching or beating teacher accuracy with thousands of times fewer parameters.","lead":"This paper introduces TREX, a method for compressing large pretrained PDE foundation models into small, fast surrogate models by training the small model on long rollouts generated by the large one. If the results hold, accurate neural PDE surrogates could run in real time with thousands of times fewer parameters and an order of magnitude faster inference.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The 'student surpasses teacher' headline may rest on unequal training budgets: students receive 25k batch updates (Sec. 5) while the Poseidon protocol in Sec. F is 200 epochs, a much smaller budget at low N.","rationale":"I read the central claim as the abstract's capability claim that a compact TREX student can match or surpass the fine-tuned teacher's accuracy with far fewer parameters and faster inference. The reader's weakest assumption is the quality of teacher rollouts as supervision; that concern is real but it is explicitly acknowledged in the Limitations, and the w/o-GT ablation and the student-vs-teacher error-growth plots already provide partial evidence that the student does not simply copy teacher drift. The most load-bearing unaddressed issue is the training budget asymmetry: the student is trained for a fixed 25k batch updates, while the only stated Poseidon protocol is 200 epochs, which at the lowest data counts is roughly two orders of magnitude fewer updates. If the teacher baseline is undertrained, the headline comparison overstates what TREX achieves. This concern is concrete, not acknowledged in the paper, and can be settled by a single budget-matched fine-tuning experiment. I therefore disagree with the reader's choice of weakest assumption, while agreeing with the overall CONDITIONAL verdict. The recommendation is UNCHANGED because the concern reinforces the existing conditional status rather than moving it to a different verdict.","tokens_in":19030,"tokens_out":22811,"duration_ms":244597,"concrete_test":"Fine-tune Poseidon-L on the same N in {4,8,16,32} trajectory subsets for the same 25k batch updates used for the students, using the Sec. F optimizer settings and the original n-to-n loss, then re-evaluate the Teacher rows of Table 5 with the same median relative L1 metric. If the budget-matched teacher errors fall below the corresponding TREX student medians at N=4 or N=8, the 'match or surpass teacher' headline should be qualified or the paper should adopt a stronger teacher baseline; if the teacher errors remain above the student's, the concern is resolved.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The central quantitative claim (Fig. 3, Table 5, Sec. 5.1) compares a TFNO student trained for 25k batch updates with a fine-tuned Poseidon-L teacher. Section 5 states that the number of epochs for all student methods is adjusted so that each method receives 25k batch updates when using only ground truth. The only Poseidon training budget reported in Sec. F is 200 epochs at batch size 16; at N=4 this is roughly 400 optimizer steps, and even at N=32 it is about 3,200 steps, versus 25,000 for the student. If this 200-epoch protocol is also the teacher fine-tuning budget, the Teacher column in Fig. 3 and Table 5 is an undertrained baseline, and the student's match-or-surpass result could be explained by extra optimization on the same few trajectories rather than by TREX's knowledge transfer. The paper does not report a budget-matched teacher baseline or the teacher fine-tuning budget separately, so the headline comparison is not isolated from training effort. This does not affect the TREX-vs-IC-KD comparison, since both use the same student budget, but it does affect the abstract's claim that the distilled student can match or surpass the teacher's accuracy.","agreement_with_reader":"disagree"},"referee_report":{"model":"deepseek-v4-flash","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.","tokens_in":19268,"tokens_out":8993,"duration_ms":88325,"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":[{"comment":"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.","section":"§5 and §F"},{"comment":"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.","section":"Table 5 and Fig. 3"},{"comment":"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.","section":"Sec. 5, task selection"}],"minor_comments":[{"comment":"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.","section":"§5.1"},{"comment":"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.","section":"Table 3"},{"comment":"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.","section":"Eq. (7)"},{"comment":"There is a missing space in 'Time stepst>7' and a few other minor typographical errors; a careful proofreading pass would improve readability.","section":"§5.1"},{"comment":"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.","section":"§G"}],"recommendation":"major_revision","confidential_remarks":"The main experimental gap is the teacher baseline budget. If the authors can add a fine-tuned teacher trained for 25k updates, ideally with the same autoregressive objective, without changing the student results, the paper would be suitable; otherwise the abstract's 'surpass the teacher' statement should be removed or qualified to 'match or improve in aggregate on selected tasks'. The paper is otherwise careful and honest, including its limitations section, and I do not see grounds for rejection."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Max,\n\nI read the TREX paper on distilling PDE foundation models. The core idea is genuinely new: rather than sampling new initial conditions as previous distillation methods (Omnifluids, SINO) require, it extends the few available ground-truth trajectories by rolling out the fine-tuned teacher, with periodic noise injection. That removes the need for an IC sampler and gives a student access to states the teacher actually visits autoregressively. This is a clean contribution, and the paper demonstrates it convincingly across four Poseidon tasks and one Walrus task, with three seeds and confidence intervals. The student TFNO is about 3600x smaller and over 10x faster, and often matches or beats the teacher's median relative L1 error. The equivariance experiments are a nice extra: the student enforces translation equivariance, and it holds up on shifted rollouts where the teacher degrades.\n\nThe writing is honest. The limitations section acknowledges that the student inherits teacher bias, and the ablations isolate the effects of the rollouts and the noise. The self-distillation control in the appendix shows the gain is not just from self-training.\n\nThe main soft spot is the unequal training budget in the headline comparison. Students are trained for 25k batch updates; the Poseidon teacher is fine-tuned for 200 epochs, which at N=4 is roughly 400 optimizer steps and only a few thousand steps at N=32. The paper never reports a budget-matched teacher baseline, so the 'student surpasses teacher' claim could partly be an artifact of more optimization on the same trajectories, rather than of TREX's knowledge transfer. This does not affect the TREX-vs-IC-KD comparison, since both use the same student budget, and that is the key new comparison. But the abstract's claim is weaker than it looks without a matched teacher.\n\nOther concerns are minor: the downstream tasks were selected because the teacher outperforms a ground-truth-only student, which is disclosed but limits generality; no code is released; and there is no simple teacher-free noise-augmentation control to show that teacher rollouts are doing more than generic data augmentation.\n\nStill, this is a solid paper for the PDE surrogate community. The method is new, the experiments are mostly careful, and the open issues are addressable in a revision. I'd send it to peer review. The main thing I'd ask the authors to add is a budget-matched teacher baseline and a teacher-free augmentation control. Short of that, the central claim should be phrased more modestly.\n\nI'd bring it to our reading group. I'll probably cite it.\n\nLet's chat if you want to hash out the budget issue further.","headline":"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.","tokens_in":19828,"tokens_out":4787,"would_cite":true,"duration_ms":45483,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"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…","keywords":["knowledge distillation","PDE foundation models","neural operators","autoregressive rollout","translation equivariance","surrogate modeling","low-data regime","tensorized Fourier neural operator"],"falsifier":"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.","tokens_in":18806,"feed_emoji":"🌊","tokens_out":10207,"duration_ms":89964,"temperature":0.7,"pith_summary":"Large pretrained PDE foundation models adapt to new physics with only a few trajectories, but their size and inference cost make them impractical as fast surrogates. This paper proposes Teacher Rollout Extension (TREX): after fine-tuning the foundation model on a handful of target trajectories, the teacher is rolled out for long horizons (with periodic Gaussian noise injection) and a compact student is trained on the resulting teacher-labeled transitions. TREX does not need a sampler for the downstream initial-condition distribution, because it reuses the available trajectories as starting points. On Poseidon and Walrus benchmarks, the resulting students match or beat the teacher's median relative $L_1$ error while using about 3,604 times fewer parameters and cutting per-step inference from roughly 270 ms to about 22 ms.","feed_headline":"Teacher rollouts shrink PDE foundation models 3,604x","feed_subtitle":"TREX distillation trains a compact surrogate that matches the teacher at a fraction of the size and cost.","key_machinery":"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.","core_discovery":"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.","pith_inferences":["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."],"forward_implications":["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."],"supporting_citations":[{"why":"Supplies the Poseidon foundation model and the four fluid-dynamics benchmark tasks used in the main distillation experiments.","marker":"Herde et al., 2024"},{"why":"Supplies the Walrus multi-frame foundation model used as teacher in the Kolmogorov-flow distillation experiment.","marker":"McCabe et al., 2025"},{"why":"Introduces the tensorized FNO architecture used as the main compact student.","marker":"Kossaifi et al., 2024"},{"why":"Provides the implementation of the TFNO used in the experiments.","marker":"Kossaifi et al., 2025"},{"why":"Defines relational knowledge distillation, the feature-based baseline against which TREX is compared.","marker":"Park et al., 2019"},{"why":"Defines the IC-KD baseline that assumes access to the initial-condition distribution, the key comparison for TREX's no-sampler advantage.","marker":"Zhang et al., 2025"},{"why":"Introduces the Fourier neural operator used as a student architecture in the model-ablation experiments.","marker":"Li et al., 2021"},{"why":"Provides the Kolmogorov-flow dataset used to test distillation of the multi-frame Walrus teacher.","marker":"Li et al., 2022"}],"fun_headline_variants":["TREX distillation shrinks PDE foundation models 3,604x","Teacher rollouts train compact PDE surrogate with 12x faster inference","PDE foundation model distilled to 3,604x smaller student","Rollout-based distillation yields fast, small PDE surrogate","Equivariant student from TREX matches teacher at fraction of size"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"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.","fun_headline_variants_meta":{"raw":{"variants":["TREX distillation shrinks PDE foundation models 3,604x","Teacher rollouts train compact PDE surrogate with 12x faster inference","PDE foundation model distilled to 3,604x smaller student","Rollout-based distillation yields fast, small PDE surrogate","Equivariant student from TREX matches teacher at fraction of size"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000427,"raw_usage":{"total_tokens":2216,"prompt_tokens":1007,"completion_tokens":1209,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":623,"completion_tokens_details":{"reasoning_tokens":1120}},"tokens_in":623,"tokens_out":1209,"duration_ms":11154,"temperature":1.0,"reasoning_tokens":1120,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-16T00:22:03.465470+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"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.","supporting_citations":[{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Provides the implementation of the TFNO used in the experiments."}],"review_version":1}