🤖 machine learning

Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning

本文提出了一种结构剪枝方法，通过移除不可约表示块，在原子基础模型中保持SO(3)等变性，从而显著降低参数量和计算成本，同时在精度和下游微调效率方面均优于从头训练的小规模模型。

原作者： Chen Wang, Siyu Hu, Guangming Tan, Weile Jia

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Chen Wang, Siyu Hu, Guangming Tan, Weile Jia

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你有一位能完美烹饪世间任何菜肴的大厨。这位大厨才华横溢，但体型庞大：他背着一个装满数千种香料、巨型锅具和复杂工具的背包。虽然他能做出完美的佳肴，但烹饪耗时很长，能耗巨大，且需要一间巨大的厨房来运作。

现在，想象你想开设一个小型快餐摊。你并不需要这位大厨那整套庞大的工具箱；你只需要掌握快速制作美味汉堡和薯条的核心技能。

本文介绍了一种新方法，旨在将那位“大厨”（一种用于预测原子行为的大型高精度人工智能模型）通过手术式修剪，转化为“快餐厨师”（一种更小、更快的模型），同时不丧失其烹饪美味佳肴的能力。

以下是他们如何实现这一目标的分解说明，使用了简单的类比：

1. 问题：“沉重的背包”

当前化学领域最佳的人工智能模型（称为SO(3) 等变模型）就像那些完美理解旋转物理原理的大厨。如果你旋转一个分子，该模型能确切知道原子之间如何相对移动。这使得它们极其精准。

然而，为了实现这一点，它们背负着由复杂数学（高阶张量）构成的“沉重背包”。

权衡取舍：背包越大，烹饪越精准，但运行速度越慢，成本越高。
旧方法：如果你想要一位更小的大厨，通常必须从头开始训练一位全新的小厨师。这位新厨师毫无经验，学习时间长，且最终做出的食物往往不如那位大厨美味。

2. 解决方案：“结构剪枝”（手术式修剪）

作者没有从头开始训练新厨师，而是对现有的大厨进行了“结构剪枝”。这可以想象为一次非常精细的理发，或者是对背包的修剪。

挑战：你不能随意剪掉一个香料罐或工具。在这些模型中，“工具”（数学部分）紧密相连。如果你切断工具的某一部分，整个工具就会损坏，模型也会失去理解旋转的能力（不再具备“等变性”）。
创新：作者发现如何一次性剪掉整个工具块。他们将一组特定的关联工具视为一个单一的“原子单元”。如果决定移除一个单元，就将其整体移除，从而确保剩余的工具仍能完美协同工作。

3. 如何决定剪掉什么（“敏感性测试”）

你如何知道该保留哪些工具，该丢弃哪些？不能仅凭猜测。

作者使用了一种基于能量和力的巧妙测试：

想象模型正拿着一件精致的玻璃雕塑（分子）。
他们问道：“如果我移除这个特定工具，雕塑会破裂或倒塌吗？”
他们测量了当移除一个工具时，模型对“能量”和“力”的预测发生了多大变化。
规则：如果移除一个工具几乎不改变结果，那它就是一个“懒惰工具”，会被剪掉。如果移除它会导致模型出错，那它就是一个“关键工具”，会被保留。

4. 过程：四步食谱

本文概述了将巨型模型转化为紧凑模型的四个步骤：

校准：让大厨做几道测试菜，观察哪些工具实际上被使用了，以及它们有多重要。
剪枝：根据测试结果，手术式地移除“懒惰”的工具块。
再训练：模型现在变小了，且留下了工具被移除后的空缺。他们利用少量数据给它上一堂快速的“复习课”，帮助它适应新的、更小的身体。
微调：最后，他们教导这个新的紧凑模型执行特定任务（例如预测特定药物分子的行为）。

5. 结果：更快、更便宜，且依然美味

结果令人印象深刻：

优于从头开始：经过修剪的模型比从头开始训练的全新小模型更精准，尽管它们的规模相同。
巨大节省：
- 训练：与从头训练一个小模型相比，创建修剪后模型所需的计算机时间和成本减少了2.5 到 4 倍。
- 速度：在使用模型进行预测时，速度快了2.7 倍。
- 内存：它所需的计算机内存显著减少（最多减少 5.7 倍）。
通用性：这种方法不仅适用于一种类型的模型（MACE），也适用于其他模型（SevenNet, eSCN），证明它是这类人工智能大厨的通用食谱。

6. “额外”组合

本文还指出，这种修剪方法可以与其他效率技巧相结合：

量化：就像从高清视频切换到标清视频以节省空间。
知识蒸馏：就像让大厨教小厨师特定的技巧。
当这些方法结合使用时，模型会变得更小、更快，且不会损失质量。

总结

简而言之，作者找到了一种方法，通过仔细剪除不必要的部分，同时保留核心的“物理大脑”，从而缩小用于化学领域的庞大且昂贵的人工智能模型。其结果是一个更小、更快、更便宜的模型，其智能程度甚至超过任何从头构建的小模型。这使得先进材料发现和药物设计能够惠及更多没有超级计算机的研究人员。

技术摘要：通过结构剪枝实现紧凑的 SO(3) 等变原子基础模型

问题陈述

SO(3) 等变图神经网络（GNN）已成为原子基础模型的主导范式，通过显式编码旋转对称性实现了高精度和高数据效率。然而，这些模型依赖高阶张量积运算来耦合几何特征，导致计算成本高达 $O(L_{max}^6)$ ，其中 $L_{max}$ 是不可约表示的最大阶数。这在模型精度与推理效率之间造成了显著的权衡。

现有的模型压缩策略——特别是知识蒸馏（KD）和量化——在该领域存在局限性。KD 通常需要从零开始训练学生模型，无法直接复用教师模型的权重；而量化虽然降低了比特精度，但参数数量和张量积路径保持不变。此外，为大语言模型（LLM）开发的标准结构剪枝方法无法直接应用于 SO(3) 等变模型。简单地移除特征维度会破坏张量积路径的旋转对称性，从而违反物理一致性所必需的基本等变性。此前对原子模型进行剪枝的尝试仅限于粗粒度的块移除（例如整个相互作用层），这会减小模型的感受野，并对浅层架构尤为不利。

方法论

作者提出了一种专为 SO(3) 等变原子基础模型设计的结构剪枝框架。该方法在预训练检查点的基础上运行，生成更小的稠密模型，同时保留 SO(3) 等变性和高精度。

1. 等变剪枝粒度

核心创新在于剪枝单元的定义。在 SO(3) 等变模型中，特征被分解为阶数 $l \in [0, L_{max}]$ 的不可约表示。特定通道 $k$ 和阶数 $l$ 的特征向量由 $2l+1$ 个分量组成，索引为 $m \in [-l, l]$ 。

约束：SO(3) 中的旋转操作耦合了特定阶数 $l$ 内的所有 $m$ 分量。移除单个分量（逐元素剪枝）会破坏等变性。
解决方案：该方法将特定阶数 $l$ 和通道 $k$ 的 $2l+1$ 个分量所张成的子空间视为原子单元。剪枝操作针对元组 $(k, l)$ 进行，由二进制掩码 $z^{(t)}_{k,l}$ 表示。这确保了如果保留某个特征块，其所有旋转分量均被保留，从而严格维持 SO(3) 等变性。

2. 物理信息重要性准则

为了确定保留哪些 $(k, l)$ 块，作者提出了一种基于势能面敏感度的重要性评分。

双重目标：机器学习势函数（MLIPs）必须同时预测势能（ $E$ ）和力（ $F_i = -\nabla_{r_i} E$ ）。
指标：重要性评分 $I^{(t)}_{k,l}$ 源自移除特征块后能量变化的一阶泰勒展开。它聚合了所有原子和表示维度上特征幅度与能量对该特征梯度的乘积：
$I^{(t)}_{k,l} = \mathbb{E}_{R,Z \sim \mathcal{D}} \left| \sum_{j,m} \frac{\partial E}{\partial h^{(t)}_{j,klm}} \cdot h^{(t)}_{j,klm} \right|$
优势：该准则具有 SO(3) 不变性，且开销可忽略不计，因为梯度 $\frac{\partial E}{\partial h}$ 在标准力推理过程中已经计算得出。

3. 结构对齐与流程

该框架遵循四阶段流程：

校准：将一小部分预训练数据输入大模型以计算重要性评分。
目标驱动剪枝：定义目标架构。根据重要性评分保留排名靠前的 $(k, l)$ 块。
结构对齐：物理上改变模型结构以匹配缩减后的维度。这包括：
- 张量积重构：切片径向权重和克莱布希 - 高登（Clebsch-Gordan）路径，以与保留的通道对齐。
- 权重切片：物理上移除线性变换权重中的行和列，以创建稠密的较小权重张量。
重训练与微调：剪枝后的模型在预训练数据的子集上进行轻量级重训练，以恢复近似误差，随后在特定下游任务上进行微调。

主要贡献

等变剪枝方法：提出了一种针对高阶 SO(3) 等变张量的结构剪枝方法，涵盖通道和阶数维度，通过将 $(k, l)$ 块视为原子单元，严格保持 SO(3) 等变性。
四阶段框架：一个用于压缩基础模型的完整流程（校准、剪枝、重训练、微调）。
通用性：证明了该方法适用于多种架构，包括 MACE、SevenNet 和 eSCN。
可组合性：验证了结构剪枝可以与量化和知识蒸馏相结合，以实现复合的效率提升。

实验结果

该方法主要在 MACE-MP（无机材料）和 MACE-OFF（有机分子）上进行了评估，并在 SevenNet 和 eSCN 上进行了通用性测试。

精度与从零训练对比

优越性：剪枝模型始终优于同等大小从零训练的模型。
- 在 MACE-MP 上，与从零训练的基线相比，剪枝模型在 sAlex 数据集上实现了更低的力平均绝对误差（MAE）和更好的分布外性能。
- 在 Matbench Discovery 上，一个仅使用 25% 预训练数据的剪枝 MACE-MP 模型，在 9 项指标中的 7 项上优于官方从零训练的小型模型。
- 在 MACE-OFF 上，即使在仅使用 50% 预训练数据的情况下，剪枝模型在多样化的分子子集上也达到或超过了从零训练的基线。
下游性能：在八个无机下游数据集上对剪枝模型进行微调，与从零训练特定任务模型相比，能量和力误差分别降低了 70.1% 和 34.4%。

效率提升

参数减少：压缩后的模型参数量减少了 1.5 倍至 4 倍。
训练成本：完整流程所需的预训练计算量比从零训练一个小模型少 2.5 倍至 4 倍（例如，特定配置下为 181 GPU 小时 vs. 750 GPU 小时）。
推理速度：剪枝模型实现了 2.7 倍至 3.0 倍 的推理加速，峰值内存使用量最多减少 5.7 倍。

与其他方法的结合

量化：将剪枝与 FP32 量化相结合，实现了 9.6 倍 的吞吐量提升（相对于基础 FP64 模型）。
知识蒸馏：将剪枝与基于海森矩阵（Hessian）的 KD 相结合，进一步细化了模型性能，证明了这两种方法是互补的。

意义与主张

本文主张，结构剪枝为弥合原子基础模型中精度与效率之间的差距提供了一条自然且有效的途径。通过直接从大型预训练检查点继承权重，并在保持物理对称性的同时对其进行压缩，该方法生成的紧凑模型具有以下特点：

比同等大小从零训练的模型更准确。
训练和部署成本显著更低。
物理一致，保持了可靠分子动力学模拟所需的严格 SO(3) 等变性（通过对液态水的分子动力学模拟验证，显示了稳定的能量漂移和准确的径向分布函数）。

作者强调，核心原则——评估耦合不可约表示的重要性——广泛适用于任何将特征表示为 SO(3) 不可约表示之和的等变架构。这项工作表明，剪枝使从业者能够在资源受限的环境中部署高保真基础模型，而不会牺牲在大规模预训练期间学到的可迁移性。