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

TL;DR

This paper introduces low-bit quantization techniques for SO(3)-equivariant GNNs, enabling efficient and accurate molecular property prediction on edge devices by reducing model size and inference time while preserving symmetry and accuracy.

Contribution

It proposes novel quantization strategies for equivariant GNNs, including magnitude-direction decoupling, branch-separated training, and attention normalization, to maintain accuracy under low precision.

Findings

8-bit models match full-precision accuracy on QM9 and rMD17

Achieve 2.37--2.73x faster inference

Reduce model size by 4x

Abstract

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.