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

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

## Key facts

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

## Builds on (most cited first, 30 of 32)

- [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)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [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)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers](https://mlsanity.com/paper/2206.01861/) (2022)
- [PACT: Parameterized Clipping Activation for Quantized Neural Networks](https://mlsanity.com/paper/1805.06085/) (2018)
- [Spherical CNNs](https://mlsanity.com/paper/1801.10130/) (2018)
- [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)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [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)
- [Gauge Equivariant Convolutional Networks and the Icosahedral CNN](https://mlsanity.com/paper/1902.04615/) (2019)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules](https://mlsanity.com/paper/2011.14115/) (2020)
- [A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups](https://mlsanity.com/paper/2104.09459/) (2021)
- [Graph neural networks for materials science and chemistry](https://mlsanity.com/paper/2208.09481/) (2022)
- [QDrop: Randomly Dropping Quantization for Extremely Low-bit Post-Training Quantization](https://mlsanity.com/paper/2203.05740/) (2022)
- [e3nn: Euclidean Neural Networks](https://mlsanity.com/paper/2207.09453/) (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)
- [ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs](https://mlsanity.com/paper/2206.08515/) (2022)

## Abstract (arXiv)

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.
