Self-Supervised Pretraining of Molecular Graph Encoders with LeJEPA

Technical Summary: Self-Supervised Pretraining of Molecular Graph Encoders with LeJEPA

Problem Statement
While self-supervised pretraining has revolutionized language and vision models, its utility for molecular graph neural networks (GNNs) remains contested. In chemistry, labeled data is scarce and expensive, yet the space of synthesizable molecules is vast. The prevailing "pretrain-then-finetune" paradigm, successful in other domains, has yielded mixed results in molecular property prediction. Existing self-supervised objectives (contrastive, generative, and predictive) often struggle to provide robust downstream gains over well-tuned supervised baselines using hand-crafted descriptors. This work investigates whether pretraining a molecular graph encoder on a large unlabelled corpus using a specific architecture, LeJEPA, improves downstream property prediction, and if so, how that benefit manifests.

Methodology
The authors adapt LeJEPA (Latent-Euclidean Joint Embedding Predictive Architecture) to molecular graphs. Unlike traditional JEPAs that rely on asymmetric architectures (momentum teachers, stop-gradients, and predictor networks), LeJEPA employs a predictor-free design regularized by Sketched Isotropic Gaussian Regularisation (SIGReg).

• Architecture: The encoder maps a molecule graph $G$ to a $K$-dimensional embedding $z$. Two backbone architectures were evaluated:
  1. GPS Graph Transformer: A hybrid model combining local GINE message passing with global multi-head self-attention, utilizing random-walk positional encodings.
  2. Chemprop-style D-MPNN: A directed message-passing network.
• Pretraining Objective: The model is trained on $\sim$2.9 million bioactivity-curated molecules from ChEMBL. The objective minimizes a weighted sum of two terms:
  1. Prediction Loss ($L_{pred}$): For each molecule, multiple views are generated (2 global views with light corruption and 8 local views via BFS subgraphs). The encoder must predict the centroid of the global-view embeddings from the local-view embeddings. Crucially, physico-chemical descriptors (RDKit) are attached only to global views to prevent trivial solutions where the model ignores structure.
  2. SIGReg Loss ($L_{SIGReg}$): This regularizer prevents representation collapse by driving the batch embedding distribution toward an isotropic Gaussian $N(0, I)$. It uses the Epps–Pulley statistic to compare the empirical characteristic function of 1D projections of the embeddings against a standard normal distribution across 1024 random directions.
• Evaluation Protocol: The pretrained encoders were evaluated on two tasks under rigorous, multi-seed, bootstrap-based protocols:
  1. Antibiotic Activity: A dataset of 39,312 compounds against Staphylococcus aureus (Wong et al.), using both random and scaffold splits (Bemis-Murcko).
  2. ogbg-molhiv: A standard benchmark for HIV replication inhibition with a fixed scaffold split.
  • Probes: Evaluation included "frozen probes" (fitting a head on fixed embeddings) and full "finetuning" (updating encoder weights). Comparisons were made against random initialization and hand-crafted baselines (RDKit descriptors, Morgan fingerprints).

Key Results

1. Representational Quality vs. Finetuning:

   • Frozen Probes: Pretraining yields a significantly better representation. On ogbg-molhiv, a frozen LeJEPA embedding achieved a ROC-AUC of 0.788 vs. 0.665 for a random-initialized encoder (a +0.123 lift), reaching the band of published self-supervised GNNs. On the antibiotic task, the frozen pretrained embedding (AUPRC 0.159) outperformed both random initialization (0.082) and RDKit descriptors alone (0.112).
   • Finetuning: This representational advantage does not robustly translate to downstream finetuning gains.
     • On the antibiotic random split, finetuning showed no gain over scratch training.
     • On the antibiotic scaffold split, a single canonical partition showed a significant gain (+0.041 AUPRC), but this effect vanished when replicated across five different scaffold partitions (pooled $\Delta$ = +0.013, $p=0.095$).
     • On ogbg-molhiv and with the D-MPNN backbone, finetuning gains were null.

2. Complementarity with Fingerprints:

   • The authors found that the pretrained embedding is complementary to hand-crafted Morgan fingerprints, but only when the embedding's informative subspace is isolated.
   • The embedding saturates at 16–32 effective dimensions, whereas fingerprints improve up to 1024 bits.
   • At matched dimensionality (128), the fingerprint outperformed the embedding on validation splits but underperformed on shifted test scaffolds. The embedding showed superior robustness to scaffold shifts.
   • Feature-Level Combination: Truncating the pretrained embedding to 32 dimensions (via PCA) and concatenating it with a 1024-bit Morgan fingerprint increased the ogbg-molhiv test ROC-AUC from 0.803 to 0.830 ($\Delta$ = +0.027, $p=0.014$).
   • Crucially, this gain was specific to the pretrained embedding; concatenating an untrained (random) embedding provided no benefit, confirming the gain stems from the learned structure, not merely added dimensions.

3. Baseline Performance:

   • On the antibiotic task, a simple tree ensemble over 217 RDKit descriptors (without any graph network) matched or exceeded the performance of a graph-based Chemprop model and the published baseline, suggesting graph networks are not strictly necessary for this specific task.

Significance and Claims
The paper concludes that LeJEPA pretraining produces a measurably better molecular representation that captures physico-chemical structure and generalizes well to novel chemotypes (as evidenced by frozen probes and scaffold-shift robustness). However, the authors argue that this benefit is not realized through standard finetuning, which often fails to exploit the pretrained weights better than random initialization or hand-crafted features.

Instead, the value of the pretraining is realized through feature-level combination. The pretrained embedding provides robustness to distribution shifts (scaffold changes) that hand-crafted fingerprints lack, but only when the embedding is truncated to its informative subspace and combined with established representations. The authors position their work as a mechanistic study demonstrating that while the representation is superior, its downstream utility is conditional on how it is integrated, rather than a universal replacement for supervised baselines. The work emphasizes that the "gain" is often invisible on validation sets and only apparent under distribution shifts, making it difficult to select for via standard hyperparameter tuning.

Limitations Acknowledged
The authors note that the finetuning evidence is the weakest part of the study, as the single significant scaffold-split result did not replicate across multiple partitions. They also acknowledge that their compact encoder does not achieve state-of-the-art absolute scores on benchmarks compared to larger, supervised models, and that the pretraining corpus does not fully cover the downstream chemical space, capping the potential transfer ceiling.