Enhancing molecular property prediction with quantized GNN models

Abstract

Efficient and reliable prediction of molecular properties, such as water solubility, hydration free energy, lipophilicity, and quantum mechanical properties, is essential for rational compound design in the chemical and pharmaceutical industries. While Graph Neural Networks (GNNs) have significantly advanced molecular property prediction tasks, their high memory footprint, computational demands, and inference latency are often overlooked. These challenges hinder the deployment of property prediction models on resource-constrained devices such as smartphones and IoT devices. Therefore, optimizing storage, reducing resource consumption, and improving inference speed are crucial. This paper presents a systematic approach to molecular networks by integrating GNN models with the DoReFa-Net quantization algorithm. The proposed method aims to enhance computational efficiency while maintaining predictive performance, enabling lightweight yet effective models suitable for molecular task. The study investigates the impact of different bitwidth quantization levels on model performance, using metrics such as RMSE and MAE. Results show that, for physical chemistry datasets, the effectiveness of quantization is highly dependent on the model architecture. Notably, the quantum mechanical dipole moment task maintains strong performance up to 8-bit precision, achieving similar or slightly better results. However, extreme quantization, particularly at 2-bit precision, severely degrades performance, highlighting the limitations of aggressive compression.