Atomas: Hierarchical Alignment on Molecule-Text for Unified Molecule Understanding and Generation

TL;DR

Atomas introduces a hierarchical alignment framework for molecule-text representation learning, capturing fine-grained details and improving performance across diverse molecular understanding and generation tasks.

Contribution

It proposes a novel hierarchical alignment model that learns fine-grained fragment correspondence between molecules and text, enhancing molecular representation capabilities.

Findings

Outperforms 11 baselines on 12 tasks across 11 datasets

Demonstrates robustness and scalability in scaling experiments

Qualitative analysis confirms chemical relevance of learned representations

Abstract

Molecule-and-text cross-modal representation learning has emerged as a promising direction for enhancing the quality of molecular representation, thereby improving performance in various scientific fields. However, most approaches employ a global alignment approach to learn the knowledge from different modalities that may fail to capture fine-grained information, such as molecule-and-text fragments and stereoisomeric nuances, which is crucial for downstream tasks. Furthermore, it is incapable of modeling such information using a similar global alignment strategy due to the lack of annotations about the fine-grained fragments in the existing dataset. In this paper, we propose Atomas, a hierarchical molecular representation learning framework that jointly learns representations from SMILES strings and text. We design a Hierarchical Adaptive Alignment model to automatically learn the fine-grained fragment correspondence between two modalities and align these representations at three semantic levels. Atomas's end-to-end training framework supports understanding and generating molecules, enabling a wider range of downstream tasks. Atomas achieves superior performance across 12 tasks on 11 datasets, outperforming 11 baseline models thus highlighting the effectiveness and versatility of our method. Scaling experiments further demonstrate Atomas's robustness and scalability. Moreover, visualization and qualitative analysis, validated by human experts, confirm the chemical relevance of our approach. Codes are released on https://github.com/yikunpku/Atomas.