Exploring Optimal Transport-Based Multi-Grained Alignments for Text-Molecule Retrieval

TL;DR

This work introduces the Optimal Transport-based Multi-grained Alignments model (ORMA), a novel approach that facilitates multi-grained alignments between textual descriptions and molecules, and is the first attempt to explore alignments at both the motif and multi-token levels.

Abstract

The field of bioinformatics has seen significant progress, making the cross-modal text-molecule retrieval task increasingly vital. This task focuses on accurately retrieving molecule structures based on textual descriptions, by effectively aligning textual descriptions and molecules to assist researchers in identifying suitable molecular candidates. However, many existing approaches overlook the details inherent in molecule sub-structures. In this work, we introduce the Optimal TRansport-based Multi-grained Alignments model (ORMA), a novel approach that facilitates multi-grained alignments between textual descriptions and molecules. Our model features a text encoder and a molecule encoder. The text encoder processes textual descriptions to generate both token-level and sentence-level representations, while molecules are modeled as hierarchical heterogeneous graphs, encompassing atom, motif, and molecule nodes to extract representations at these three levels. A key innovation in ORMA is the application of Optimal Transport (OT) to align tokens with motifs, creating multi-token representations that integrate multiple token alignments with their corresponding motifs. Additionally, we employ contrastive learning to refine cross-modal alignments at three distinct scales: token-atom, multitoken-motif, and sentence-molecule, ensuring that the similarities between correctly matched text-molecule pairs are maximized while those of unmatched pairs are minimized. To our knowledge, this is the first attempt to explore alignments at both the motif and multi-token levels. Experimental results on the ChEBI-20 and PCdes datasets demonstrate that ORMA significantly outperforms existing state-of-the-art (SOTA) models.

Figures (7)

• Figure 1: The text-molecule retrieval task is designed to retrieve molecules based on text queries, while molecule-text retrieval task does the opposite. The red box indicates the ground-truth retrieval result.
• Figure 2: The illustration of our model. In short, our model comprises a text encoder and a molecule encoder and designs three alignment losses $\mathcal{L}_{ta}, \mathcal{L}_{mm}, \mathcal{L}_{sm}$ at token-atom, multitoken-motif, and sentence-molecule levels, respectively. On the left side of the figure is a molecular graph, where the nodes at the top represent atom nodes. In the middle, each node with the solid edge is a motif node, connected to the atom nodes it contains. At the bottom, the node with the dashed edge is the molecule node, connected to all motif nodes.
• Figure 3: The process of token-atom alignments. Within a batch, we update the atom representations of each sample and then calculate the token-atom level similarities between samples. Then we use contrastive learning to maximize the similarities between matched text-molecule pairs while minimizing those between unmatched pairs.
• Figure 4: The process of obtaining multi-token representations through OT. Considering the alignments between token representations and motif representations as an optimal transport problem, we fuse the token representations into multi-token representations corresponding to specific motifs.
• Figure 5: The visualization for the alignments of motifs and multi-tokens. The motifs marked in red and green on the left correspond to the multi-tokens highlighted in red and green on the right, respectively.