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

Haoyu Zhou; Ping Xue; Hao Zhang; Tianfan Fu

cs.LGMar 5, 2026

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

Authors: Haoyu Zhou, Ping Xue, Hao Zhang, Tianfan Fu

Organizations: Nanjing University, Nanjing, China. · Gusu Lab, Suzhou, China.

Abstract

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.

Figures & tables

Fig. 1: Overview of the proposed equivariant quantization framework. (a) MDDQ : Decouples equivariant vectors into magnitude $\|\mathbf{v}\|$ and direction $\mathbf{u}$ 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 and temperature scaling $\tau$ . (d) Results Summary : Our method achieves $2.37$ – $2.73\times$ faster inference and $\sim 4\times$ memory reduction with ultra-low equivariance error ( $\text{LEE}\approx 0.15$ ).

Fig. 2: Detailed architecture and quantization pipeline. Inputs $Z$ and $\mathbf{r}$ are processed through decoupled invariant and equivariant paths, employing specialized quantization strategies (Linear vs. MDDQ) for each branch.

Architecture	$C_{\text{full}}$ (FP32)	$\ell_{\max}$	$C_{\text{quant}}$ (k-bit)	Gain $C_{\text{quant}}/C_{\text{full}}$
PaiNN [ 28 ]	$O(n\langle N\rangle 4F)$	1	$\rho_{k}$	$O(n\langle N\rangle 4F)\,\rho_{k}$
SpookyNet [ 29 ]	$O(n\langle N\rangle(\ell_{\max}+1)^{2}F)$	2	$\rho_{k}$	$O(n\langle N\rangle(\ell_{\max}+1)^{2}F)\,\rho_{k}$
NequIP [ 1 ]	$O(n\langle N\rangle(\ell_{\max}+1)^{6}F)$	3	$\rho_{k}$	$O(n\langle N\rangle(\ell_{\max}+1)^{6}F)\,\rho_{k}$
So3krates [ 2 ]	$O(n\langle N\rangle((\ell_{\max}+1)^{2}+F))$	1	$\rho_{k}$	$O(n\langle N\rangle((\ell_{\max}+1)^{2}+F))\,\rho_{k}$

TABLE I: Complexity with and without quantization.

C_{\text{full}}

is per-layer asymptotic cost in full precision;

C_{\text{quant}}

is the cost when both weights and activations are

k

-bit (here

\rho_{k}\equiv k/32

). Quantization yields a constant-factor speedup

\approx\rho_{k}

without changing scaling in

n

\langle N\rangle

F

, or

\ell_{\max}

. For So3krates, we utilize

\ell_{\max}=1

as discussed in [ 2 ] for efficiency and compatibility with our vector quantization scheme, though the architecture supports higher orders.

∗ Validation loss stagnated due to gradient fracture.
Method	Bits (W/A)	E-MAE	F-MAE	Stability
FP32 Baseline	32 / 32	23.20	21.20	Stable
Naive INT8	8 / 8	118.20	102.39	Degraded
SVQ-KMeans	8 / 8	226.20 ∗	N/A	Diverged
Degree-Quant [ 22 ]	8 / 8	63.20	58.90	Stable
Ours (GAQ)	4 / 8	9.31	22.60	Stable

TABLE II: Performance comparison on Azobenzene (rMD17). Energy MAE in meV and Force MAE in meV/Å. ”Stability” indicates whether the training converged.

Fig. 3: Energy conservation in NVE dynamics (1 ns). The Naive INT8 model ( red ) exhibits catastrophic energy divergence and explosion within 100 ps due to symmetry breaking. In contrast, our W4A8 model ( orange ) maintains excellent stability comparable to the FP32 baseline ( gray ). Note that the realistic thermal fluctuations observed in our model confirm that the simulation retains correct physical dynamics without long-term drift ( $<0.15$ meV/atom/ps).

Method	LEE (meV/Å)	Remark
FP32 Baseline	$\approx 0.0$	Exact equivariance
Naive INT8	5.23	Broken symmetry
Degree-Quant	2.10	Partially preserved
Ours (W4A8)	0.15	Preserved

TABLE III: Symmetry analysis. LEE measures the sensitivity of predictions to input rotation (lower is better). Our method reduces the symmetry error by

>30\times

compared to naive quantization.

Operation	FP32	Ours (W4A8)	Speedup
Memory I/O (Weights)	120.5	30.1	4.0 $\times$
Compute (GEMM)	45.0	25.0	1.8 $\times$
Quant Overhead	0.0	5.2	-
Attention	15.2	15.2	1.0 $\times$
Total Latency	180.7	75.5	2.39 $\times$

TABLE IV: Latency breakdown. Time in

\mu

s. Our method achieves speedup primarily by reducing memory I/O.

Equivariant GNNs

Equivariant graph neural networks encode how node and edge features transform under symmetries to make geometry-aware predictions.

See all 42 papers

Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning
May 9, 2026Chen Wang, Siyu Hu, Guangming Tan +1Equivariant GNNs Structured Pruning
State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences
GEqTrain: A Configuration-Driven Framework for Retargeting Equivariant Graph Neural Networks Across 3D Scientific Tasks
Jul 21, 2026Daniele Angioletti, Marco Nobile, Vittorio LimongelliFlow Matching Generative Modeling
Faculty of Biomedical Sciences, Euler Institute, Universit´a della Svizzera italiana, 6900 Lugano, Switzerland
Triply-Scalable Equivariant Gaussian Process Modeling
Sep 17, 2026Tim Steinert, David GinsbourgerSparse Gaussian Processes Gaussian Processes
Institute of Mathematical Statistics and Actuarial Science, University of Bern, Bern, Switzerland

May 9, 2026J/K move · Enter open · S save

Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning

Chen Wang, Siyu Hu, Guangming Tan, Weile Jia

State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences

Equivariant GNNs Structured Pruning Neural Network Compression Neural Network Pruning

Open paper AskPDF arXiv

SO(3) equivariant graph neural networks have become the dominant paradigm for atomistic foundation models, achieving high accuracy and data efficiency by building rotational symmetry directly into the architecture. Yet the computational cost of their higher-order tensor operations creates a tough trade-off between model accuracy and inference efficiency. In this paper, we propose a structural pruning method for SO(3) equivariant atomistic foundation models to bridge this accuracy-efficiency gap. The pruning is applied along the channel and order dimensions, with each irreducible representation kept or removed as a complete block, thereby retaining SO(3) equivariance. Starting from a large checkpoint, the pruned model substantially reduces the inference cost while retaining higher accuracy than an independently trained small model. The pruned MACE-MP model outperforms the official from-scratch trained small model on 7 of 9 metrics on the Matbench Discovery leaderboard. In terms of efficiency, compressed MACE-MP and MACE-OFF models contain 1.5

\times

to 4

\times

fewer parameters and require 2.5

\times

to 4

\times

less pre-training compute than training a small model from scratch. For downstream applications, fine-tuning the pruned model reduces energy and force errors by 70.1% and 34.4% compared to training task-specific models from scratch across eight representative downstream datasets. We demonstrate that the method generalizes to other SO(3) equivariant architectures (SevenNet, eSCN) and can be combined with quantization and knowledge distillation for further gains.

Abstract

Equivariant GNNs

Explore similar work

Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning

GEqTrain: A Configuration-Driven Framework for Retargeting Equivariant Graph Neural Networks Across 3D Scientific Tasks

Triply-Scalable Equivariant Gaussian Process Modeling