面向高效分子性质预测的量化 SO(3) 等变图神经网络
机器学习
2026-03-04 v2
摘要
在边缘设备上部署具有3D旋转等变性（即 SO(3）群的图神经网络（GNN）具有挑战，因其高计算成本。本文通过低位量化技术压缩和加速一个 SO(3) 等变 GNN 来解决这一问题。具体而言，我们为量化等变转换器引入了三项创新：（1）一种幅度-方向解耦的量化方案，分别对等变（向量）特征的范数和方向进行量化；（2）一种分支分离的量化感知训练策略，在注意力机制的 SO(3)-GNN 中对不变和等变特征通道采用不同处理方式；（3）一种增强鲁棒性的注意力归一化机制，以稳定低精度注意力计算。实验在 QM9 和 rMD17 分子基准上表明，8 位模型在能量和力预测方面精度相当于全精度基线，同时显著提升效率。我们还进行消融研究，以量化每个组件对维持精度和等变性贡献，使用 Local error of equivariance（LEE）指标。所提出的技术使在不牺牲精度或物理对称性的前提下，将对称感知 GNN 部署到实际化学应用中成为可能，推理速度提升 2.37--2.73 倍，模型规模缩小 4 倍。
引用
@article{arxiv.2601.02213,
title = {Quantized SO(3)-Equivariant Graph Neural Networks for Efficient Molecular Property Prediction},
author = {Haoyu Zhou and Ping Xue and Hao Zhang and Tianfan Fu},
journal= {arXiv preprint arXiv:2601.02213},
year = {2026}
}