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

TL;DR

This work tackles the challenge of deploying $SO(3)$-equivariant GNNs for molecular property prediction on resource-limited devices by introducing an equivariance-aware quantization framework. The approach combines Magnitude-Direction Decoupled Quantization ($MDDQ$), branch-separated quantization-aware training for invariant vs. equivariant features, and robust attention normalization, together with an equivariance-preserving loss (LEE). On QM9 and rMD17 benchmarks, 8-bit models achieve accuracy close to FP32 while delivering $2.37$--$2.73\times$ faster inference and about $4\times$ memory reduction; aggressive 4-bit weight/8-bit activation quantization (W4A8) can further improve efficiency with stable molecular dynamics. These results demonstrate practical, symmetry-preserving low-bit deployment of GNNs for chemistry, enabling on-device molecular analyses and paving the way for broader deployment of equivariant models on edge hardware.

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.

Figures (4)

• Figure 1: Overview of the proposed equivariant quantization framework. (a) MDDQ: Decouples equivariant vectors into magnitude $r$ and direction $\hat{\mathbf{h}}$ 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 of queries and keys. (d) Results Summary: Our method achieves $2.37$--$2.73\times$ faster inference and $\sim 4\times$ memory reduction with accuracy comparable to FP32 models.
• Figure 2: Detailed architecture and quantization pipeline. Inputs $Z$ and $\mathbf{r}$ are processed through decoupled invariant and equivariant paths.
• Figure 3: MD stability and per-structure force error distribution on rotated rMD17--Ethanol, demonstrating robustness under aggressive quantization.
• Figure 4: Accuracy--efficiency bar chart on rMD17--Ethanol. Grouped bars show energy MAE (left $y$-axis, meV; hatched “//”) and force MAE (right $y$-axis, meV/Å; dotted hatch). Numbers above the energy bars indicate the relative speedup over FP32, $S = t_{\mathrm{FP32}} / t_{\mathrm{method}}$ (higher is faster). INT8 scalar/mixed use the GPU ms/structure latency protocol, while TorchAO w4a8 uses the current CPU-only ms/batch; to compare across devices, we report the dimensionless speedup $S$. The legend is placed outside (right) to avoid covering bars.