mlsanityData to 7 Oct 2026

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

arXiv 2601.02213 ↗First posted 5 Jan 2026cs.LGcited by 0

Not on the map. Why?

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.

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

Builds on

32 of 39 references.

References · 32

  1. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  2. 2017Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inferencecited by 1,548
  3. 2022SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Modelscited by 1,338
  4. 2016Group Equivariant Convolutional Networkscited by 1,154
  5. 2021E(n) Equivariant Graph Neural Networkscited by 951
  6. 2018Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point cloudscited by 735
  7. 2021E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentialscited by 701
  8. 2023Scaling Vision Transformers to 22 Billion Parameterscited by 690
  9. 2019Learned Step Size Quantizationcited by 604
  10. 2020SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networkscited by 602
  11. 2016DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradientscited by 557
  12. 2020Directional Message Passing for Molecular Graphscited by 556
20 more
  1. 2022ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformerscited by 517
  2. 2018PACT: Parameterized Clipping Activation for Quantized Neural Networkscited by 456
  3. 2018Spherical CNNscited by 452
  4. 2021Equivariant message passing for the prediction of tensorial properties and molecular spectracited by 438
  5. 20183D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Datacited by 348
  6. 2021GemNet: Universal Directional Graph Neural Networks for Moleculescited by 342
  7. 2016Harmonic Networks: Deep Translation and Rotation Equivariancecited by 293
  8. 2020Query-Key Normalization for Transformerscited by 266
  9. 2019Gauge Equivariant Convolutional Networks and the Icosahedral CNNcited by 262
  10. 2020Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Moleculescited by 254
  11. 2022Learning Local Equivariant Representations for Large-Scale Atomistic Dynamicscited by 254
  12. 2021A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groupscited by 185
  13. 2022Graph neural networks for materials science and chemistrycited by 181
  14. 2022QDrop: Randomly Dropping Quantization for Extremely Low-bit Post-Training Quantizationcited by 175
  15. 2022e3nn: Euclidean Neural Networkscited by 171
  16. 2021SpookyNet: Learning Force Fields with Electronic Degrees of Freedom and Nonlocal Effectscited by 103
  17. 2020Degree-Quant: Quantization-Aware Training for Graph Neural Networkscited by 93
  18. 2022ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphscited by 69
  19. 2022The Lie Derivative for Measuring Learned Equivariancecited by 51
  20. 2022So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systemscited by 39

Built on by

Not cited yet.