arXiv:2603.05343cs.LG2026-03

提出几何感知量化方法,让旋转对称的分子模型更高效且不失准。

Preserving Continuous Symmetry in Discrete Spaces: Geometric-Aware Quantization for SO(3)-Equivariant GNNs

  • 分离长度与方向分量进行量化,保持几何结构不变。
  • 4位权重+8位激活在基准上误差仅9.31 meV,比原版低23.20 meV。
  • 适合做高精度分子模拟的轻量化部署,尤其适合算力受限场景。

等变图神经网络对物理一致的分子模拟至关重要,但高阶表示带来高昂计算与内存开销。低比特量化虽可缓解,但直接应用于旋转敏感特征会破坏SO(3)等变结构,导致显著误差和守恒律违反。本文提出几何感知量化(GAQ)框架,在离散空间中严格保持连续对称性。核心贡献包括:(1)幅度-方向解耦量化(MDDQ),分离不变长度与等变方向以维持几何保真;(2)针对标量与矢量特征采用不同量化策略的训练方法;(3)鲁棒注意力归一化机制,稳定低比特下的梯度。在rMD17基准上,我们的W4A8模型准确率媲美FP32基线(9.31 meV vs. 23.20 meV),局部等变误差降低30倍以上。在消费级硬件上实现2.39倍推理加速与4倍内存压缩,支持纳秒级稳定、能量守恒的分子动力学模拟。

原文摘要 · Abstract (English)

Equivariant Graph Neural Networks (GNNs) are essential for physically consistent molecular simulations but suffer from high computational costs and memory bottlenecks, especially with high-order representations. While low-bit quantization offers a solution, applying it naively to rotation-sensitive features destroys the SO(3)-equivariant structure, leading to significant errors and violations of conservation laws. To address this issue, in this work, we propose a Geometric-Aware Quantization (GAQ) framework that compresses and accelerates equivariant models while rigorously preserving continuous symmetry in discrete spaces. Our approach introduces three key contributions: (1) a Magnitude-Direction Decoupled Quantization (MDDQ) scheme that separates invariant lengths from equivariant orientations to maintain geometric fidelity; (2) a symmetry-aware training strategy that treats scalar and vector features with distinct quantization schedules; and (3) a robust attention normalization mechanism to stabilize gradients in low-bit regimes. Experiments on the rMD17 benchmark demonstrate that our W4A8 models match the accuracy of FP32 baselines (9.31 meV vs. 23.20 meV) while reducing Local Equivariance Error (LEE) by over 30x compared to naive quantization. On consumer hardware, GAQ achieves 2.39x inference speedup and 4x memory reduction, enabling stable, energy-conserving molecular dynamics simulations for nanosecond timescales.

等变神经网络量化分子模拟

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