通过结构修剪的紧凑 SO(3) 等变原子基础模型

SO(3) 等变图神经网络已成为原子基础模型的主导范例，通过将旋转对称直接构建到架构中来实现高精度和数据效率。然而，高阶张量运算的计算成本在模型精度和推理效率之间造成了艰难的权衡。在本文中，我们提出了一种 SO(3) 等变原子基础模型的结构剪枝方法，以弥补这种精度-效率差距。修剪沿着通道和顺序维度进行，每个不可约表示作为一个完整的块保留或删除，从而保留 SO(3) 等方差。从大的检查点开始，修剪后的模型大大降低了推理成本，同时比独立训练的小模型保持了更高的准确性。修剪后的 MACE-MP 模型在 Matbench Discovery 排行榜上的 9 个指标中的 7 个指标上优于官方从头开始训练的小模型。在效率方面，与从头开始训练小型模型相比，压缩的 MACE-MP 和 MACE-OFF 模型包含的参数少了 1.5$\times$ 到 4$\time$，并且需要的 预训练 计算量少了 2.5$\time$ 到 4$\time$。对于下游应用，与在八个代表性下游数据集上从头开始训练特定任务模型相比，微调 修剪模型将能量和力误差减少了 70.1% 和 34.4%。我们证明该方法可以推广到其他 SO(3) 等变架构（SevenNet、eSCN），并且可以与量化和 知识蒸馏 相结合以获得进一步的增益。