Distillation of Foundation Models for Time-dependent PDEs

📄 arXiv: 2608.11937v1 📥 PDF

作者: Daniel Musekamp, Boshra Ariguib, Andrei Manolache, Mathias Niepert

分类: cs.LG

发布日期: 2026-08-12

💡 一句话要点

提出TREX框架以高效蒸馏时间依赖PDE模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 知识蒸馏 偏微分方程 模型压缩 合成数据 高效推理

📋 核心要点

  1. 现有的基础模型在时间依赖的PDE任务中虽然表现良好，但其庞大的体积和高计算需求限制了实际应用。
  2. TREX框架通过知识蒸馏，将预训练的教师模型的能力转移到一个更小的学生模型，同时生成合成数据以增强训练。
  3. 实验结果表明，使用TREX的学生模型在准确性上可以匹敌或超越教师模型，同时参数数量大幅减少，推理速度显著提升。

📝 摘要（中文）

针对时间依赖的偏微分方程（PDE），基础模型在多样化的物理系统上训练，能够有效泛化到新的下游任务。通过仅在目标领域的少量轨迹上进行微调，这些模型在低数据环境下也能实现较高的准确性。然而，这些模型通常体积庞大且计算密集，限制了其作为数值求解器快速替代品的实用性。为此，本文提出了教师回放扩展（TREX）框架，通过知识蒸馏将预训练基础模型的预测能力转移到一个紧凑高效的学生模型上。TREX通过生成长的合成轨迹来增强有限的下游数据，学生模型能够在参数数量上减少几个数量级，并在推理速度上实现超过一个数量级的提升。

🔬 方法详解

问题定义：本文旨在解决时间依赖的偏微分方程（PDE）模型在实际应用中的高计算成本和大模型体积问题。现有的基础模型虽然在准确性上表现优异，但其庞大的参数量和计算需求使得在低数据环境下的应用受到限制。

核心思路：论文提出的TREX框架通过知识蒸馏，将预训练的教师模型的预测能力转移到一个紧凑的学生模型上。通过生成长的合成轨迹，TREX能够在不依赖初始条件分布的情况下增强学生模型的训练数据，从而提高其泛化能力。

技术框架：TREX的整体架构包括教师模型的微调、合成轨迹的生成以及学生模型的训练三个主要模块。教师模型首先在目标领域进行微调，然后通过回放生成合成数据，最后利用这些数据训练学生模型。

关键创新：TREX的创新在于其通过教师回放生成合成数据的能力，使得学生模型能够接触到长时间范围的状态和局部恢复行为，从而提升了模型的学习效果。这一方法与传统的蒸馏方法不同，后者通常依赖于真实数据。

关键设计：在TREX中，学生模型可以引入任务特定的归纳偏置，例如等变性，这些偏置教师模型并不一定强制执行。此外，合成轨迹的生成过程中可以选择性地注入周期性噪声，以增强模型的鲁棒性。实验中，学生模型的参数数量减少了几个数量级，推理速度提升超过一个数量级。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，使用TREX框架训练的学生模型在多个PDE基准测试中，准确性能够匹敌或超越教师模型，同时参数数量减少了几个数量级，推理速度提升超过一个数量级，展现出显著的性能优势。

🎯 应用场景

该研究的潜在应用领域包括科学计算、工程模拟和实时物理系统建模等。通过提供高效的模型，TREX能够在低数据环境下实现快速且准确的预测，具有重要的实际价值。未来，该方法可能会推动更多领域的基础模型应用，尤其是在资源受限的情况下。

📄 摘要（原文）

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.