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

2026-01-05 · Haoyu Zhou, Ping Xue, Hao Zhang, Tianfan Fu 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.

📄 PDF Abstract BibTeX arXiv:2601.02213

Code (0)

등록된 구현이 없습니다.

Tasks

Molecular Property Prediction

A Complete Guide to Spherical Equivariant Graph Transformers

2025-12-15 · Sophia Tang arxiv

Spherical equivariant graph neural networks (EGNNs) provide a principled framework for learning on three-dimensional molecular and biomolecular systems, where predictions must respect the rotational symmetries inherent i…

Molecular Property Prediction

Equivariant Graph Attention Networks for Molecular Property Prediction

2022-02-20 · Tuan Le, Frank Noé, Djork-Arné Clevert

Learning and reasoning about 3D molecular structures with varying size is an emerging and important challenge in machine learning and especially in drug discovery. Equivariant Graph Neural Networks (GNNs) can simultaneou…

Drug Discovery Graph Attention Molecular Property Prediction Prediction+1

Relevance of Rotationally Equivariant Convolutions for Predicting Molecular Properties

2020-08-19 · Benjamin Kurt Miller, Mario Geiger, Tess E. Smidt, Frank Noé

Equivariant neural networks (ENNs) are graph neural networks embedded in $\mathbb{R}^3$ and are well suited for predicting molecular properties. The ENN library e3nn has customizable convolutions, which can be designed t…

Molecular Property Prediction Property Prediction

QT-Net: Rethinking Evaluation of AI Models in Atomic Chemical Space

2026-05-11 · Pablo Martínez Crespo, Stefano Ribes, Martin Rahm, Richard Beckmann 외 arxiv

Atomic properties such as partial charges or multipoles encode chemically meaningful information that can inform downstream molecular property prediction, but their evaluation as machine learning targets has been complic…

Molecular Property Prediction Graph Neural Network

Using Multiple Vector Channels Improves E(n)-Equivariant Graph Neural Networks

2023-09-06 · Daniel Levy, Sékou-Oumar Kaba, Carmelo Gonzales, Santiago Miret 외

We present a natural extension to E(n)-equivariant graph neural networks that uses multiple equivariant vectors per node. We formulate the extension and show that it improves performance across different physical systems…