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

会议: NeurIPS 2026
arXiv: 2605.08885
领域: 模型压缩
关键词: 原子基座模型, SO(3)等变图神经网络, 结构化剪枝, 机器学习原子间势函数, 张量积重构

一句话总结¶

针对 SO(3) 等变原子基座模型高阶张量计算开销极大的瓶颈，本文提出一种按不可约表示完整块剪枝并结合能量-力敏感度准则的结构化剪枝与微调框架，在严格保持物理旋转等变性的同时实现最高 4 倍参数压缩与 2.7 倍推理加速，精度超越从头训练的官方同尺寸小模型。

研究背景与动机¶

机器学习原子间势函数（MLIPs）近年来已发展成为计算化学与材料科学的基础范式，以 MACE、SevenNet 和 eSCN 为代表的原子基座模型通过在大规模量子化学和材料数据集（如 MPtrj、SPICE）上预训练，展现出跨元素周期表的强大泛化能力。为了确保物理规律的一致性与极高的数据效率，当前前沿模型广泛采用 SO(3) 等变图神经网络架构，直接在特征空间显式编码三维旋转对称性。然而，这类等变性通常依赖于高阶几何不可约表示（irreducible representations）之间的 Clebsch-Gordan 张量积运算，其理论计算复杂度随着最大阶数 $L_{\max}$ 呈 $\mathcal{O}(L_{\max}^6)$ 急剧爆炸。高阶张量计算的高昂代价导致大模型虽然精度优异，但推理延迟过高且显存开销巨大，在大尺度分子动力学（MD）长时间模拟与高通量晶体筛选中面临严峻的部署瓶颈。

现有的模型压缩方案在处理等变几何模型时存在显著局限。知识蒸馏依赖于重新训练全新架构的学生模型且无法直接复用教师模型的预训练权重，量化仅降低数值精度而不减少张量积路径与参数量。最关键的是，直接将大语言模型（LLM）中成熟的结构化剪枝方法照搬到等变模型完全行不通：普通神经网络处理标量特征，而 SO(3) 等变模型的操作对象是高阶不可约表示张量；若简单按单个特征维度进行元素级通道裁剪，将破坏 Wigner-D 矩阵在分量间混合的群表示变换法则，直接摧毁模型的等变对称性。此外，现有针对原子模型的剪枝工作大多停留在粗粒度的“层级剪枝”（Depth Pruning），然而在消息传递机制中，模型感受野半径严格正比于网络层数（$R_{\text{rec}} = T \times r_c$），裁剪层数会直接截断原子局域感知的物理范围，在浅层模型上造成严重的灾难性精度退化。

面对如何在保持三维旋转等变性与物理感受野的前提下对大基座模型进行高精度压缩的核心矛盾，本文提出了沿通道维度和阶数维度的细粒度结构化剪枝范式。核心 idea：将阶数 $l$ 的 $2l+1$ 个基底分量作为一个不可分割的原子几何块 $(k,l)$ 进行剪枝以严格保等变性，结合一阶能量-力敏感度重要性评分与张量积切片重构，实现紧凑等变基座模型的高效继承与下游适配。

方法详解¶

整体框架¶

本文构建了一个端到端的四阶段结构化剪枝与下游迁移流水线：校准（Calibration）、目标驱动剪枝与结构对齐（Target-Driven Pruning & Structural Alignment）、轻量再训练（Lightweight Retraining）以及下游微调（Fine-tuning）。系统输入预训练好的大规模等变基座模型检查点与少量无标注结构样本，首先利用前向势能与反向原子受力的中间梯度评估各几何特征块的物理重要性；随后依据预定义的目标紧凑架构对几何块进行筛选，物理切除冗余参数并重构 Clebsch-Gordan 张量积耦合路径；再经过短期预训练数据重校准恢复截断误差，最终生成兼具高精度与高吞吐的专用下游势函数。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["预训练等变基座模型<br/>大尺寸检查点 + 校准集"] --> B["1. 等变保对称剪枝粒度<br/>划分 (k,l) 不可约表示完整块"]
    B --> C["2. 能量-力敏感度物理重要性评估<br/>一阶泰勒展开能量梯度乘特征"]
    C --> D["3. 物理切片与张量积重构<br/>重构 CG 路径与密集参数切片"]
    D --> E["4. 轻量重适应再训练<br/>利用小比例预训练数据恢复精度"]
    E --> F["下游紧凑专用模型<br/>微调适配材料/分子动力学模拟"]

关键设计¶

1. 等变保对称剪枝粒度：以不可约表示完整块为最小剪枝单元消除等变破缺

常规通道剪枝独立裁剪矩阵中的某一维，但 SO(3) 等变模型中第 $t$ 层原子 $j$ 的通道 $k$、阶数 $l$ 特征包含 $2l+1$ 个分量 $h_{j,klm}^{(t)}$（其中 $m \in [-l, l]$）。三维空间旋转群作用通过稠密的 Wigner-D 矩阵 $D^l(g)$ 将这 $2l+1$ 个分量紧密耦合。如果采用任意对角掩码对各分量独立赋权，由于 $D^l(g)$ 具有非零非对角项，掩码操作与旋转变换将无法交换，导致等变性被彻底破坏。为此，本文将属于同一通道 $k$ 与阶数 $l$ 的整个 $2l+1$ 维子空间视为最小不可分割的剪枝单元，引入二值块掩码 $z_{k,l}^{(t)} \in \{0, 1\}$： $$\hat{h}_{j,klm}^{(t)} = z_{k,l}^{(t)} \cdot h_{j,klm}^{(t)}$$ 因为标量因子与线性旋转群矩阵乘法自然对易，该块级剪枝严格保证了等变变换不变性。同时从球面调和函数谱分析角度，阶数 $l$ 对应球面拉普拉斯算子特征值 $\lambda_l = l(l+1)$，表征原子局域环境的角度高频分量。该粒度允许模型保留特定通道的低阶标量（$l=0$）而剔除极耗费算力的高阶几何张量（$l>0$），等价于在角度空间进行自适应低通滤波，保留了粗粒度能量主导结构。

2. 能量-力敏感度物理重要性评估：免额外反向传播的物理信息导向评分

势函数模型不仅要预测系统的总标量势能 $E$，更核心的是预测原子守恒受力 $\mathbf{F}_i = -\nabla_{\mathbf{r}_i} E$。因此纯激活值幅度或单纯权重的启发式评分无法反映特征对势能面的物理敏感度。本文基于势能面的一阶泰勒展开，推导特征块被置零所引起的能量变化量 $\Delta E \approx -\sum_{j,m} \frac{\partial E}{\partial h_{j,klm}^{(t)}} \cdot h_{j,klm}^{(t)}$。由于求解原子受力本身就必须将能量梯度通过计算图反向传播至原子坐标，中间特征梯度 $\frac{\partial E}{\partial h}$ 在标准受力推理过程中天然存在，这一评估准则几乎无需额外的计算开销。在校准集 $\mathcal{D}$ 上对所有原子 $j$ 和表示分量 $m$ 聚合求绝对值期望，得到最终物理重要性得分： $$I_{k,l}^{(t)} = \mathbb{E}_{\mathbf{R},\mathbf{Z}\sim\mathcal{D}} \left| \sum_{j,m} \frac{\partial E}{\partial h_{j,klm}^{(t)}} \cdot h_{j,klm}^{(t)} \right|$$ 数学证明显示，由于势能 $E$ 具有旋转不变性且 Wigner-D 矩阵为酉矩阵，该重要性得分在三维坐标系整体旋转下保持严格不变，确保了剪枝决策不受分子构型空间朝向的影响。

3. 物理切片与张量积重构：稠密紧凑结构重组与元素专用压缩

为了使剪枝直接转化为硬件加速，模型不能依赖稀疏掩码，而必须在内存布局上转换为物理紧凑的小架构。对于等变线性层权重 $W_{k\tilde{k}l}^{(t)}$，直接依据保留通道的索引物理切除无用行与列；对于核心的张量积模块，消息聚合依赖径向权重 $R_{kl_1l_2l_3}^{(t)}$ 与 Clebsch-Gordan 路径，当输入阶数 $l_2$ 的通道数缩减后，重新索引并物理裁剪对应的径向基投影参数与多体收缩（GCG）耦合张量，抹除无效分支计算图。进一步地，针对残差层中与元素类别绑定的线性投影参数 $W_{z_i, k\tilde{k}l}^{(t)}$，系统在切片通道的同时，能够根据下游应用体系的具体化学组成，直接物理剔除目标数据集中不存在的化学元素对应的权重参数，实现化学空间与几何空间的联合紧凑化。

4. 轻量重适应再训练：小算力预算下的精度恢复与下游迁移

粗暴剪枝后的模型权重存在一定的泰勒截断误差。与从头预训练大模型不同，本文提出仅利用小比例（如 9% 到 25%）的预训练数据，在相同训练超参数下进行短期微调再适应。实测表明，再训练后保留通道的重要性排序与剪枝前呈现高达 0.81 的 Spearman 秩相关系数，且前 12.5% 最重要通道在再训练前后的重合率达到 87.5%，证明核心结构路径在轻量重训练后保持稳定。经过这一短暂过渡后，紧凑基座模型即可作为下游微调的优质初始权重，避免下游任务在冷启动时因特征严重破坏而陷入局部极小。

实验关键数据¶

主实验¶

主实验在无机材料基座模型 MACE-MP（基于 MPtrj 数据集预训练）上开展，评估模型压缩至不同尺寸后的能量（E, meV/atom）、受力（F, meV/Å）、应力（S, meV/Å³）平均绝对误差（MAE），以及材料科学大规模热力学稳定性基准 Matbench Discovery 的 25 万个候选晶体筛选表现。

模型与方法	架构配置	参数量	预训练数据量	能量 MAE (E)	受力 MAE (F)	应力 MAE (S)	sAlex 能量 MAE
From Scratch (大模型)	Layer2 L1 C256	15,847,440	100%	/	/	/	81.7
From Scratch (小模型基线)	Layer2 L0 C128	3,847,696	100%	29.6	52.3	1.8	78.0
L 剪枝 (本文)	Layer2 L0 C128	3,847,696	25%	30.9	50.8	1.9	74.1
L 剪枝 (本文)	Layer2 L0 C128	3,847,696	100%	25.5	47.3	1.7	71.9
L & Channel 剪枝 (本文)	Layer2 L0 C128	3,847,696	40%	18.9	51.7	1.9	67.2

在极具挑战的真实下游微调任务中（涵盖固态电解质 SSE、液态水 H2O、各类纯金属及合金共 8 个无机体系），本文剪枝初始化的紧凑模型相比从头训练展现出显著优势。

下游体系微调初始化	SSE 能量/受力 RMSE	H2O 能量/受力 RMSE	AgAu 能量/受力 RMSE	AlMgCu 能量/受力 RMSE	平均误差对比
From Scratch (从头训练)	1.8 / 29.9	79.9 / 29.7	369.1 / 34.5	7.7 / 42.9	基线 (100%)
Pre-training (100% 数据大基座)	0.6 / 29.1	1.1 / 36.8	15.4 / 15.8	3.4 / 12.2	强迁移上限
Pruning (25% 数据紧凑基座，本文)	0.5 / 25.9	0.9 / 31.6	6.6 / 15.1	3.2 / 11.6	能量误差降 70.1%，受力降 34.4%

消融实验¶

重要性准则对比消融（在 MACE-MP 上从 Layer2 L1 C256 剪枝至 Layer2 L1 C128，9% MPtrj 再训练）：

重要性准则配置	能量 MAE (meV/atom)	受力 MAE (meV/Å)	应力 MAE (meV/Å³)	说明
Random（随机剪枝）	42.7	68.6	2.4	无物理准则指导
Magnitude（权重幅度剪枝）	47.7	68.9	2.5	无法感知等变多体相关性，误差最高
Activation（纯激活值幅度）	36.4	64.4	2.3	捕捉特征尺度但缺少力场梯度敏感度
本文能量-力敏感度准则	35.9	64.2	2.3	能量泰勒一阶展开与受力反向传播一致，综合精度最优

剪枝与量化正交组合加速实验（NVIDIA RTX 3090，吞吐量单位：atom/s）：

模型配置	FP64 吞吐量	FP32 吞吐量	相对基线整体加速比	说明
Base (Layer2 L2 C128)	2,728	5,576	1.0× / 2.0×	基础未压缩模型
Pruned (Layer2 L1 C128)	4,976	12,268	1.8× / 4.5×	结构剪枝剔除 $L=2$ 阶张量积
Pruned (Layer2 L0 C128)	8,271	26,144	3.0× / 9.6×	剪枝与低精度量化复合，吞吐量提升近 10 倍

关键发现¶

宽度/阶数剪枝完胜深度剪枝：实验表明，裁剪层数的 Depth Pruning（从 2 层减到 1 层）导致能量 MAE 暴增 50%（升至 44.4 meV/atom），因为网络感受野半径与层数线性锁定（$R_{\text{rec}} = T \cdot r_c$），削减深度直接摧毁了长程相互作用建模能力；而宽度与高阶张量剪枝保留了完整感受野，仅平滑局域高频角度细节。
低算力反超满血从头训练：使用本文剪枝方案并仅配以 25%–40% 预训练数据再训练的紧凑模型，在受力精度和 OOD（分布外）测试集 sAlex 上全线超越使用 100% 全量数据从头训练的官方同构模型，端到端 GPU 耗时从 785 小时锐减至 216 小时（节省约 4 倍算力）。
复合压缩具有正交乘数效应：结构化剪枝（削减 FLOPs 与参数）、低精度量化（削减位宽）与 Hessian 知识蒸馏（软标签监督）作用于完全不同的优化维度，三者可无缝组合，在 RTX 3090 上实现最高 9.6 倍极限吞吐加速。

亮点与洞察¶

将群表示理论与网络剪枝精准契合：敏锐指出 SO(3) 群旋转作用混叠不可约表示各子分量的数学本质，将最小剪枝单元约束在完整块，从根本上解决了高阶几何等变张量被非均匀裁剪导致物理对称性破缺的难题。
巧妙复用保守力场的反向传播计算图：利用分子力场中原子受力由势能梯度自然推导（$\mathbf{F}_i = -\nabla E$）的物理特性，使泰勒展开所需特征敏感度项成为前向受力推理的副产物，实现了几乎零额外计算开销的全局物理评分。
化学元素维度的精准瘦身：结合真实材料模拟体系元素单一的实际场景，在重构等变线性投影层时顺带物理清除目标体系不存在的化学元素参数，为跨体系大模型向特定材料/药物场景落地提供了精巧的定制范式。

局限与展望¶

目标压缩架构仍需人工预先指定：当前方法依赖使用者手动定义剪枝后的目标通道数和目标最大阶数 $L_{\max}$，未能结合神经架构搜索（NAS）根据具体算力约束实现跨层、跨阶数的自适应参数分配。
对非线性激活依赖极深的多体模型适配细节复杂：如在 NequIP/SevenNet 中门控非线性机制对标量辅助特征产生硬依赖，需额外引入强制保留约束，未来需要更通用的等变依赖图分析器来自动化处理复杂依赖。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统性解决高阶 SO(3) 等变几何张量的保对称结构化剪枝难题。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖无机材料、有机分子、大尺度 MD 动力学模拟及三类不同等变架构，验证极其扎实。
写作质量: ⭐⭐⭐⭐⭐ 数学证明严谨，从对称性守恒到拉普拉斯谱分解的物理意义阐述极为清晰透彻。
价值: ⭐⭐⭐⭐⭐ 极大降低了原子基座模型的预训练与推理部署门槛，对材料与药物计算具有重大的实际工业价值。