Learning Invariant Molecular Representation in Latent Discrete Space

Xiang Zhuang · Qiang Zhang · Keyan Ding · Yatao Bian · Xiao Wang · Jingsong Lv · Hongyang Chen · Huajun Chen

Great Hall & Hall B1+B2 (level 1) #106
[ ]
Thu 14 Dec 8:45 a.m. PST — 10:45 a.m. PST


Molecular representation learning lays the foundation for drug discovery. However, existing methods suffer from poor out-of-distribution (OOD) generalization, particularly when data for training and testing originate from different environments. To address this issue, we propose a new framework for learning molecular representations that exhibit invariance and robustness against distribution shifts. Specifically, we propose a strategy called ``first-encoding-then-separation'' to identify invariant molecule features in the latent space, which deviates from conventional practices. Prior to the separation step, we introduce a residual vector quantization module that mitigates the over-fitting to training data distributions while preserving the expressivity of encoders. Furthermore, we design a task-agnostic self-supervised learning objective to encourage precise invariance identification, which enables our method widely applicable to a variety of tasks, such as regression and multi-label classification. Extensive experiments on 18 real-world molecular datasets demonstrate that our model achieves stronger generalization against state-of-the-art baselines in the presence of various distribution shifts. Our code is available at

Chat is not available.