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

arXiv: https://arxiv.org/abs/2603.05343 · HTML: https://mlsanity.com/paper/2603.05343/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 31 references that are in the map.
- First posted on arXiv 5 Mar 2026.

## Builds on (most cited first, 25 of 25)

- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://mlsanity.com/paper/1712.05877/) (2017)
- [SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://mlsanity.com/paper/2211.10438/) (2022)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Learned Step Size Quantization](https://mlsanity.com/paper/1902.08153/) (2019)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients](https://mlsanity.com/paper/1606.06160/) (2016)
- [ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers](https://mlsanity.com/paper/2206.01861/) (2022)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data](https://mlsanity.com/paper/1807.02547/) (2018)
- [Harmonic Networks: Deep Translation and Rotation Equivariance](https://mlsanity.com/paper/1612.04642/) (2016)
- [Query-Key Normalization for Transformers](https://mlsanity.com/paper/2010.04245/) (2020)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups](https://mlsanity.com/paper/2104.09459/) (2021)
- [QDrop: Randomly Dropping Quantization for Extremely Low-bit Post-Training Quantization](https://mlsanity.com/paper/2203.05740/) (2022)
- [SpookyNet: Learning Force Fields with Electronic Degrees of Freedom and Nonlocal Effects](https://mlsanity.com/paper/2105.00304/) (2021)
- [Degree-Quant: Quantization-Aware Training for Graph Neural Networks](https://mlsanity.com/paper/2008.05000/) (2020)
- [The Lie Derivative for Measuring Learned Equivariance](https://mlsanity.com/paper/2210.02984/) (2022)
- [So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems](https://mlsanity.com/paper/2205.14276/) (2022)
- [Exact conservation laws for neural network integrators of dynamical systems](https://mlsanity.com/paper/2209.11661/) (2022)
- [Optimal Invariant Bases for Atomistic Machine Learning](https://mlsanity.com/paper/2503.23515/) (2025)

## Abstract (arXiv)

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.