瞬态偏微分方程基础模型的 蒸馏

瞬态偏微分方程 (PDE) 的基础模型在大量且多样化的物理系统集合上进行训练，并且可以有效地推广到新的下游任务。在仅对目标域的几个轨迹进行 微调 后，它们就可以在低数据情况下实现很高的准确性。然而，这些模型通常很大且计算密集，限制了它们作为数值求解器的快速替代品的有用性。我们提出了 Teacher Rollout Extension (TREX)，这是一个 知识蒸馏 框架，可将预训练基础模型的预测能力转移到紧凑而高效的学生中。从微调教师开始，TREX 通过教师预测轨迹生成长合成轨迹，可选地定期注入噪声，从而增强有限的下游数据。该过程从教师诱导的预测轨迹分布中进行采样，而不需要明确了解初始条件分布，同时使学生暴露于长视野状态和自回归预测期间遇到的状态周围的局部恢复行为。学生可以进一步纳入特定于任务的归纳偏差，例如老师不一定强制执行的等方差。我们在多个 PDE 基准上评估 TREX。由此产生的学生可以达到或超过老师的准确性，同时将参数数量减少几个数量级，并实现超过一个数量级的推理加速。