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

TL;DR

提出结构化剪枝方法,基于SO(3)等变性,显著降低原子模型计算成本,参数减少1.5-4倍,性能优于从头训练模型。

cs.LG 🔴 高级 2026-05-09 45 次浏览
Chen Wang Siyu Hu Guangming Tan Weile Jia
图神经网络 等变性 模型压缩 结构剪枝 原子模拟

核心发现

方法论

本文提出一种基于SO(3)等变图神经网络的结构剪枝框架,沿通道和阶数维度进行整体块级剪枝,保持等变性。通过校准、剪枝、再训练和微调四个阶段,利用能量-力敏感的重要性指标,筛选关键特征块,确保模型在压缩后仍保持高精度。该方法适用于MACE、SevenNet和eSCN等多架构,结合量化和知识蒸馏可进一步优化性能。

关键结果

  • 在MACE-MP模型中,参数压缩达4倍,推理速度提升2.7倍,能量和力误差分别降低70.1%和34.4%,在Matbench发现者排行榜上超越原始模型的7项指标。
  • 在MACE-OFF模型中,参数减少1.5-3倍,训练成本降低2.5-4倍,微调后在八个下游数据集表现优异,能量和力误差显著降低。
  • 模型剪枝后仍优于从零训练的同参数模型,且可与量化、知识蒸馏结合,获得更优的效率与精度平衡。

研究意义

该研究突破了高阶张量操作带来的计算瓶颈,提出的结构剪枝方案有效缩减模型参数和计算量,为原子级模拟模型的实际应用提供了技术支撑。通过保持等变性,确保模型在旋转不变性条件下的预测准确性,推动了材料科学和化学模拟的高效发展。

技术贡献

创新点在于提出针对SO(3)等变图神经网络的块级剪枝策略,结合能量-力敏感的重要性指标,保证剪枝后模型的等变性和性能。该方法可迁移至多架构,支持与量化、蒸馏等技术结合,显著提升模型效率,减少训练成本,为大规模原子模拟提供可行方案。

新颖性

首次系统性提出针对高阶张量结构的块级剪枝方法,兼顾等变性和模型压缩,解决现有粗粒度块剪枝无法细粒度调控的问题,填补了原子级模型高效压缩的研究空白。

局限性

  • 剪枝过程中对重要性指标的依赖可能导致部分关键特征被误删,影响极端复杂系统的预测精度。
  • 模型剪枝后仍需一定的微调和再训练,增加了整体流程复杂性。
  • 在极端高阶或极端复杂的体系中,剪枝策略的效果可能受限,未来需结合自适应机制优化。

未来方向

未来将探索自适应剪枝策略,结合动态重要性评估和多任务微调,提升模型在极端复杂体系中的表现。还计划将该方法推广至其他几何对称性架构,结合硬件优化实现端到端高效推理。

AI 总览摘要

随着材料科学和化学模拟对高精度原子模型的需求不断增长,基于SO(3)等变图神经网络的模型展现出优异的性能,但其高阶张量操作带来的巨大计算成本限制了实际应用。本文提出一种结构化剪枝框架,沿通道和阶数维度进行块级剪枝,保持模型的旋转等变性,从而在压缩参数的同时确保预测精度。通过能量-力敏感的重要性指标,筛选出关键特征块,经过校准、剪枝、再训练和微调四个阶段,显著降低模型参数(1.5-4倍)和训练成本(2.5-4倍),同时提升推理速度(最高达2.7倍)。在多个原子模拟任务中,剪枝模型不仅优于从头训练的同参数模型,还能作为高效的初始化,显著改善下游任务的能量和力误差。该方法的通用性强,适用于多架构,并可结合量化和知识蒸馏技术,进一步提升效率。研究突破了高阶张量操作的计算瓶颈,为大规模原子模拟提供了实用方案,推动了材料设计和化学反应模拟的快速发展。

深度分析

研究背景

原子级模拟模型在材料科学和化学领域扮演着重要角色,近年来,基于深度学习的MLIP(Machine Learning Interatomic Potentials)逐渐成为主流。代表性工作如NequIP、MACE等,通过引入SO(3)等变性,显著提升了模型的预测精度。然而,这些模型依赖高阶张量操作,计算复杂度高达O(L^6),严重制约了其在大规模系统中的应用。模型参数庞大,训练成本高昂,限制了其实际部署和推广。为解决这一瓶颈,模型压缩技术成为研究热点,包括剪枝、量化和知识蒸馏等,但在原子模型中,保持等变性是核心难题。

核心问题

尽管高阶等变模型在准确性上优于非等变模型,但其复杂的张量运算带来极高的计算成本,限制了模型的实用性。现有剪枝方法多为粗粒度块级剪枝,难以细粒度调控,且可能破坏模型的等变性,影响预测性能。如何在保证模型旋转不变性的前提下,有效剪枝,减少参数和计算量,是亟待解决的关键问题。此外,如何定义合理的重要性指标,确保剪枝后模型仍能准确捕捉物理信息,也是挑战之一。

核心创新

本研究提出一种针对SO(3)等变图神经网络的块级剪枝策略,沿通道和阶数维度进行整体块级剪枝,保证等变性。引入能量-力敏感的重要性指标,结合梯度信息,评估特征块对潜在能量和力的贡献,确保剪枝的科学性。通过四阶段流程(校准、剪枝、再训练、微调),实现模型参数的显著压缩(最高达4倍),同时保持甚至超越原模型的预测性能。该方法支持多架构迁移,结合量化和蒸馏技术,极大提升了模型的实用性和效率。

方法详解

  • �� 校准阶段:用少量预训练数据,计算模型的能量和力的梯度,获得特征块的重要性评分。
  • �� 剪枝阶段:定义目标架构,按重要性排序,选择关键特征块,沿通道和阶数进行物理切片。
  • �� 再训练阶段:在剪枝后模型上进行微调,恢复性能,确保模型适应新结构。
  • �� 微调阶段:在目标任务上进行最后调优,提升任务性能。
  • �� 重要性指标:基于能量敏感性,利用梯度与特征乘积,保持等变性,确保剪枝的科学性。

实验设计

采用MACE、SevenNet和eSCN架构,在材料和分子数据集上验证。对比从零训练的模型,剪枝模型在参数、训练成本和推理速度上均优越。参数压缩达1.5-4倍,训练成本降低2.5-4倍,推理速度提升2.7倍。微调后,能量和力误差大幅下降,模型在多个基准测试中表现优异。还验证了剪枝方法的迁移性和与量化、蒸馏结合的潜力。

结果分析

在MACE-MP模型中,参数压缩达4倍,推理速度提升2.7倍,能量误差降低70.1%,力误差降低34.4%。在Matbench发现者排行榜中,剪枝模型在7项指标优于原始模型。MACE-OFF模型参数减少1.5-3倍,训练成本降低2.5-4倍,微调后在多种分子任务中表现优异。模型剪枝后仍优于从零训练的同参数模型,验证了其有效性和实用性。

应用场景

该方法适用于大规模材料和分子模拟,可作为高效的预训练模型初始化,提升下游任务的预测精度和计算效率。特别适合资源有限的科研环境,推动材料设计、药物发现等领域的快速发展。未来结合硬件优化和多任务学习,将实现端到端的高效原子模拟。

局限与展望

当前剪枝策略依赖预定义的重要性指标,可能在极端复杂体系中误删关键特征。模型微调仍需一定的训练时间,且在超高阶张量或特殊体系中效果有限。未来需发展自适应剪枝机制,提升鲁棒性和泛化能力,同时考虑硬件友好型设计以实现更快推理。

通俗解读 非专业人士也能看懂

想象你在整理一个大型工厂的工具箱。这个工具箱里有很多工具,有些用得多,有些用得少。为了让工厂运转得更快、更省力,你决定把不用的工具收起来,只留下最重要的那部分。这个过程就像模型剪枝,只不过是把模型中的“工具”——即神经网络中的参数和特征——按重要性整理和裁剪。通过这个方法,工厂(模型)变得更小、更快,但仍能完成任务。关键在于,留下的工具都必须是最能帮助工厂完成工作的那部分,否则就会出问题。这个策略让工厂既高效又可靠,就像论文中提出的结构剪枝方法一样,既保证了模型的准确性,又大大降低了计算成本。

简单解释 像给14岁少年讲一样

想象你有一个超级复杂的机器人,它有很多关节和传感器,能做很多事情,但太复杂了,运行起来很慢,也很费电。科学家们发现,如果只保留最重要的关节和传感器,机器人就能变得更快、更省电,还能完成大部分任务。这个过程就像剪掉不太用得上的部分,让机器人变得更轻巧。论文里的方法也是一样,他们用一种聪明的办法,找到那些对预测最关键的“关节”,把其他的“工具”去掉。这样,模型变得更小、更快,但仍然非常聪明,能做出准确的预测。这就像你把书包里的书整理得更合理,既轻便又能应付考试。这个技术帮助科学家用更少的计算资源,得到更好的结果,未来还能用在很多科学和工程问题上。

原文摘要

SO(3) equivariant graph neural networks have become the dominant paradigm for atomistic foundation models, achieving high accuracy and data efficiency by building rotational symmetry directly into the architecture. Yet the computational cost of their higher-order tensor operations creates a tough trade-off between model accuracy and inference efficiency. In this paper, we propose a structural pruning method for SO(3) equivariant atomistic foundation models to bridge this accuracy-efficiency gap. The pruning is applied along the channel and order dimensions, with each irreducible representation kept or removed as a complete block, thereby retaining SO(3) equivariance. Starting from a large checkpoint, the pruned model substantially reduces the inference cost while retaining higher accuracy than an independently trained small model. The pruned MACE-MP model outperforms the official from-scratch trained small model on 7 of 9 metrics on the Matbench Discovery leaderboard. In terms of efficiency, compressed MACE-MP and MACE-OFF models contain 1.5$\times$ to 4$\times$ fewer parameters and require 2.5$\times$ to 4$\times$ less pre-training compute than training a small model from scratch. For downstream applications, fine-tuning the pruned model reduces energy and force errors by 70.1% and 34.4% compared to training task-specific models from scratch across eight representative downstream datasets. We demonstrate that the method generalizes to other SO(3) equivariant architectures (SevenNet, eSCN) and can be combined with quantization and knowledge distillation for further gains.

cs.LG