Towards Cross-Modal Text-Molecule Retrieval with Better Modality Alignment

Jia Song,Wanru Zhuang,Yujie Lin,Liang Zhang,Chunyan Li,Jinsong Su,Song He,Xiaochen Bo

2024-10-31

Abstract:Cross-modal text-molecule retrieval model aims to learn a shared feature space of the text and molecule modalities for accurate similarity calculation, which facilitates the rapid screening of molecules with specific properties and activities in drug design. However, previous works have two main defects. First, they are inadequate in capturing modality-shared features considering the significant gap between text sequences and molecule graphs. Second, they mainly rely on contrastive learning and adversarial training for cross-modality alignment, both of which mainly focus on the first-order similarity, ignoring the second-order similarity that can capture more structural information in the embedding space. To address these issues, we propose a novel cross-modal text-molecule retrieval model with two-fold improvements. Specifically, on the top of two modality-specific encoders, we stack a memory bank based feature projector that contain learnable memory vectors to extract modality-shared features better. More importantly, during the model training, we calculate four kinds of similarity distributions (text-to-text, text-to-molecule, molecule-to-molecule, and molecule-to-text similarity distributions) for each instance, and then minimize the distance between these similarity distributions (namely second-order similarity losses) to enhance cross-modal alignment. Experimental results and analysis strongly demonstrate the effectiveness of our model. Particularly, our model achieves SOTA performance, outperforming the previously-reported best result by 6.4%.

Information Retrieval

What problem does this paper attempt to address?

The problem that this paper attempts to solve is to achieve the efficiency and accuracy of cross - modal text - molecule retrieval in drug design. Specifically, the author points out two main deficiencies in existing methods when dealing with text and molecule modalities:

1. **Insufficient capture of modality - shared features**: When dealing with the significant differences between text sequences and molecule graphs, existing methods cannot fully extract modality - shared features.
2. **Dependence only on first - order similarity**: Existing methods mainly rely on contrastive learning and adversarial training to achieve cross - modal alignment, but these methods mainly focus on first - order similarity and ignore second - order similarity which can capture more structural information.

To solve these problems, the author proposes a new cross - modal text - molecule retrieval model, which improves the effect of modality alignment by introducing a memory - bank - based feature projector and second - order similarity loss. The specific improvements are as follows:

- **Memory - bank - based feature projector**: A memory - bank - based feature projector with learnable memory vectors is stacked on top of the two modality - specific encoders to better extract modality - shared features.
- **Second - order similarity loss**: During the model training process, four similarity distributions (text - to - text, text - to - molecule, molecule - to - molecule, molecule - to - text) for each instance are calculated, and cross - modal alignment is enhanced by minimizing the distances between these similarity distributions (i.e., second - order similarity loss).

Through these improvements, the model has achieved significant performance improvements on multiple datasets. In particular, its performance on the ChEBI - 20 dataset is better than the previous best result, with a 6.4% improvement in the Hits@1 metric. This indicates that the model has high effectiveness and practicality in the cross - modal text - molecule retrieval task in drug design.