MolX: A Geometric Foundation Model for Protein-Ligand Modelling

MolX is a scalable, interpretable E(3)-equivariant Graph Transformer foundation model that jointly learns geometric and chemical representations of protein-ligand interactions from over 3 million pockets and 5 million molecules, achieving state-of-the-art performance across diverse drug discovery benchmarks through a hybrid pretraining paradigm.

The Gist

Imagine you are trying to figure out why a specific key fits perfectly into a specific lock. In the world of medicine, the "lock" is a protein inside your body, and the "key" is a drug molecule. If they fit together just right, the drug can fix a problem (like killing a cancer cell). If they don't fit, the drug does nothing.

For a long time, computer programs trying to predict this "fit" have been like people looking at a 2D drawing of the key and a 2D drawing of the lock. They can see the shapes, but they miss the crucial 3D depth, the angles, and how the materials feel against each other.

Enter MolX, a new AI model described in this paper. Think of MolX as a super-smart 3D sculptor that has spent years studying millions of real-life keys and locks to understand exactly how they interact.

Here is a simple breakdown of how MolX works and why it's a big deal:

1. The Problem: The "Decoupled" Mistake

Previous AI models often studied the key and the lock separately. It's like trying to guess if a puzzle piece fits by looking at the piece in one hand and the puzzle board in the other, without ever bringing them close together. They missed the subtle "dance" that happens when the two actually touch.

2. The Solution: A Unified 3D Dance Floor

MolX changes the game by looking at the key and the lock together in a 3D space.

• The Analogy: Imagine a dance floor where the protein (lock) and the drug (key) are partners. MolX doesn't just watch them separately; it watches how they move in relation to each other.
• The Magic Trick (E(3)-Equivariance): This is a fancy math term that basically means MolX understands that a lock is the same lock whether you hold it upside down, sideways, or walk around it. It respects the laws of physics and geometry, so it doesn't get confused by the angle of view.

3. How It Learned: The "Denoising" Gym

Before MolX could help doctors, it had to train. The researchers used a clever method to teach it:

• The Game: They took perfect 3D models of keys and locks, then intentionally scrambled them. They moved the atoms around randomly and hid some of the atom types (like covering the teeth of a key with mud).
• The Task: MolX had to look at the scrambled mess and try to rebuild the original, perfect structure.
• The Result: By playing this "fix the mess" game millions of times, MolX learned the deep, hidden rules of how atoms naturally want to sit next to each other. It learned the "physics" of molecules without being explicitly told the rules.

4. The "X-Ray Vision" (Interpretability)

One of the coolest parts of MolX is that it doesn't just give you a "Yes/No" answer; it can explain why.

• The Analogy: Most AI models are "black boxes." You put a drug in, and they spit out a score. You have no idea why they gave that score.
• MolX's Approach: MolX comes with a special tool called a Sparse Autoencoder. Think of this as a high-tech highlighter. When MolX makes a prediction, this tool can highlight exactly which part of the drug molecule and which part of the protein were responsible.
  • Example: It can say, "I predicted this drug works because the red ring on the drug is hugging the blue pocket on the protein." This helps scientists understand the mechanism, not just get a number.

5. Why It Matters

The paper tested MolX on some of the hardest problems in drug discovery, like designing PROTACs (drugs that tag bad proteins for destruction) and Antibody-Drug Conjugates (drugs that deliver a payload to a specific cell).

• The Result: MolX beat all the previous best models. It was more accurate at predicting if a drug would stick, how strong the bond would be, and even the chemical properties of the molecule.
• The Impact: This means scientists can use MolX to screen millions of potential drugs faster and more accurately, potentially speeding up the discovery of new cures for diseases.

In a Nutshell

MolX is a new AI that learns to understand drugs and proteins by studying their 3D shapes together, rather than separately. It trains by fixing scrambled 3D models, and it can explain its own reasoning by highlighting the specific parts of the molecule that matter. It's like upgrading from a flat map to a full 3D GPS system for drug discovery.

Technical Summary

1. Problem Statement

Structure-based drug discovery relies on accurately modeling the interactions between small molecules (ligands) and protein binding pockets. Existing computational approaches suffer from two primary limitations:

• Decoupled Representations: Many models encode proteins and ligands separately or rely on simplified structural representations that fail to explicitly model cross-entity spatial relationships.
• Lack of Geometric Awareness: Sequence-based methods (using SMILES or amino acid sequences) omit 3D structural information, while many existing 3D models treat components independently or focus only on local atomic geometry, missing the higher-order interaction patterns essential for complex tasks like degradation (e.g., PROTACs).
• Interpretability Gap: Current foundation models often act as "black boxes," lacking mechanisms to decompose latent representations into interpretable biological components to explain why a prediction was made.

2. Methodology: The MolX Framework

MolX is a Graph Transformer foundation model designed to jointly learn geometric and chemical representations of protein pockets and ligands from large-scale 3D structural data.

Architecture

• Input Representation: Both protein pockets and ligand molecules are represented as 3D graphs, where nodes are atoms and edges represent chemical bonds.
• E(3)-Equivariant Graph Transformer: The core architecture employs dual E(3)-equivariant graph Transformer encoders. This ensures that the learned representations are invariant to rotation, translation, and reflection, preserving spatial geometry and chemical context.
• Attention Mechanism: The model integrates three types of encodings into the attention mechanism:
  1. Spatial Encoding: Uses pairwise Euclidean distances to modulate attention weights, ensuring the model prioritizes geometrically relevant local interactions while capturing global dependencies.
  2. Edge Encoding: Represents chemical bond types.
  3. Centrality Encoding: Captures node importance based on graph topology.
• Joint Learning: Unlike previous models that process entities separately, MolX updates pocket and ligand representations simultaneously through mutual attention, allowing for the modeling of interaction-specific patterns.

Pretraining Strategy

MolX utilizes a hybrid learning paradigm combining supervised and self-supervised objectives on a dataset of over 3 million protein pockets and 5 million molecules:

• Supervised Objectives:
  • Regression of physicochemical properties: LogP (lipophilicity) and HOMO–LUMO energy gap.
• Self-Supervised Objectives:
  • Coordinate Reconstruction: Randomly perturbs (noises) atomic 3D coordinates and trains the model to recover the original positions.
  • Atom-Type Prediction: Randomly masks atom types and trains the model to predict the correct identity.

Interpretability Module

To address the "black box" issue, MolX integrates a Sparse Autoencoder (SAE).

• This module decomposes the dense latent representations from the Transformer layers into a sparse set of interpretable activation features.
• These features are mapped to specific protein regions (e.g., binding domains, E3 ligase interfaces) and molecular substructures (e.g., aromatic rings, functional groups), creating a "feature dictionary" that links neural activations to biological motifs.

3. Key Contributions

1. Unified Geometric Foundation Model: MolX is the first foundation model to jointly encode protein pockets and ligands as 3D graphs within an E(3)-equivariant framework, capturing interface-level geometric constraints arising from co-organization.
2. Hybrid Pretraining Paradigm: The combination of supervised biochemical regression with self-supervised 3D coordinate reconstruction and atom masking fosters highly generalizable, structure-informed representations.
3. Mechanistic Interpretability: The integration of a sparse autoencoder allows for the decomposition of predictions into specific, interpretable biological and chemical components, revealing the driving forces behind model decisions (e.g., specific E3-target interactions in PROTACs).
4. Spatial Bias Mechanism: The introduction of a spatial positional bias in the Transformer attention mechanism explicitly encodes 3D geometric dependencies, acting as a geometry-aware gating mechanism that aligns model saliency with physical binding determinants.

4. Results

MolX was evaluated across eight downstream benchmarks, including classification (ADC, PROTAC, Molecular Glue, LIT-PCBA) and regression (Binding Affinity, Physicochemical properties).

• Classification Performance:

  • PROTAC: Achieved an AUC of 0.9211, significantly outperforming the best baseline (MolE at 0.700) by +22.1 percentage points.
  • Molecular Glue: Reached near-saturation performance with an AUC of 0.9962.
  • ADC: Achieved an AUC of 0.9807, outperforming MolE by +9.7 percentage points.
  • Robustness: Consistently outperformed baselines across fine-grained subsets (e.g., specific target-E3 pairs), demonstrating strong generalization even in data-scarce scenarios.

• Regression Performance:

  • Binding Affinity (PDBbind): Achieved the lowest Mean Absolute Error (MAE) and Root Mean Square Error (RMSE) for $K_d$, $K_i$, and $IC_{50}$ predictions. For example, reduced $K_d$ RMSE to 1.5043 (vs. 1.5504 for MolE).
  • Physicochemical Properties (MISATO): Set new state-of-the-art results for Electron Affinity, Electronegativity, and Hardness, with MAE reductions of up to 69.9% compared to FradNMI.

• Ablation Studies:

  • Removing the 3D coordinate noising objective caused the most significant performance drop, confirming its critical role in learning geometric priors.
  • Removing spatial bias degraded performance, proving the necessity of explicit 3D distance encoding in the attention mechanism.

• Interpretability Validation:

  • The SAE successfully identified known binding motifs (e.g., VHL and CRBN interfaces) and chemically relevant substructures (e.g., aromatic rings in ligands).
  • Counterfactual analysis showed that modifying high-activation substructures significantly altered predictions, validating the model's sensitivity to functionally relevant chemical motifs.

5. Significance

MolX establishes a new paradigm for molecular representation learning by bridging the gap between geometric deep learning and biological interpretability.

• For Drug Discovery: It provides a scalable, unified framework for predicting complex small-molecule-protein interactions, particularly for challenging modalities like PROTACs and ADCs where 3D geometry is critical.
• For AI in Science: It demonstrates that foundation models can be made interpretable without sacrificing performance, offering mechanistic insights into how models learn interaction rules.
• Future Impact: By providing a pre-trained model that understands both the topology and the continuous 3D manifold of molecular structures, MolX accelerates the design of novel therapeutics and reduces reliance on trial-and-error experimental screening.