在离散空间中保持连续对称性:面向 SO(3) 等变图神经网络的几何感知量化
机器学习
2026-03-17 v2
摘要
等变图神经网络 (GNNs) 对物理性一致的分子模拟至关重要,但在高阶表示下会受到高计算成本和内存瓶颈的制约。虽然低位量化提供了解方方案,但对旋转敏感特征直接应用量化会破坏 SO(3) 等变结构,导致显著误差并违反守恒定律。为此,本工作提出了保持离散空间中连续对称性的几何感知量化 (GAQ) 框架,在压缩和加速等变模型的同时严格保持连续对称性。我们的方案引入了三个关键贡献:(1) Magnitude-Direction Decoupled Quantization (MDDQ) 方案,将不变长度与等变方向分离以保持几何保真;(2) 针对标量和向量特征的对称感知训练策略;(3) 稳定低位情景中梯度的鲁棒注意力归一化机制。在 rMD17 数据集上进行的实验表明,我们的 W4A8 模型在精度上与 FP32 基准相当(9.31 meV vs. 23.20 meV),同时将局部等变误差 (LEE) 降低了 30 倍以上。
引用
@article{arxiv.2603.05343,
title = {Preserving Continuous Symmetry in Discrete Spaces: Geometric-Aware Quantization for SO(3)-Equivariant GNNs},
author = {Haoyu Zhou and Ping Xue and Hao Zhang and Tianfan Fu},
journal= {arXiv preprint arXiv:2603.05343},
year = {2026}
}