SELFormerMM: multimodal molecular representation learning via SELFIES, structure, text, and knowledge graph integration

The paper introduces SELFormerMM, a multimodal framework that integrates SELFIES, structural graphs, textual descriptions, and knowledge graph data to generate superior molecular representations for drug discovery, outperforming existing unimodal models across various property prediction tasks.

The Gist

Imagine you are trying to describe a complex recipe to a friend so they can cook it perfectly. If you only give them the list of ingredients (the sequence), they might miss the cooking technique. If you only give them a photo of the finished dish (the structure), they won't know the order of steps. If you only give them a food critic's review (text), they won't know the exact measurements. And if you only give them a map of how this dish interacts with other foods (knowledge graph), they won't know what the dish actually is.

To truly understand the recipe, you need all four of these perspectives combined.

This is exactly the problem scientists face when trying to teach computers to understand molecules (the tiny building blocks of drugs). For a long time, computer models could only look at one thing at a time, like just the ingredient list or just the photo. This paper introduces a new AI model called SELFormerMM that acts like a "super-chef" who can look at all four perspectives simultaneously to understand a molecule better than ever before.

Here is a simple breakdown of how it works:

1. The Four "Languages" of Molecules

The researchers realized that molecules speak four different "languages," and previous AI models only spoke one. SELFormerMM learns all of them:

• The Recipe List (SELFIES): Molecules are often written as strings of letters (like a chemical sentence). The old way (SMILES) was like writing a recipe with typos that sometimes made the dish explode. The new way (SELFIES) is a "grammar-proof" language that guarantees the recipe always makes sense.
• The Blueprint (Structure): This is a 2D map showing how the atoms are connected, like a wiring diagram for a house.
• The Review (Text): Molecules have descriptions in scientific papers and databases. This is like reading a food critic's article about how a dish tastes or feels.
• The Social Network (Knowledge Graph): Molecules don't exist in a vacuum; they interact with proteins, genes, and diseases. This is like a map showing who the molecule "friends" with in the body (e.g., "This molecule talks to Protein X, which helps cure Disease Y").

2. How SELFormerMM Works: The "Group Hug"

The model uses a clever training trick called Contrastive Learning. Imagine you have four different people describing the same person:

1. A photographer (Structure)
2. A writer (Text)
3. A DNA tester (Sequence)
4. A social worker (Knowledge Graph)

In the past, these four people worked in separate rooms. SELFormerMM puts them in the same room and says, "You are all describing the same molecule! You need to agree on what it looks like."

The AI forces these four different views to "hug" each other in a shared mental space. If the "writer" says the molecule is "toxic" and the "social worker" says it "attacks a specific protein," the AI learns to connect those dots. By forcing these different views to align, the AI builds a much richer, 3D understanding of the molecule than any single view could provide.

3. The Results: A Better Drug Detective

The researchers tested this new "super-chef" on a massive dataset of 3 million molecules and then asked it to predict real-world drug behaviors, such as:

• Can this drug cross the blood-brain barrier? (Like asking, "Can this delivery truck get through the city gates?")
• Is this drug toxic? (Like asking, "Will this dish make the customer sick?")
• How well does it stick to a protein? (Like asking, "How strong is the magnet?")

The verdict: SELFormerMM beat almost every other model that only looked at one or two of these languages.

• It was particularly good at predicting side effects and brain penetration.
• It proved that when you combine the "recipe," the "blueprint," the "review," and the "social network," you get a much smarter prediction than if you just looked at the blueprint alone.

4. Why This Matters

Think of drug discovery as finding a needle in a haystack.

• Old AI: Looked at the needle from one angle. It might miss the needle if it's hidden from that specific angle.
• SELFormerMM: Spins the needle, looks at it under a microscope, reads the label on it, and checks its history. It finds the needle faster and with more confidence.

The Bottom Line

This paper presents a new tool that teaches computers to understand molecules the way a human expert does: by combining chemistry, structure, language, and biology all at once. It's a step toward faster, cheaper, and safer drug discovery, helping scientists find life-saving medicines without wasting time on dead ends.

The best part? The creators made this tool free for everyone to use, so other scientists can start building better medicines with it today.

Technical Summary

1. Problem Statement

Molecular representation learning is fundamental to computational drug discovery, yet current state-of-the-art models predominantly rely on unimodal inputs (e.g., only molecular sequences like SMILES/SELFIES, or only structural graphs).

• Limitations of Unimodal Approaches: These models capture only a limited aspect of molecular behavior. For instance, sequence models may miss topological constraints, while graph models may lack semantic context.
• The Challenge: Unifying heterogeneous modalities—specifically molecular sequences (SELFIES), structural graphs, natural language descriptions, and biological interaction networks (Knowledge Graphs)—into a coherent framework is non-trivial. Existing multimodal approaches often focus on limited subsets of these modalities or fail to effectively align them in a shared latent space, hindering the creation of rich, biologically grounded representations.

2. Methodology

The authors propose SELFormerMM, a unified multimodal framework designed to learn shared molecular embeddings by integrating four distinct data modalities.

A. Data Construction

The model is pre-trained on a large-scale dataset comprising approximately 3 million molecules, aggregated from:

• Sequence: ChEMBL v36 (SMILES converted to SELFIES to ensure 100% syntactic validity).
• Text: M3-20M dataset (natural language descriptions derived from PubChem, physicochemical properties, and GPT-3.5 summaries).
• Knowledge Graph (KG): CROssBARv2-KG, a subgraph containing compounds, proteins, genes, and drugs with interactions (CTI, DTI, PPI, orthology).
• Structure: 2D molecular graphs derived from SMILES.
• Note: A multi-stage mapping procedure (InChIKey, exact SMILES, canonical SMILES) was used to align molecules across these disparate sources.

B. Model Architecture

SELFormerMM consists of four modality-specific encoders and a shared projection space:

1. Sequence Encoder: Uses SELFormer (a RoBERTa-based transformer) trained on SELFIES tokens.
2. Text Encoder: Uses SciBERT (pretrained on scientific text) to encode molecular descriptions.
3. Structure Encoder: Uses Uni-Mol (a 3D molecular pretraining framework) to generate structural embeddings.
4. KG Encoder: Uses DMGI (Deep Mutual Graph Infomax) to encode biological interaction networks.
5. Projection & Alignment:
   • Text, Structure, and KG encoders are initialized from pretrained checkpoints and kept frozen.
   • Each modality passes through a dedicated non-linear projection network (MLP) to map its native embedding space into a shared 768-dimensional hidden space (matching the SELFormer dimension).
   • Missing modalities are handled by feeding zero vectors through the projection networks.

C. Training Strategy

The training occurs in two stages:

1. Multimodal Pretraining (Self-Supervised):
   • Objective: Multi-view, multi-positive Supervised Contrastive Learning using SINCERELoss.
   • Mechanism: The model aligns embeddings of the same molecule across different modalities (positives) while pushing apart embeddings of different molecules (negatives).
   • Scale: Trained for 25 epochs on ~3M molecules using mixed-precision optimization.
2. Supervised Finetuning:
   • The pretrained backbone is adapted to downstream tasks (classification/regression) by attaching task-specific heads.
   • Strategy: Partial freezing is employed (first 9 transformer blocks fixed) to preserve alignment while allowing task-specific refinement in higher layers.

3. Key Contributions

• First Integration of Four Modalities: SELFormerMM is among the first to simultaneously integrate SELFIES sequences, 2D/3D structural graphs, natural language text, and large-scale biological knowledge graphs into a single representation space.
• Robust Sequence Handling: By utilizing SELFIES instead of SMILES, the model guarantees chemically valid tokenization, addressing the syntactic instability of traditional sequence models.
• Scalable Pretraining: The framework is pre-trained on a massive scale (~3M molecules), demonstrating that multimodal learning can be effective even with sparse coverage of text and KG data across the dataset.
• Open Source: The authors released the code, datasets, pretrained models, and embeddings as a programmatic tool.

4. Results

The model was evaluated on the MoleculeNet suite, covering both classification (BBBP, HIV, BACE, Tox21, SIDER) and regression (ESOL, FreeSolv, Lipophilicity, PDBbind) tasks.

• Classification Performance:
  • SELFormerMM achieved state-of-the-art (SOTA) or competitive results on several benchmarks.
  • SIDER (Side Effects): Achieved the best ROC-AUC (0.749), outperforming unimodal baselines and other multimodal models.
  • BBBP (Blood-Brain Barrier): Achieved 0.918 ROC-AUC, ranking second overall and demonstrating strong generalization without relying on hand-crafted fingerprints.
  • HIV & BACE: Showed improvements over the unimodal SELFormer baseline, though performance was slightly limited by the low availability of text/KG data in these specific subsets.
• Regression Performance:
  • Lipophilicity: Achieved the best RMSE (0.574) among multimodal models.
  • PDBbind (Binding Affinity): Improved significantly over the unimodal baseline (RMSE 1.328 vs. 1.437), likely due to the integration of structural and KG context.
  • ESOL: Remained competitive (RMSE 0.603), though slightly lower than the unimodal SELFormer (0.386), suggesting that for solubility, structural signals alone may be dominant.
• Ablation Studies:
  • Modality Contribution: Structural graph information was the most consistent and robust signal. Text and KG modalities provided complementary gains but were sensitive to data sparsity.
  • Finetuning: Finetuning significantly improved performance over the raw pretrained embeddings, validating the transferability of the learned representations.

5. Significance and Future Directions

• Biological Relevance: Case studies (e.g., predicting BBB permeability for dextroamphetamine vs. benserazide) confirmed that SELFormerMM captures biologically meaningful signals consistent with pharmacological knowledge.
• Framework Utility: The study proves that integrating heterogeneous data sources (chemical, structural, semantic, and relational) creates richer representations that outperform unimodal approaches, particularly for complex tasks involving biological context (e.g., side effects, binding affinity).
• Future Work: The authors suggest improving strategies for handling missing modalities (moving beyond zero-padding), training encoders end-to-end rather than freezing them, and curating higher-quality textual data from literature to enhance coverage.

In conclusion, SELFormerMM establishes a new benchmark for multimodal molecular representation learning, offering a scalable and flexible foundation for hypothesis-driven drug discovery.