Timezone: »

 
Spotlight
Learning Substructure Invariance for Out-of-Distribution Molecular Representations
Nianzu Yang · Kaipeng Zeng · Qitian Wu · Xiaosong Jia · Junchi Yan

Thu Dec 08 09:00 AM -- 11:00 AM (PST) @
in Lightning Talks 5B-3 »

Molecule representation learning (MRL) has been extensively studied and current methods have shown promising power for various tasks, e.g., molecular property prediction and target identification. However, a common hypothesis of existing methods is that either the model development or experimental evaluation is mostly based on i.i.d. data across training and testing. Such a hypothesis can be violated in real-world applications where testing molecules could come from new environments, bringing about serious performance degradation or unexpected prediction. We propose a new representation learning framework entitled MoleOOD to enhance the robustness of MRL models against such distribution shifts, motivated by an observation that the (bio)chemical properties of molecules are usually invariantly associated with certain privileged molecular substructures across different environments (e.g., scaffolds, sizes, etc.). Specifically, We introduce an environment inference model to identify the latent factors that impact data generation from different distributions in a fully data-driven manner. We also propose a new learning objective to guide the molecule encoder to leverage environment-invariant substructures that more stably relate with the labels across environments. Extensive experiments on ten real-world datasets demonstrate that our model has a stronger generalization ability than existing methods under various out-of-distribution (OOD) settings, despite the absence of manual specifications of environments. Particularly, our method achieves up to 5.9\% and 3.9\% improvement over the strongest baselines on OGB and DrugOOD benchmarks in terms of ROC-AUC, respectively. Our source code is publicly available at \url{https://github.com/yangnianzu0515/MoleOOD}.

Keywords: Out-of-Distribution · invariant learning · Molecule Representation Learning

Author Information

Nianzu Yang (Shanghai Jiao Tong University)
Kaipeng Zeng (Shanghai Jiaotong University)
Qitian Wu (Shanghai Jiao Tong University)
Xiaosong Jia (University of California, Berkeley)
Junchi Yan (Shanghai Jiao Tong University)

Related Events (a corresponding poster, oral, or spotlight)

More from the Same Authors