Distillation of Foundation Models for Time-dependent PDEs¶

会议: NeurIPS 2026
arXiv: 2608.11937
领域: 模型压缩
关键词: 知识蒸馏, 偏微分方程, 基础模型, 等变性, 自回归外推

一句话总结¶

针对时间依赖偏微分方程（PDE）大模型推理开销极高且缺乏物理归纳偏置的瓶颈，提出教师演化延伸框架 TREX，通过带周期性高斯噪声的自回归演化扩增下游低数据样本，蒸馏出参数量降低超 3000 倍且具备平移等变性的轻量学生模型。

研究背景与动机¶

偏微分方程（PDE）是流体力学、天气预报、天体物理和工程设计等科学计算领域的核心数学建模工具。传统的数值求解器（如有限体积法、有限元法）在处理复杂非线性动态或大尺度三维系统时计算开销巨大，而基于深度学习的神经算子代理模型不仅推理可微分，还能直接融合观测数据，展现出极高的应用价值。为了缓解神经算子在小样本下游场景下严重依赖数值求解器生成海量训练数据的困境，近年来基于大规模物理系统预训练的 PDE 基础模型（如 Poseidon、Walrus）应运而生。这些大模型在数个轨迹样本微调后即可实现极佳的少样本泛化性能。

然而，现存的 PDE 基础模型通常基于拥有数亿至数十亿参数的庞大 Vision Transformer 架构，单步与全轨迹自回归推理延迟极高且显存占用巨大，无法满足工业制造、风电实时控制等闭环低延迟场景的严苛要求。更为关键的是，为了适配多样化物理方程，通用大模型通常舍弃了具体的物理归纳偏置（如平移对称性、旋转对称性、频域衰减特性等），使得模型在长时自回归推演时误差容易累积发散，且对微小的空间平移极为敏感。

如果尝试直接用传统的知识蒸馏方法压缩大模型，现存范式通常假设已获知下游任务初始条件（IC）的显式分布，从而通过大量重新采样初值来让教师模型生成合成标签。但在真实物理实验和部分可观测系统中，研究者往往仅能获取极少数固定的短观测轨迹，根本无法得知甚至无法采样底层的初始分布；此外对于 Walrus 等依赖多帧上下文窗口输入的大模型，单纯采样单帧初值也无法脱离原始仿真器构造有效的自回归启动状态。本文的切入角度是：耗散型 PDE 与混沌动力系统通常具有吸引子或不变测度结构，即便从少数锚点轨迹出发，只要沿时间轴进行足够长的自回归外推，系统状态便会逐渐覆盖教师动力学所诱导的状态占有测度空间。核心 idea：提出教师演化延伸蒸馏框架 TREX，无需初值分布先验，仅利用少量下游真实轨迹作为演化锚点，通过注入周期性噪声的教师长程自回归推演生成丰富且具备抗扰动恢复能力的状态跃迁样本，并将通用无偏大模型蒸馏进内嵌平移等变性与频域正则化约束的极紧凑神经算子学生模型中。

方法详解¶

整体框架¶

TREX 解决的是极低数据（如仅 4-32 条短轨迹样本）下游场景下，如何将庞大的通用 PDE 基础模型有效压缩为超轻量、低延迟且物理一致的专用代理模型。系统整体包含三个阶段：下游微调阶段、教师演化延伸生成数据阶段以及物理约束学生联合蒸馏阶段。

首先，在大规模多物理数据上预训练的基础模型作为教师网络 $\mathcal{T}$，在给定的少量真实短轨迹数据集 $\mathcal{D}_{\mathrm{GT}}$ 上完成全参数或高效微调。接着，进入 TREX 的核心生成环节：以真实轨迹的初始状态为锚点，运行教师模型自回归多步演化，且每隔固定步长注入受控尺度的通道高斯噪声，以促使演化轨迹探索更广阔的状态空间并在扰动后向未受扰动流形自发恢复，形成合成轨迹集 $\mathcal{D}_{\mathrm{TREX}}$。最后，构建内嵌任务物理先验（如平移等变性、频域截断）的紧凑神经算子（如张量化傅里叶神经算子 TFNO）作为学生模型 $\mathcal{S}$，在真实数据自回归损失与合成数据自回归蒸馏损失的混合目标下完成训练。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["下游极少真实观测轨迹<br/>D_GT: 仅4-32条短轨迹"] --> B["阶段1：基础模型教师微调<br/>预训练大模型在D_GT上适配目标域"]
    B --> C["阶段2：教师演化延伸 TREX<br/>以真实初值为锚点进行长程自回归推演"]
    C --> D["阶段3：周期性高斯噪声注入<br/>每隔T_Noise步扰动状态并由教师自恢复演化"]
    D --> E["生成合成丰富状态集 D_TREX<br/>获取教师诱导长时占有测度样本"]
    E --> F["阶段4：内嵌物理归纳偏置学生模型<br/>构建紧凑平移等变TFNO神经算子"]
    A --> G["阶段5：多步自回归联合蒸馏训练<br/>真实数据监督 + 教师外推样本联合优化"]
    F --> G
    G --> H["轻量高效且保对称性的学生算子<br/>参数减少>3000倍，推理加速>10倍"]

关键设计¶

1. 教师诱导状态占有测度的长程自回归外推：摆脱初值分布依赖

在下游实验和复杂物理观测中，真实初值分布 $p_0$ 往往不可见或难以解析采样。针对该痛点，TREX 放弃了重新采样初始条件的传统数据生成逻辑，转而利用动力系统的长时渐近特性。许多耗散和混沌偏微分方程在相空间中存在吸引子或不变测度 $\mu_{\mathcal{T}}$，满足 $\mathcal{T}_{\#}\mu_{\mathcal{T}} = \mu_{\mathcal{T}}$。给定样本经验初值分布 $\nu_0 = \frac{1}{N}\sum_{n=1}^N \delta_{\mathbf{u}_n^0}$，标准单步训练仅能在有限的观测时刻覆盖狭窄流形，而 TREX 构造了一个跨越 $K$ 步演化的教师状态占有测度（Occupancy Measure）： $$\rho_{\mathcal{T}}^K = \frac{1}{K}\sum_{t=0}^{K-1} (\mathcal{T}^t)_{\#}\nu_0$$ 式中 $(\mathcal{T}^t)_{\#}\nu_0$ 表示初值分布在教师算子作用 $t$ 步后的前推测度。通过长时间演化，不仅让学生模型得以见识自回归测试期间必然遭遇的累积漂移状态，更通过动力系统的遍历与混合特性让轨迹自发探索流形中的广阔代表性区域，彻底消除了对仿真器额外初值调用的依赖。

2. 周期性高斯噪声注入与扰动自恢复监督：强化局部状态空间覆盖

如果仅仅执行完全确定性的长程推演，样本依然严格受限于教师单一预测轨迹上，对轨迹附近的离轴状态缺乏感知，导致学生在遭遇微小扰动时缺乏纠偏能力。TREX 提出了间隔为 $T_{\mathrm{Noise}}$ 的周期性高斯噪声注入机制：在每个周期节点 $t \equiv 0 \pmod{T_{\mathrm{Noise}}}$ 且 $t>0$ 时，对当前状态 $\mathbf{z}^t$ 施加与各个通道全局标准差 $s_c$ 成比例的空间独立高斯扰动： $$\tilde{\mathbf{z}}^t = \mathbf{z}^t + m_t \boldsymbol{\xi}^t, \quad (\boldsymbol{\xi}^t)_c(\mathbf{x}) \sim \mathcal{N}(0, \sigma^2 s_c^2), \quad m_t = \mathbf{1}\{t>0 \land t \equiv 0 \pmod{T_{\mathrm{Noise}}}\}$$ 关键机制在于，传统数据增强通常仅在训练输入中加入噪声而保持真实标签不动，而 TREX 在状态扰动后立即由教师模型根据受扰状态计算下一步真实转移：$\mathbf{z}^{t+1} = \mathcal{T}(\tilde{\mathbf{z}}^t)$。在随后的 $T_{\mathrm{Noise}}$ 步无扰动推演中，教师模型展示了动力学如何从扰动偏离态自然回落、收敛至未扰动的典型流形。这种设计让学生同时掌握了无噪长程外推轨迹与局部扰动恢复动态，大幅提升自回归稳定性。

3. 结构化物理归纳偏置内嵌：打破教师模型的对称性缺陷

现存 Vision Transformer 基础模型因追求通用性而缺少物理对称性硬约束。在具有周期性边界条件的均匀介质中，控制方程具有天然的空间平移等变性，即对空间平移算子 $\tau_a$ 满足 $\Phi_{\Delta t}(\tau_a \mathbf{u}) = \tau_a \Phi_{\Delta t}(\mathbf{u})$。若由大模型单独推演，微小的对称性破缺误差会在数十步后急剧放大。TREX 将教师知识蒸馏至张量化傅里叶神经算子（TFNO）中。TFNO 在频域直接进行谱乘积截断并应用 Tucker 张量分解，天然具备连续空间平移等变性，同时显式截断了高频高斯噪声模态。通过知识蒸馏，教师网络向学生传递从海量预训练中习得的复杂动力学演化知识，而学生模型的网络架构本身则强制保证了严格平移不变性与高频正则化，形成了“大模型经验知识”与“小模型物理归纳偏置”的高效互补。

损失函数 / 训练策略¶

学生模型 $\mathcal{S}$ 采用多步残差自回归预测结构 $\hat{\mathbf{u}}^{t+1} = \mathcal{S}(\mathbf{u}^t) = \mathcal{M}(\mathbf{u}^t) + \mathbf{u}^t$。训练总损失由真实轨迹监督损失 $\mathcal{L}_{\mathrm{GT}}$ 和教师生成演化数据损失 $\mathcal{L}_{\mathrm{TREX}}$ 线性加权组成： $$\mathcal{L}_{\mathcal{S}} = \mathcal{L}_{\mathrm{GT}} + \lambda_{\mathrm{TREX}} \mathcal{L}_{\mathrm{TREX}}$$ 默认权重 $\lambda_{\mathrm{TREX}} = 1$。为了避免单步训练引起的自回归漂移，两部分损失均在长度为 $k$ 的自回归子序列（如 $k=2$）上计算。对于 $\mathcal{D}_{\mathrm{TREX}}$，损失函数定义为： $$\mathcal{L}_{\mathrm{TREX}} = \frac{1}{|\mathcal{D}_{\mathrm{TREX}}|} \sum_{(\mathbf{z}^i, \dots, \mathbf{z}^{i+k}) \in \mathcal{D}_{\mathrm{TREX}}} \frac{1}{k} \sum_{t=i}^{i+k-1} \ell\left(\mathbf{z}^{t+1}, \mathcal{S}(\hat{\mathbf{u}}^t)\right), \quad \hat{\mathbf{u}}^i = \mathbf{z}^i, \ \hat{\mathbf{u}}^{t+1} = \mathcal{S}(\hat{\mathbf{u}}^t)$$ 其中误差度量 $\ell$ 沿用 Poseidon 的归一化相对 $L^1$ 损失（在感兴趣物理量 QOI 上跨空间和通道求和归一）。优化器使用 Adam，基础学习率 $10^{-3}$，配合余弦退火衰减至 $10^{-6}$。

实验关键数据¶

主实验¶

实验采用 6.29 亿参数的 Poseidon-L 作为教师模型，蒸馏为仅 17.4 万参数的 TFNO 学生模型（参数量降低 3604 倍）。在 2D 可压缩欧拉方程（CE-RPUI）与不可压缩 Navier-Stokes 方程的三种不同初值配置（NS-PwC、NS-SVS、NS-BB）上评估。主表汇报了在仅有 $N=4$ 条真实下游训练轨迹的极端少样本设定下，最后一步推演的相对 $L^1$ 中位数误差（三组随机种子均值）以及计算效率对比。

方法 / 模型	CE-RPUI ($N=4$)	NS-SVS ($N=4$)	NS-PwC ($N=4$)	NS-BB ($N=4$)	参数量	单步推演耗时	全轨迹推演耗时	显存占用 (全轨迹)
Poseidon-L (Teacher, autoreg.)	0.446	0.043	0.070	0.046	628.6M	269.4 ms	1887.7 ms	3749 MB
Only GT (TFNO 学生纯真值训练)	0.608	0.171	0.282	0.178	174.4K	22.1 ms	152.1 ms	1171 MB
Relational KD (特征关系蒸馏)	0.608	0.178	0.175	0.111	174.4K	22.1 ms	152.1 ms	1171 MB
IC-KD (需采样真实初值分布)	0.440	0.045	0.067	0.041	174.4K	22.1 ms	152.1 ms	1171 MB
TREX (本文方法，无初值依赖)	0.421	0.043	0.056	0.036	174.4K	22.1 ms	152.1 ms	1171 MB

此外，在基于 1.3B 参数的大规模多帧上下文物理大模型 Walrus 蒸馏至 TFNO 的柯尔莫哥洛夫湍流（Kolmogorov Flow, $Re=5000$）实验中，TREX 同样展现出惊人的效率提升：全序列推演耗时从 Walrus 的 21.425 秒急剧缩短至 0.414 秒（加速超 51 倍），且长程自回归误差 $\text{VRMSE}_{T \in [1:20]}$ 从 0.214 显著降至 0.123。

消融实验¶

在 NS-PwC ($N=4$) 任务上，作者对真实数据监督与周期性噪声注入进行了系统性剥离，探究 TREX 为何能够实现超越教师模型的性能。

配置	中位数相对 $L^1$ 误差	说明
TREX 完整方法	0.056 ± 0.003	融合真值 + 教师带噪长程演化数据
w/o GT (完全移除真实轨迹监督)	0.060 ± 0.002	仅使用教师生成的演化数据，仍超越教师
w/o GT & w/o Noise (无真值且演化无噪声)	0.064 ± 0.006	仅使用纯确定性教师演化数据
Teacher (Poseidon-L 教师基线)	0.070 ± 0.011	微调后的教师模型自身精度

针对平移对称性破缺的专门评估（输入随机二维空间滚动位移 $\tau_a$）证明了等变归纳偏置的关键作用：

数据集	教师等变误差	学生等变误差	教师原始终步误差	教师平移终步误差	学生原始终步误差	学生平移终步误差
NS-BB	0.026 ± 0.013	0.000 ± 0.001	0.0651	0.0676	0.0504	0.0504
CE-RPUI	0.101 ± 0.006	0.001 ± 0.001	0.4189	0.4108	0.3549	0.3549
NS-PwC	0.023 ± 0.006	0.000 ± 0.001	0.1095	0.1041	0.0744	0.0744
NS-SVS	0.110 ± 0.005	0.000 ± 0.001	0.0820	0.2398 (严重恶化)	0.0755	0.0755 (保持恒定)

关键发现¶

学生反超教师的自正则化机制：即便完全移除真实轨迹监督（w/o GT），仅依靠教师模型带噪推演生成的合成数据训练，学生模型的误差（0.060）仍然显著低于教师本身（0.070）。这一反超源于动力学吸引子的高密度采样覆盖以及轻量学生算子所带来的高频滤波与参数自正则化效应。
高噪声幅度的反直觉增益：噪声尺度敏感度实验表明，将相对高斯标准差缩放因子设为高达 $\sigma = 1.0$（即扰动幅度等同于通道标准差本身）时取得最佳蒸馏效果。教师网络在强噪声注入后能迅速在随后的子轨迹演化中自发将流体场“去噪修复”，为学生提供了宝贵的多尺度相空间纠错样本。
对称性破缺的灾难性消除：在 NS-SVS 剪切层问题上，未加等变约束的 Vision Transformer 教师模型在初值发生微小空间平移后，最终步相对误差从 0.0820 剧增至 0.2398（误差放大近 3 倍）；而 TFNO 学生模型由于网络算子内置平移等变性，平移前后误差完全一致（0.0755），彻底消除了长时自回归推演中的漂移崩溃。

亮点与洞察¶

动力学占有测度视角的蒸馏创新：传统大模型蒸馏依赖采样新样本，而时间依赖系统天然具备时间轴演化能力；TREX 巧妙地利用耗散动力系统的遍历性，将“初值空间探索”转化为“沿时间轨道的相空间遍历”，在不调用外部仿真器的情况下自举出海量高质量离轴物理状态。
大模型泛化性与轻量算子归纳偏置的解耦共融：基础大模型参数多、表征强，擅长吸收跨物理系统的先验；而下游特定算子参数小、速度快，能够精准硬编码对称性与频域先验。TREX 将两者优势有机结合，为科学基础模型走向工业级边缘部署开辟了明确范式。

局限与展望¶

受限于教师模型的长期吸引子稳定性：TREX 极度依赖微调后教师长程推演的物理保真度。若教师基础模型自身存在数值耗散过快、伪震荡或假吸引子坍缩等长期累积偏差，这些病态伪影将作为系统性标签噪声全盘传导至学生。
对非遍历或强瞬态系统的有效性递减：对于迅速衰减至零解的平庸系统、具有多个互不连通相空间的复杂系统、或具有强间歇性瞬态的冲击波流场，纯粹延长时间维度的推演无法有效覆盖全局物理相空间。
高斯加性噪声的物理非保真性：当前直接注入的空间各向同性高斯白噪声并未显式满足无散度约束（$\nabla \cdot \mathbf{u} = 0$）或边界条件，未来可探索保对称性或沿能量级联规律注入的物理一致性扰动机制。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从动力系统占有测度与周期自恢复角度切入偏微分方程基础模型蒸馏，思想十分精妙。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 Poseidon 与 Walrus 两大顶尖大模型，囊括可压缩欧拉、三类 Navier-Stokes 以及柯尔莫哥洛夫湍流，涵盖等变性与消融全面验证。
写作质量: ⭐⭐⭐⭐⭐ 动力系统数学表述严谨，机理解释透彻，实验数据详实清晰。
价值: ⭐⭐⭐⭐⭐ 成功将亿级参数大模型压缩超 3000 倍且加速超 10-50 倍，对科学智能的工程落地与实时控制具有重大实用价值。