Self-Supervised Pretraining of Molecular Graph Encoders with LeJEPA

1.1 Pretraining for molecular graph neural networks

• Molecular GNN pretraining targets scarce labels, but evidence is mixed across objectives and endpoints, and supervised descriptor baselines can be difficult to beat.Contrastive methods require designed augmentations and many negatives; generative methods reconstruct input details that may not help downstream.Source: §1.1 · pp. 1–2

1.2 World models and LeJEPA

• JEPAs predict held-out views in latent space, while generative models reconstruct inputs and contrastive methods rely on negative samples [7].The motivation is to predict useful features without reconstructing incidental detail or requiring large numbers of negatives [7].Source: §1.2 · p. 2
• LeJEPA [10] replaces a JEPA’s teacher, stop-gradient and predictor with SIGReg, and this paper applies it to molecular graphs.SIGReg drives embeddings toward an isotropic Gaussian; the paper notes LeJEPA had previously been demonstrated on image data.Source: §1.2 · p. 2

1.3 Antibiotics and antimicrobial resistance

• Antibiotic activity is highly imbalanced, with about 1% of screened compounds active, making it a low-label transfer setting.The Wong et al. [1] benchmark contains 39,312 compounds with four binary endpoints.Source: §1.3 · p. 2

1.4 A standard benchmark: ogbg-molhiv

• ogbg-molhiv provides a standard benchmark of about 41,000 molecules, a fixed scaffold split and published scratch-versus-pretrained baselines.It evaluates HIV-replication inhibition using ROC-AUC.Source: §1.4 · p. 3

2.1 Overview

• The encoder maps each molecular graph to a K-dimensional embedding and learns from multiple views of each molecule.The objective combines a prediction loss with SIGReg, weighted by a single scalar λ set to 0.2.Source: §2.1 · p. 3
• The prediction loss makes views of the same molecule agree, while SIGReg keeps the embedding distribution full-rank and well-spread.Together, these terms make the objective informative and non-degenerate.Source: §2.1 · p. 3
• LeJEPA uses no predictor network, momentum teacher or stop-gradients; gradients flow through the view embeddings and prediction target.The loss is a convex combination of prediction and distributional regularization.Source: §2.1 · p. 3

2.2 View generation

• Each molecule supplies two global views and eight local views for the prediction task.The global views are lightly corrupted whole molecules; the local views are connected BFS subgraphs covering about 60% of the atoms.Source: §2.2 · p. 3
• Two global views retain nearly the whole molecule: one masks features on 15% of atoms, while the other drops 10% of bonds.Their embeddings serve as a stable reference because the corruptions preserve molecule identity.Source: §2.2 · p. 3
• Eight connected BFS local views cover 60% of a molecule’s atoms and must predict a near-whole-molecule reference.This fragment-to-reference task is intended to make the encoder learn how local substructure relates to molecular identity.Source: §2.2 · p. 3, Figure 1 · p. 4
• The 217 z-scored RDKit descriptors are attached only to global views, leaving local views to be explained from molecular structure.Giving descriptors to every view could let the encoder match views using their shared descriptor signal instead of learning structure.Source: §2.2 · p. 3

2.3 Prediction loss

• The target is the centroid of the global-view embeddings, and every view embedding is pulled toward that centroid.The near-whole-molecule global views provide an informative target, while local views must infer it from fragments.Source: §2.3 · p. 4
• The prediction target aggregates the global views rather than using an individual view embedding.The method defines the target as their centroid.Source: §2.3 · p. 4
• Gradients flow through both the centroid and each view embedding, with no separate predictor network or stop-gradient.Because the global views carry auxiliary descriptors, the bulk of the learning signal comes from local-view terms.Source: §2.3 · p. 4

2.4 SIGReg loss

• SIGReg prevents the prediction objective’s trivial constant-embedding solution by driving embeddings toward a full-rank isotropic Gaussian.The target has unit variance in every direction and is uncorrelated across dimensions.Source: §2.4 · p. 5
• For each of 1,024 random unit directions, SIGReg compares projected embeddings’ empirical characteristic function with a standard normal using an Epps–Pulley statistic.The statistic uses a Gaussian window and is evaluated over t ∈ [−5, 5].Source: §2.4 · p. 5
• By the Cramér–Wold theorem, matching every projection to N(0, 1) matches the joint distribution to N(0, I), preventing collapse and keeping representations well-conditioned.The effective rank rises toward K as SIGReg takes effect.Source: §2.4 · p. 5
• The per-direction statistics are averaged across projections, and the objective is optimised end-to-end with λ = 0.2.Under distributed training, empirical characteristic functions are averaged across replicas for the global batch.Source: §2.4 · p. 5

2.5 Encoder architectures

• The objective is encoder-agnostic, requiring only a mapping from a graph to a fixed-size embedding.The study compares two backbones.Source: §2.5 · p. 5
• GPS combines bond-aware GINE message passing with global self-attention, then mean-pools atom embeddings and projects them to K = 128 dimensions.Its six-layer, four-head configuration has hidden dimension 128 and ∼2M parameters; random-walk positional encodings break message-passing symmetries.Source: §2.5 · p. 5
• The Chemprop-style D-MPNN passes messages along directed bonds, excludes each edge’s reverse edge from its message, and pools final atom states into a graph embedding.The architecture control uses hidden width 512 and depth 5 to match capacity at which it is competitive on supervised tasks.Source: §2.5 · p. 5

2.6 Pretraining corpus and optimisation

• ChEMBL contains 90.8% of ogbg-molhiv training scaffolds and 44.7% of test scaffolds, with median test-molecule nearest-neighbour Tanimoto similarity 0.43.About 55% of test scaffolds are absent from the corpus.Source: §2.6 · p. 6
• The ∼2.9M-molecule ChEMBL corpus was chosen over a larger enumerated library because it covers substantially more downstream scaffolds relevant to the tasks.Its coverage matches or exceeds that of the antibiotic task, supporting interpretation of null transfer as a method result rather than corpus mismatch.Source: §2.6 · p. 6
• Pretraining uses Adam with λ = 0.2, Vg = 2, Vl = 8, Ns = 1024 random projections, no learning-rate warmup, and four-GPU distributed training.Code, configurations, and final checkpoints are released for reproducibility.Source: §2.6 · p. 6

2.7 Qualitative effect of pretraining

• Pretrained embeddings form a well-spread manifold with smooth gradients in molecular weight, lipophilicity, and ring count.The projections use pure graph embeddings, with auxiliary RDKit descriptors zeroed so they do not directly organise either embedding.Source: §2.7 · p. 6
• Random-initialisation embeddings recover only a coarse molecular-size axis, consistent with atom count correlating with molecular weight.They lack the organised lipophilicity and ring-count gradients seen in the pretrained manifold.Source: §2.7 · p. 6
• The pretrained manifold’s organised property gradients are consistent with SIGReg spreading the representation into an isotropic, chemistry-organised distribution.Figure 3 describes well-trained embeddings as having standard-normal projections and T → 0.Source: §2.7 · p. 6, Figure 3 · p. 7

• The same pretrained checkpoint is compared with the same architecture trained from scratch, making pretraining weights the only difference between arms.Source: §2.8 · p. 7
• The antibiotic task uses 39,312 compounds, including 512 actives, and tests both a stratified random split and leakage-free scaffold splits.The scaffold evaluation includes the canonical Bemis-Murcko partition and four additional partitions; AUPRC is the metric for this imbalanced task.Source: §2.8 · p. 7
• Random-split antibiotic predictions are pooled across five independent splits and compared head-to-head with paired bootstrap resampling.Models are five-seed ensembles; the protocol uses 10,000 resamples and one-sided p-values, with published and retrained Chemprop comparisons.Source: §2.8 · p. 9
• ogbg-molhiv provides a second test on approximately 41,000 molecules with a fixed scaffold split and test ROC-AUC averaged over five seeds.Epochs are selected on validation data, and learning rate and frozen-probe epochs are swept with each model tuned independently.Source: §2.8 · p. 9
• Frozen probes assess encoder representations without updating encoder weights, using logistic regression for antibiotic scaffolds and a random forest for ogbg-molhiv.References are random-initialised encoders and task-specific handcrafted representations.Source: §2.8 · p. 9
• Width sweeps compare fingerprints and truncated embeddings under the same forest, then test their concatenation across fingerprint width, embedding dimensions, and feature-sampling settings.PCA is fitted on training data only; sweep sizes are smaller than the main probe, so absolute values are compared only within each sweep.Source: §2.8 · p. 10

3.1 Descriptor tree-ensembles match the graph model and exceed the published baseline

• The descriptor RandomForest significantly outperforms retrained Chemprop on the random split (∆AUPRC = +0.046, p = 0.006) and is statistically level with it on the scaffold split.The random-split table summarizes pooled test AUPRC across five splits; the scaffold-split RandomForest reaches AUPRC 0.339 versus Chemprop 0.312 (p = 0.17).Source: §3.1 · p. 10, Table 1 · p. 11

3.2 A frozen probe shows that pretraining learns a better representation

• Frozen probes show better pretrained than random-initialised representations on both tasks; on ogbg-molhiv, ROC-AUC is 0.788 ± 0.003 versus 0.665 ± 0.004.The antibiotic scaffold-split embedding also exceeds the random-initialised encoder and descriptors alone under a frozen linear probe.Source: §3.2 · pp. 10–11, Table 2 · p. 11, Table 3 · p. 12

3.3 The representational gain does not robustly convert into a finetuning gain

• Across five leakage-free antibiotic scaffold partitions, the pretrained-versus-scratch finetuning gain is partition-dependent; pooled ∆AUPRC is +0.013 (95% CI [−0.006, +0.032], one-sided p = 0.095), not significant.Finetuning is also null on the random split, ogbg-molhiv, and a D-MPNN backbone.Source: §3.3 · p. 12

3.4 The representation is complementary to a fingerprint

• At matched 128-dimensional width, fingerprints lead validation ROC-AUC (0.799 vs 0.782), but trail on test scaffolds (0.759 vs 0.788); embeddings saturate by roughly 16–32 dimensions.Fingerprints lose 0.040 ROC-AUC across the scaffold shift, while embeddings lose none.Source: §3.4 · pp. 13–14, Figure 7 · p. 16
• A 1024-bit fingerprint plus a 32-dimensional pretrained embedding reaches test ROC-AUC 0.830±0.004 versus 0.803±0.004 for the fingerprint baseline, with paired ∆= +0.027 (p = 0.014).The gain appears on shifted test scaffolds, not validation; restoring diluting dimensions to k = 64 removes most of the effect.Source: §3.4 · pp. 14–15, Table 6 · p. 15

3.5 Positioning against published baselines

• The paper positions its molhiv result as a controlled internal comparison: its compact backbone is below published self-supervised absolute scores, while the frozen probe reaches that band.The strongest measured single representation remains a hand-engineered fingerprint; the combination narrows the gap to the leaderboard fingerprint ensemble.Source: §3.2 · p. 11, §3.4 · p. 14, §3.5 · p. 16

4 Discussion and conclusion

• 0.788 vs 0.665 ROC-AUC under a frozen random-forest probe on ogbg-molhiv shows a +0.123 pretrained-representation lift; antibiotic frozen probes also favor pretraining, 0.159 vs 0.082 AUPRC.The authors describe these frozen-probe results as reproducible representation improvements, distinct from finetuning outcomes.Source: §4 · p. 17
• The canonical antibiotic scaffold split gains +0.041 AUPRC, but five partitions yield a pooled +0.013 estimate (95% CI [−0.006, +0.032], p = 0.095), while other finetuning comparisons are null.The D-MPNN canonical-split comparison is also null (∆AUPRC = −0.000, 95% CI [−0.045, +0.046], p = 0.48).Source: §4 · p. 17
• Truncating the embedding and concatenating it with a 1024-bit fingerprint raises ogbg-molhiv test ROC-AUC from 0.803 to 0.830 (∆= +0.027, 95% CI [+0.003, +0.054], p = 0.014).The gain disappears with the k = 64 dilution control, and an untrained backbone gains nothing (∆= −0.003, p = 0.59).Source: §4 · p. 17
• The combined representation is nearly tied with the fingerprint on validation (0.847 vs 0.845) but loses less across the scaffold shift (−0.017 vs −0.042), supporting a distribution-shift-specific advantage.The authors interpret the result as robustness to chemotype change, not an improvement on data drawn like training data.Source: §4 · p. 17
• The concatenation result rests on one benchmark, one split, and 130 test actives, with its truncation dimension selected on validation.The authors also note that their compact encoder scores below the strongest published self-supervised graph networks and that the pretraining corpus incompletely covers downstream chemical space.Source: §4 · p. 17
• The authors favor complementary auxiliary signals and shift-focused evaluation over simply enlarging an encoder that does not fill its available dimensions.They also recommend treating pretrained embeddings as components alongside established representations rather than replacements.Source: §4 · p. 17

Declaration of generative AI use

• The authors used Claude for code, statistical scripts, and manuscript preparation, but not to generate experimental data.They state that the research questions, experimental designs, and interpretations were their own; they supervised runs, checked reported numbers against experimental outputs, and take responsibility for the article.Source: Declaration of generative AI use · p. 18