Improving molecular representation learning with metric learning-enhanced optimal transport

Fang Wu; Nicolas Courty; Shuting Jin; Stan Z Li

doi:10.1016/j.patter.2023.100714

Improving molecular representation learning with metric learning-enhanced optimal transport

Patterns (N Y). 2023 Mar 29;4(4):100714. doi: 10.1016/j.patter.2023.100714. eCollection 2023 Apr 14.

Authors

Fang Wu^{1

2}, Nicolas Courty³, Shuting Jin^{4

5}, Stan Z Li¹

Affiliations

¹ School of Engineering, Westlake University, Hangzhou 310024, China.
² Institute of AI Industry Research, Tsinghua University, Beijing 100084, China.
³ French National Centre for Scientific Research, Southern Brittany University, Lorient, France.
⁴ School of Informatics, Xiamen University, Xiamen 361005, China.
⁵ National Institute for Data Science in Health and Medicine, Xiamen University, Xiamen 361005, China.

Abstract

Training data are usually limited or heterogeneous in many chemical and biological applications. Existing machine learning models for chemistry and materials science fail to consider generalizing beyond training domains. In this article, we develop a novel optimal transport-based algorithm termed MROT to enhance their generalization capability for molecular regression problems. MROT learns a continuous label of the data by measuring a new metric of domain distances and a posterior variance regularization over the transport plan to bridge the chemical domain gap. Among downstream tasks, we consider basic chemical regression tasks in unsupervised and semi-supervised settings, including chemical property prediction and materials adsorption selection. Extensive experiments show that MROT significantly outperforms state-of-the-art models, showing promising potential in accelerating the discovery of new substances with desired properties.

Keywords: deep learning; domain adaptation; drug discovery; geometric neural network; materials synthesis; molecular representation learning; optimal transport.