arXiv:2601.02213cs.LG2026-01

8位量化让分子属性预测模型更快更小,还保持物理对称性。

Quantized SO(3)-Equivariant Graph Neural Networks for Efficient Molecular Property Prediction

  • 分解向量特征的大小与方向分别量化,提升精度。
  • 8位模型在QM9和rMD17上性能接近全精度,误差极小。
  • 适合需要高效部署的化学计算、移动端分子建模场景。

将对3D旋转对称(SO(3))不变的图神经网络部署到边缘设备面临高计算开销挑战。本文通过低比特量化压缩并加速此类SO(3)-等变GNN。提出三项创新:(1) 量化解耦范数与方向的向量特征;(2) 区分不变与等变通道的量化感知训练策略;(3) 增强注意力计算鲁棒性的归一化机制。在QM9和rMD17数据集上的实验表明,8位模型在能量与力预测上性能接近全精度基线,效率显著提升。消融实验使用局部等变误差(LEE)评估各组件贡献。所提方法使对称感知GNN在实际化学应用中实现2.37–2.73倍推理加速、4倍模型压缩,且不损失精度或物理对称性。

原文摘要 · Abstract (English)

Deploying 3D graph neural networks (GNNs) that are equivariant to 3D rotations (the group SO(3)) on edge devices is challenging due to their high computational cost. This paper addresses the problem by compressing and accelerating an SO(3)-equivariant GNN using low-bit quantization techniques. Specifically, we introduce three innovations for quantized equivariant transformers: (1) a magnitude-direction decoupled quantization scheme that separately quantizes the norm and orientation of equivariant (vector) features, (2) a branch-separated quantization-aware training strategy that treats invariant and equivariant feature channels differently in an attention-based $SO(3)$-GNN, and (3) a robustness-enhancing attention normalization mechanism that stabilizes low-precision attention computations. Experiments on the QM9 and rMD17 molecular benchmarks demonstrate that our 8-bit models achieve accuracy on energy and force predictions comparable to full-precision baselines with markedly improved efficiency. We also conduct ablation studies to quantify the contribution of each component to maintain accuracy and equivariance under quantization, using the Local error of equivariance (LEE) metric. The proposed techniques enable the deployment of symmetry-aware GNNs in practical chemistry applications with 2.37--2.73x faster inference and 4x smaller model size, without sacrificing accuracy or physical symmetry.

分子预测等变网络量化边缘计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。