MolFusion: Multimodal Fusion Learning for Molecular Representations via Multi-granularity Views

Numerous molecular representation learning methods are based on single molecular representations, such as SMILES [13], molecular graphs [14], and 3D molecular representations [14]. There are relatively comprehensive methods for both molecular-level and atomic-level training for unimodal molecular representation learning. In molecular-level methods, contrastive learning [15, 16, 17] is a common training technique. MolCLR [18], MoCL [19], and GraphCL [20] perform contrastive learning on augmented molecular graphs after data augmentation operations. Knowledge-based BERT [21] generates various SMILES from canonical SMILES for data augmentation. In MolGNet [22], molecular graphs are split into two halves. Two halves from the same molecule are used as positive samples, and halves from different molecules are used as negative samples. KCL [23] enhances molecular graphs using knowledge graphs, adding corresponding attribute nodes to atoms and connecting them with edges. The molecular graphs before and after knowledge enhancement are encoded separately, with the encodings of the same molecule as positive samples and the encodings of different molecules as negative samples.

Beyond contrastive learning, SMILES Transformer [24] uses the autoencoding method to learn molecular representations. Similarly, X-Mol [25] designs Transformer layers, using bidirectional self-attention layers in the first half to learn SMILES representations and unidirectional self-attention layers in the second half to regenerate SMILES from the learned embeddings. Grover [26] and knowledge-based BERT propose a molecular features prediction task, predicting molecular features, such as the presence of specific functional groups, from the embeddings of the entire molecule.

For atomic-level methods, the most classic approach is masked language models [27, 28, 29]. CHEM-BERT [30], for example, masks atoms in SMILES and uses molecular contextual information to predict the types of masked atoms. For molecular graphs, atom features are masked, and other molecular features, such as bond features, are used to predict the types of atoms [31, 22, 32, 33].

Beyond atom masking tasks, Grover proposes using the learned atomic features to predict the contextual properties of atoms. In the context prediction task [31], atom neighborhood encodings are used as node embeddings, and surrounding graph structures are represented as context embeddings. The task is to predict whether the node embeddings and context embeddings come from the same molecule. InfoGraph [34] proposes to predict which atoms belong to which molecules based on representations of all molecules and the atoms they contain.

Notably, Hu et al. [31] suggests strategies for pre-training at the molecular or atomic level are found to offer limited improvements and may even degrade the performance of many downstream tasks. Therefore, most works [35, 26, 22] attempt to retain both molecular-level and atomic-level training strategies in their studies, employing multi-granularity methods to ensure optimal performance.

Researchers [8, 9, 10] find that different molecular representations can provide complementary information, making the effective utilization of multiple molecular representations a key focus in drug property prediction. For instance, DMP [8] and PanGu [11] utilize molecular linear representations and molecular graph information, GraphMVP [9] leverages molecular graphs and 3D molecular representations, while MEMO [10] processes multiple representations, including SMILES, molecular graphs, 3D representations, and fingerprints. In multimodal methods, contrastive learning and its variants are prevalent. For example, GraphMVP and DMP treat different representations of the same molecule as positive samples and representations from different molecules as negative samples for contrastive learning or its variants. MEMO treats single representations and aggregated representations from the same molecule as positive samples and those from different molecules as negative samples.

Beyond contrastive learning, self-reconstructing is another approach. GraphMVP reconstructs one representation from the other between molecular graphs and 3D molecular representations. Similarly, PanGu uses molecular graphs to reconstruct the linear representation of SELFIES [36, 37], which is similar to SMILES.

Unfortunately, these works lack exploration of multimodal atomic-level approaches and also ignore the use of existing trained, powerful, unimodal models. For instance, DMP is a typical work that integrates SMILES and molecule graphs, using atomic-level masking methods and dual-view molecular-level contrastive learning on two single modalities. However, DMP trains two encoders from scratch, neglecting the use of the current powerful single modal encoders, resulting in the need for a large amount of data: 110M of training data, which is about 4.4 times our training data and consumes significant computational resources.

Therefore, our approach aims to propose a multi-granularity molecular representation fusion learning, utilizing atomic-level alignment to enhance the complementary information between different molecular representations.

Contrastive learning [15] is a typically molecular-level method for multimodal fusion, labeling different representations of the same molecule as 1 and those of different molecules as 0. This approach implies that any two different molecules are entirely dissimilar [38, 39]. However, as shown in Figure 2, although Aspirin and Paracetamol are different molecules, they share similar properties, such as solubility in water and ethanol and analgesic effects. To better capture molecular similarities, we introduce MolSim, which uses continuous similarity measures instead of traditional binary labels. MolSim thus captures subtle molecular relationships more accurately.

For accurate similarity measurement, we utilize Morgan fingerprints [40] to compute the Tanimoto coefficient [41], which captures detailed molecular information, including atom types, bond types, and functional groups. The coefficient effectively reflects structural similarity, providing a reasonable measure of molecular relationships.

In the MolSim component, we replace binary labels with continuous similarity measures by calculating the Tanimoto coefficient between Morgan fingerprints, as shown in Figure 3 (a). This continuous similarity measure allows MolSim to learn more effectively and leverage complementary information.

The similarity matrix $T$, calculated by using the Tanimoto coefficient between Morgan fingerprints, as mentioned above, serves as the target labels for learning. Given the features from SMILES and molecule graphs, the computed similarity matrix $Q$ is computed as:

where $\tau^{-1}$ is the inverse of the temperature parameter $\tau$, and $S$ and $G$ represent the feature matrices for SMILES and molecule graphs, respectively.

In MolSim, we use the mean squared error to calculate the loss between model predictions and the molecular similarity matrix $T$:

where $N$ is the number of samples in the matrix.

In addition to learning molecular-level representations, we introduce AtomAlign for atomic-level fusion. Through atomic alignment, the model can align identical information contained in different molecular representations and obtain unique molecular information from each representation. This is consistent with the Complementary Information Assumption. In AtomAlign, we employ atomic-level masking to achieve alignment of atomic features between different molecular representations, as illustrated in Figure 3 (b).

In AtomAlign, we randomly mask atoms in SMILES and encode the masked SMILES. We then subtract the masked SMILES encodings from the molecular graph encodings of the same molecule. The resulting encodings are passed through a linear layer to predict both the types of masked atoms in SMILES and identify which atoms are not masked without predicting their types.

To align atomic features between different molecular representations, we first calculate the difference between the molecular graph encodings $\mathcal{T}$ and the masked SMILES encodings $\mathcal{E}_{\text{mask}}$. The resulting difference vector $\mathcal{D}$ is given by:

Our loss function integrates the prediction of masked atoms and the identification of unmasked atoms. The total loss is computed as follows:

where $\mathcal{D}_{i}$ represents the difference vector at the $i^{\text{th}}$ position, $\mathcal{I}_{i}^{\text{mask}}$ is the index label of the masked atom at position $i$ in the SMILES representation. $\mathcal{M}$ represents the masking indicator matrix, where $\mathcal{M}_{i}$ is the 0/1 label indicating whether the atom at the $i^{\text{th}}$ position is masked, and $\alpha$ is a weighting parameter that balances the contributions of the masked and unmasked loss components, with $0\leq\alpha\leq 1$.

These two components are trained synchronously, therefore:

where $\beta$ is a hyperparameter.

Our method design offers several advantages: (1) Integrating molecular-level and atomic-level fusion learning effectively leverages the complementary information from different molecular representations. (2) Our method significantly reduces data and computational resource consumption by efficiently using the existing powerful unimodal encoders. (3) Due to the accessibility of SMILES and molecular graph data, even if only SMILES or molecular graph is provided as input in downstream tasks, we can use the RDKit tool [42] to obtain another modal representation of molecules, thereby jointly improving model performance, as shown in Figure 3 (c). (4) Our design is notably lightweight, introducing only a single linear layer in the AtomAlign task, which avoids the inclusion of excessive model parameters and reduces training costs.

We validate our proposed fusion method on drug property prediction tasks, and the experimental results are shown in Table I. Our fusion method achieves significant performance improvements.

Unlike previous work [8, 9, 10], our method fully utilizes two encoders of different modalities. Previously, in downstream tasks, only a single encoder was used, which is insufficient for utilizing molecular multimodal complementary information: using only one encoder means that there is only one modality to supplement the information of another modality without the reverse information. Therefore, our work investigates the complementary effect of information between two molecular modalities. We evaluate the SMILES encoder only, molecule graph encoder only, EWA, and CCO under the conditions of no-train, contrastive learning, DMP, and our proposed method, as shown in Table I.

The results in Table I show that, among all the compared fusion methods, our approach achieves the most significant improvement by aggregating multimodal representations, as opposed to relying solely on single-modal representation. For example, in the BBBP dataset, compared to the maximum performance of the encoder-only method, our method shows an improvement of 4.64% using aggregation operations, while the no-train condition results in a decrease of 1.60%, contrastive learning achieves an improvement of 2.93%, and DMP shows a minor increase of 0.12%. This clearly demonstrates the effectiveness of our method in leveraging complementary multimodal information. Compared to our method, the effectiveness of other fusion methods decreases to varying degrees. Comparing DMP with our method, we can analyze that multimodal atomic-level methods can better integrate molecular multimodal information than single-mode atomic-level methods.

In order to better understand the roles of the two components in MolFusion, we conduct ablation experiments. We run MolSim and AtomAlign components separately and combine them with atomic-level masked learning model and molecular-level contrastive learning, respectively, to form a multi-granularity method. We select SIDER, BBBP, Tox21, and Toxcast as the ablation experimental datasets and record the maximum experimental result of four aggregation operations as the result of each fusion method. The results demonstrate the effectiveness of two components in our method.

As shown in Figure 5, the experimental results show that our method is optimal among many combinations of molecular-level and atomic-level methods. Neither MolSim nor AtomAlign components alone can effectively enable the model to learn the complementary information from molecular modalities. For instance, in the SIDER dataset, MolSim alone achieves a ROC-AUC of 54.23% and AtomAlign 50.21%, but combined in our method, the ROC-AUC improves to 56.09%. In addition, whether we replace the molecular-level component in our method with contrastive learning or replace the atomic-level component in our method with a single-modal masked learning model, the effect decreases to varying degrees. For example, in the Tox21 dataset, replacing the molecular-level and atomic-level components in our approach with single-modal methods results in a performance drop of 12.09% and 6.98%, respectively. These prove the rationality and effectiveness of the two components in MolFusion, confirming that multi-granularity fusion method specifically designed for multimodal data can more effectively learn complementary information between different modalities.

To visually demonstrate the effectiveness of our proposed fusion method in integrating two molecular representations, we use t-SNE for dimensionality reduction visualization. We select 500 molecular vectors from the ZINC dataset and visualize them under three conditions: no-train, contrastive learning, and our method. The results validate the Complementary Information Assumption.

As shown in Figure 4, our method results in partially overlapping vector spaces, consistent with Complementary Information Assumption. This leads to the most significant improvement when fusing representations of the two modalities.