Spectral-Guided Physical Dynamics Distillation¶
会议: ICLR 2026
OpenReview: https://openreview.net/forum?id=P6F4MxtOKp
领域: 物理动力学建模 / 知识蒸馏
关键词: 物理动力学预测, 谱图分析, 知识蒸馏, 特权信息, 时空表征
一句话总结¶
针对"只给初始状态、要预测粒子长时程 3D 轨迹"这一难题,本文提出 SGDD:用一个能看到未来轨迹的教师编码器作为"特权信息",在统一时空谱域上自适应加权关键频率分量,再把这种富含动力学信息的表征蒸馏给只看初始状态的学生编码器,从而在分子、蛋白质、人体运动等多尺度系统上实现更准更稳的长时程预测。
研究背景与动机¶
领域现状:物理动力学预测——给定粒子(分子原子、蛋白质骨架、人体关节)的初始状态,预测其未来 3D 轨迹——是科学与工程中的基础任务。近年来主流做法是设计等变神经网络(EGNN、ClofNet、SE(3)-Transformer 等)来捕捉物理系统的对称性;更进一步的工作(ESTAG、EGNO、GF-NODE)开始引入频率感知,用傅里叶 / 图傅里叶手段建模时间或空间上的周期结构。
现有痛点:从初始状态做长时程预测时误差会被急剧放大。根本原因是长时程会把全局低频趋势和局部高频振荡深度纠缠在一起。现有频率感知方法有两个具体缺陷:其一,它们把时间和空间分开建模,只从单一维度(要么时间、要么空间)导出谱表征,无法刻画时空相互依赖中涌现的物理过程;其二,它们对所有频率分量"一视同仁",没意识到不同系统里低频/高频的重要性其实差异很大。
核心矛盾:长时程预测需要优先保住低频模式来维持稳定与长期一致性,同时又要互补地补上高频细节来提升短期精度——但低频与高频在时空上交织,单维度的谱建模既分不开它们,也不知道该给哪段频率更多权重。
本文目标:(1) 在统一时空域上联合导出谱表征,而不是时间、空间各做各的;(2) 自适应地强调任务相关的频率分量;(3) 解决推理时拿不到未来轨迹的根本约束。
切入角度:未来轨迹本身就是一种特权信息(Privileged Information)——训练时可见、测试时不可见。如果让一个教师编码器吃进真实未来轨迹、提炼出富含动力学的频率感知表征,再通过知识蒸馏把它"教"给只能看初始状态的学生编码器,就能给学生一个直接而高效的监督信号。
核心 idea:用"教师看未来 + 学生看初始 + 时空谱域蒸馏"替代单维度频率建模,让学生在推理时无需特权信息也能生成有效的动力学表征。
方法详解¶
整体框架¶
SGDD(Spectral-Guided Dynamics Distillation)要解决的是"只有初始状态 \(G_0\),却要预测整条未来轨迹 \(\{x_1,\dots,x_T\}\)"。整体范式是 编码器→动力学表征 z→解码器:解码器是一个现成的物理动力学模型(本文用 EGNO 或 GF-NODE),表征 \(z\) 由编码器从 \(G_0\) 产生,浓缩了对未来演化的预判。SGDD 的全部创新都在"如何造一个好的 \(z\)"上。
具体地,框架并行训练两个编码器:动力学编码器 \(E_{dyn}\) 吃进特权未来序列 \(G_{1:T}\) 与 \(G_0\),产出 \(z_{dyn}\);初始编码器 \(E_{init}\) 只吃 \(G_0\),产出 \(z_{init}\)。两者的表征都送进谱引导增强模块,该模块用一组时空联合基把表征投到谱域、自适应地重加权最相关的频率分量,得到 \(z^{sg}_{dyn}\) 与 \(z^{sg}_{init}\)。训练时通过知识蒸馏强迫 \(z^{sg}_{init}\) 去模仿 \(z^{sg}_{dyn}\)(同时在时空域和谱域对齐);推理时教师整支退场,解码器仅靠 \(z^{sg}_{init}\) 预测轨迹。整个框架端到端训练,并采用分阶段训练策略保证收敛稳定。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["初始状态 G0"] --> C["初始编码器 E_init<br/>仅看 G0"]
B["未来轨迹 G_1:T<br/>(特权信息·仅训练)"] --> D["动力学编码器 E_dyn<br/>看未来+初始"]
C --> E["谱引导增强 SGE<br/>时空谱域自适应重加权"]
D --> E
E --> F["学生表征 z_init^sg"]
E --> G["教师表征 z_dyn^sg"]
G -->|双层蒸馏对齐<br/>(时空+谱)| F
F --> H["物理动力学解码器<br/>推理仅用 z_init^sg"]
H --> I["未来轨迹预测 x_1:T"]
关键设计¶
1. 双编码器 + 特权信息蒸馏:让学生"偷看"未来再独立答题
痛点在于:标准的"编码-表征-解码"管线如果输入输出都是完整未来轨迹,编码器能轻松造出富信息的隐表征;但本文的目标场景里推理时根本没有未来轨迹这种特权监督。SGDD 的做法是把这道鸿沟显式拆成师生两支:动力学编码器 \(E_{dyn}\) 在训练时观察真实未来序列 \(\{G_1,\dots,G_T\}\),产出捕捉了低频长期趋势与高频瞬时变化的 \(z_{dyn}\in\mathbb{R}^{N\times T\times d_z}\);初始编码器 \(E_{init}\) 只看 \(G_0\),但为了让它的输出落在和 \(z_{dyn}\) 同一个时空空间,作者把 \(G_0\) 的节点特征经全连接层从 \(\mathbb{R}^{N\times d}\) 扩展成 \(\mathbb{R}^{N\times T\times d}\) 构造一个"人造时空输入",得到 \(z_{init}\in\mathbb{R}^{N\times T\times d_z}\)。这样学生表征与教师表征维度、空间一致,才谈得上后续对齐。区别于把知识蒸馏仅当作模型压缩,这里蒸馏的本质是跨可观测性的信息迁移——把"能看未来"的能力压进"只能看现在"的学生里。
2. 时空联合基:在统一谱域里同时解开空间与时间频率
针对"现有方法时空分离、单维度谱建模"的痛点,本文构造一个时空联合基把表征投到谱域。设空间图与时间图归一化拉普拉斯的特征向量矩阵分别为 \(U_s\in\mathbb{R}^{N\times N}\)、\(U_t\in\mathbb{R}^{T\times T}\)(即 \(L_s=U_s\Lambda_s U_s^\top\)、\(L_t=U_t\Lambda_t U_t^\top\),特征值按升序排列,小值对应低频平滑分量、大值对应高频振荡分量),联合基由克罗内克积给出:
它把时空表征沿正交维度同时投影,从而在一个统一基里解耦空间与时间的频率分量,而不是各算各的。由于完整 \(NT\) 个基向量计算代价过高,作者只保留 \(K\) 个最小特征值对应的列,得到截断基 \(B_K=[b_1,\dots,b_K]\in\mathbb{R}^{NT\times K}\)。这一截断天然抑制了与大特征值绑定的高方差、高频内容,把"优先低频"的物理先验直接写进了基的选择里;\(K\) 是关键超参,控制能被自适应加权的频段范围。
3. 谱引导增强(SGE):用可学习的频率权重自适应放大关键频段
即便有了联合基,"哪段频率更重要"仍需自适应决定——这是 SGE 要解决的。对表征 \(z\)(reshape 成 \(\mathbb{R}^{d_z\times(NT)}\)),由于 \(B_K\) 正交,\(P=B_K B_K^\top\) 是到 \(\mathrm{span}(B_K)\) 的正交投影,于是可分解为 \(z = Pz + (I-P)z\):前者捕捉被选中的谱模式,后者是截断子空间之外的残差信息。增强过程先把 \(z\) 投到基上得谱系数 \(a:=B_K^\top z\in\mathbb{R}^{d_z\times K}\),再用一组可学习、频率特定的权重 \(w\in\mathbb{R}^K\) 调制并投回时空域:
最后把残差加回来重建:\(z^{sg}:=\tilde z + (I-P)z\)。这样 \(z^{sg}\) 既融合了被重加权的主导谱分量、又保住了截断之外的残差,得到一个频率感知的更丰富表征。可学习的 \(w\) 给了模型在谱域里直接、灵活地强调任务相关频段的能力——这正是"不同系统重要频率不同"这一观察落地的地方。重建前的谱系数 \(\tilde a_{dyn}\)、\(\tilde a_{init}\) 还会被单独留出,用于后面谱域的对齐。
4. 双层对齐 + 分阶段训练:在时空域和谱域同时蒸馏并稳住优化
简单地在位置层面模仿教师并不够,因为长时程稳定性既依赖低频全局趋势又依赖高频细节。本文因此施加双层对齐:时空域上对齐 \(z^{sg}_{dyn}\) 与 \(z^{sg}_{init}\)(\(\mathcal{L}_{rep}\)),谱域上对齐系数 \(\tilde a_{dyn}\) 与 \(\tilde a_{init}\)(\(\mathcal{L}_{spec}\)),合成对齐损失 \(\mathcal{L}_{align}=\mathcal{L}_{rep}(z^{sg}_{dyn},z^{sg}_{init})+\mathcal{L}_{spec}(\tilde a_{dyn},\tilde a_{init})\),总损失为
其中 \(\mathcal{L}_{pred}\) 是轨迹 MSE,\(\lambda\)(全实验取 1.0)加权对齐项;计算 \(\mathcal{L}_{align}\) 时对 \(z^{sg}_{dyn}\) 和 \(\tilde a_{dyn}\) 截断梯度,避免学生反过来污染教师。训练采用两阶段策略:预训练阶段 teacher forcing 比例设为 1.0、解码器只吃 \(z^{sg}_{dyn}\)(约占总轮数的 1/3);联合训练阶段比例降到 0.5,解码器在 \(z^{sg}_{dyn}\) 与 \(z^{sg}_{init}\) 间交替,且 \(E_{init}\) 此时不仅受对齐损失约束、还直接受轨迹预测损失的监督。这一渐进式安排保证了蒸馏的稳定收敛。
损失函数 / 训练策略¶
- 预测损失 \(\mathcal{L}_{pred}\):预测轨迹与真值的逐步 MSE。
- 对齐损失 \(\mathcal{L}_{align}=\mathcal{L}_{rep}+\mathcal{L}_{spec}\):表征级 MSE + 谱级 MSE,对教师侧梯度 detach。
- 权重 \(\lambda=1.0\);优化器 Adam。
- 两阶段:预训练(teacher forcing 1.0,约 1/3 轮数)→ 联合训练(teacher forcing 0.5,剩余轮数)。
- 骨干:\(E_{dyn}\) 用 STSGNN,\(E_{init}\) 用 GAT;解码器实例化为 EGNO 或 GF-NODE,对应 SGDD-EGNO 与 SGDD-GFNODE 两个变体。
实验关键数据¶
主实验¶
在分子动力学(MD17,约 10 个原子/分子)、人体运动捕捉(CMU Mocap,31 关节)、蛋白质(ADk 平衡轨迹,855 骨架节点)三类多尺度系统上评测,指标为 S2S(只看末步状态)与 S2T(整条轨迹平均),均为 MSE(\(\times 10^{-2}\),越低越好)。
| 数据集 / 设置 | 指标 | 代表样例 | 本文 SGDD | 之前 SOTA | 提升 |
|---|---|---|---|---|---|
| MD17 Benzene | S2S | EGNO 48.85 / GFNODE 4.82 | SGDD-GFNODE 2.74 | 4.82 | +43.2% |
| MD17 Aspirin | S2S | EGNO 9.18 | SGDD-GFNODE 7.29 | 7.93 | +8.1% |
| Mocap Run | S2S | EGNO 33.9 | SGDD-EGNO 28.2 | 33.9 | +16.8% |
| Mocap Walk | S2S | GFNODE 9.3 | SGDD-GFNODE 6.5 | 9.3 | +30.1% |
| Protein ADk | S2S | EGNO 2.23 | SGDD-EGNO 1.75 | 2.23 | +21.5% |
| Mocap Walk | S2T | EGNO 3.5 | SGDD-EGNO 3.2 | 3.5 | +8.6% |
在 MD17 的 S2S 评测上,SGDD 在全部 8 个分子上取得 SOTA;Benzene 上 SGDD-EGNO 相对 EGNO 提升约 72%,因为 EGNO/GFNODE 在该分子上的误差集中在低频段,而 SGDD 学到的表征能更好捕捉低频运动。蛋白质 ADk(855 节点的大图)上仍有 21.5% 提升,说明方法能扩展到大规模空间系统。
消融实验¶
在 SGDD-EGNO 上,针对谱域对齐(Freq Align)、时空域对齐(Feature Align)、谱引导增强(SGE) 三件套做消融(S2T,\(\times 10^{-2}\)):
| Freq Align | Feature Align | SGE | Ethanol | Toluene | Mocap-Walk | Mocap-Run |
|---|---|---|---|---|---|---|
| ✓ | ✓ | ✓ | 2.84 | 3.80 | 2.95 | 12.98 |
| ✓ | ✓ | - | 2.90 | 4.18 | 4.04 | 12.61 |
| ✓ | - | ✓ | 2.89 | 4.86 | 3.30 | 13.01 |
| - | ✓ | ✓ | 2.85 | 4.65 | 3.26 | 14.37 |
另有两组分析:截断参数 \(K\)(Mocap 共 \(31\times 5=155\) 个频率分量)对性能呈非单调影响——太小则只盯最低频、重要频段无法强调,增大可加权更多模式但不保证单调下降,需挑合适的 \(K\);编码器组合上,\(E_{dyn}\)=STSGNN + \(E_{init}\)=GAT 的组合在 MD17 上整体最优。
关键发现¶
- 双层对齐两支单独用时没有一致优劣,但同时用才最好:时空对齐给整体结构提供鲁棒归纳偏置,谱对齐细化频率优先级以抑制噪声与不稳定,二者互补。
- 去掉 SGE(对比表 4 第 1、2 行)在 Toluene、Mocap-Walk 等多数情形掉点明显,证明可学习谱权重确实能在对齐时放大有用频段。
- 提升在 S2T(整条轨迹)比 S2S(末步)更大,说明 SGDD 主要受益于更可靠的长时程表征,越往后预测优势越明显。
亮点与洞察¶
- 把"未来轨迹"当特权信息蒸馏:这是最巧妙的一招——训练时让教师合法地"看答案",再把这种能力压进只看初始状态的学生,绕开了推理无未来信息的硬约束,而非去硬猜未来。
- 克罗内克积时空联合基:用 \(B=U_t\otimes U_s\) 在一个统一谱域里同时解耦空间与时间频率,干净地修正了现有方法"时空分离、单维度建模"的结构缺陷,可迁移到任何"图 + 时序"的时空预测任务。
- 谱域 + 时空域双重对齐:蒸馏不止比位置,还比频率系数,把"低频管稳定、高频管细节"的物理直觉直接编码进损失,对长时程稳定性贡献显著。
- 解码器即插即用:SGDD 是表征学习框架,可包在 EGNO、GF-NODE 等不同物理动力学解码器外面成为通用增强,复用性强。
局限与展望¶
- 作者承认在 MD17 的 S2T 评测上部分分子落后于先前模型,归因为框架把这些模型当解码器、却无法完全复现其原报告结果;横向比较因此需谨慎(不同解码器实现差异会混入结果)。
- 截断参数 \(K\) 的选择对性能非单调敏感,缺乏一个先验最优准则,目前靠经验/逐数据集调;这在新系统上是额外调参成本。
- 训练需要完整未来轨迹作为特权监督,对只有稀疏/部分观测轨迹的真实场景适用性存疑。
- 时空联合基依赖固定的图拉普拉斯特征分解,对边集随时间变化、或超大规模图(特征分解代价 \(O((NT)^3\) 级)的系统可能是瓶颈,可探索近似/可学习谱基。
相关工作与启发¶
- vs EGNO / GF-NODE:它们用傅里叶时间卷积或图傅里叶 + Neural ODE 做频率感知,但只在时间或空间单维度建谱、且对各频率等权;本文在统一时空联合基上建谱并用可学习权重自适应重加权,且借特权蒸馏注入长程结构。本文实际把它们当解码器、在其外再套增强。
- vs 等变网络(EGNN / ClofNet / SE(3)-Transformer):它们强调空间等变对称性,对时间演化与多尺度频率纠缠着墨少;本文直接在时空域建模动力学并学频率感知表征。
- vs 知识蒸馏 + 特权信息(人体运动蒸馏未来位姿、learning-to-rank 的特权特征蒸馏):本文沿用"训练可见、测试不可见"的特权范式,但首次把它和时空谱域表征结合用于物理动力学长时程预测,蒸馏对象是频率感知的动力学表征而非软标签或原始位姿。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 时空联合谱基 + 特权信息蒸馏的组合在物理动力学预测中是新颖且自洽的切入点。
- 实验充分度: ⭐⭐⭐⭐ 覆盖分子/蛋白/人体三类多尺度系统并配三项消融,但部分 S2T 落后且依赖复现的 caveat 削弱了横向说服力。
- 写作质量: ⭐⭐⭐⭐ 动机—方法—公式链条清晰,谱域推导交代完整。
- 价值: ⭐⭐⭐⭐ 即插即用的表征增强框架,对长时程时空预测有较好迁移潜力。