Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Haoyu Zhou, Ping Xue, Hao Zhang, Tianfan Fu

TL;DR

A Geometric-Aware Quantization (GAQ) framework that compresses and accelerates equivariant models while rigorously preserving continuous symmetry in discrete spaces is proposed, enabling stable, energy-conserving molecular dynamics simulations for nanosecond timescales.

Abstract

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.

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

TL;DR

A Geometric-Aware Quantization (GAQ) framework that compresses and accelerates equivariant models while rigorously preserving continuous symmetry in discrete spaces is proposed, enabling stable, energy-conserving molecular dynamics simulations for nanosecond timescales.

Abstract

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.
Paper Structure (24 sections, 1 theorem, 10 equations, 3 figures, 4 tables)

This paper contains 24 sections, 1 theorem, 10 equations, 3 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Optimal Bases and Codebooks.
  4. Measuring Equivariance.
  5. Method
  6. Preliminaries: Group Representations and Equivariance
  7. Equivariant GNN Architecture Overview
  8. Magnitude–Direction Decoupled Quantization (MDDQ)
  9. Definition 3.1 (Magnitude–Direction Decoupled Quantization).
  10. Proposition 3.4 (Angular error bound on $S^2$).
  11. Optimization on the Spherical Manifold
  12. Manifold Gradient Flow.
  13. Geometric Straight-Through Estimator (Geometric STE).
  14. Branch-Separated Scheduling.
  15. Robust Attention Normalization
  16. ...and 9 more sections

Key Result

Proposition 3.1

(Orthogonality of Updates). The Geometric STE ensures that the gradient update is strictly orthogonal to the feature vector, i.e., $\langle \mathbf{u}, \frac{\partial \mathcal{L}}{\partial \mathbf{u}} \rangle = 0$. This implies that to first order, the optimization trajectory remains on the manifold

Figures (3)

  • Figure 1: Overview of the proposed equivariant quantization framework. (a) MDDQ: Decouples equivariant vectors into magnitude $\|\mathbf{v}\|$ and direction $\mathbf{u}$ to preserve geometric orientation under low precision. (b) Branch-separated QAT: Treats invariant and equivariant features differently with a staged training schedule. (c) Robust Attention: Stabilizes dot-products via $\ell_2$ normalization and temperature scaling $\tau$. (d) Results Summary: Our method achieves $2.37$--$2.73\times$ faster inference and $\sim 4\times$ memory reduction with ultra-low equivariance error ($\text{LEE} \approx 0.15$).
  • Figure 2: Detailed architecture and quantization pipeline. Inputs $Z$ and $\mathbf{r}$ are processed through decoupled invariant and equivariant paths, employing specialized quantization strategies (Linear vs. MDDQ) for each branch.
  • Figure 3: Energy conservation in NVE dynamics (1 ns). The Naive INT8 model (red) exhibits catastrophic energy divergence and explosion within 100 ps due to symmetry breaking. In contrast, our W4A8 model (orange) maintains excellent stability comparable to the FP32 baseline (gray). Note that the realistic thermal fluctuations observed in our model confirm that the simulation retains correct physical dynamics without long-term drift ($<0.15$ meV/atom/ps).

Theorems & Definitions (1)

  • Proposition 3.1