PaperlayerView exact revision on arXiv

Source-linked AI summary

Structural Hierarchy and Geometry in Molecular Representation Learning

David Sulu, Lorenzo Di Fruscia, Jana M. Weber

arXiv:2608.29886v1cs.LGq-bio.BM

TL;DR

The paper asks whether explicitly supervising molecular embeddings with Bemis–Murcko scaffolds changes learned organization and whether geometry matters. It adds scaffold supervision to molecular contrastive pretraining, compares Euclidean and Lorentz objectives, and evaluates retrieval and property prediction. Scaffold supervision consistently organizes molecules by identical and related scaffolds and often improves prediction, while gains remain task-dependent and Lorentz geometry offers no consistent overall advantage.

  • Problem

    The study addresses whether explicit Bemis–Murcko scaffold supervision changes molecular embedding organization and whether its effects depend on embedding geometry.

  • Method

    The paper adds molecule-to-scaffold supervision to contrastive pretraining, compares Euclidean and Lorentz projection geometries, and evaluates retrieval plus frozen property prediction.

  • Results

    Scaffold supervision consistently organizes molecules by identical and related scaffolds and improves property prediction on several tasks, while gains depend on the property and Lorentz geometry has no consistent overall advantage.

  • Takeaways & Limitations

    Explicitly teaching the relation between molecules and structural cores reliably shapes molecular embedding organization, but the usefulness of that organization remains task dependent.

  • Takeaways & Limitations

    The study is a controlled comparison against matched MolCLR retrainings using only Bemis–Murcko scaffolds, not a comprehensive benchmark of broader structure-aware approaches.

Abstract

from arXiv · show

Molecular self-supervised learning uses chemical structures to guide which molecular embeddings should be similar. We study whether explicitly encoding a molecule's Bemis-Murcko scaffold and using it to supervise the molecular embedding changes what the model learns. We further test whether this effect depends on the embedding geometry by comparing Euclidean and Lorentz contrastive objectives. Across two augmentation strengths, scaffold-supervised models consistently organize molecules according to both identical and structurally related scaffolds. The resulting embeddings also improve molecular property prediction on several tasks, while the exact gains depend on the predicted property. The effect of scaffold supervision on molecular organization is stronger under Lorentz objectives, but neither geometry provides a consistent overall advantage. These results show that explicitly teaching the relation between a molecule and its structural core can reliably shape the organization of molecular embedding space, while the extent of usefulness of this organization remains task dependent.

1 Introduction

The paper motivates scaffold supervision as a way to introduce explicit cross-level structural organization into molecular contrastive learning, then tests its effects across embedding geometries and evaluation dimensions.

  • Motivation: Limited molecular property labels motivate self-supervised learning from large collections of unlabeled structures.Experimental validation is costly, restricting labelled data for property prediction.
  • Related work: Contrastive learning teaches molecular embeddings from similarity by aligning related views and separating negatives.Molecular methods create perturbed graph views through masking, bond changes, or subgraph removal.
  • Scaffold supervision: Molecular cores provide a chemically meaningful shared abstraction for organizing related compound series.Using a common core creates a many-to-one relation between compounds and their structural target.
  • Scaffold supervision: Scaffold supervision adds scaffold-level organization while leaving the molecular encoder and augmentation scheme unchanged.This isolates the effect of the additional supervision.
  • Study scope: The study tests BM scaffolds as explicit many-to-one supervision and compares Euclidean with Lorentz geometry using retrieval and frozen property prediction.These evaluations probe both molecular-neighbourhood organization and downstream transfer.

2 Methods

The method pretrains a graph encoder with molecular-view contrastive learning plus explicit molecule-to-scaffold supervision, while varying projection geometry and evaluating the frozen backbone.

  • Molecular representation: A GINE encodes each molecular graph into reusable backbone embedding h, while a projection head maps h to training embedding z.The contrastive loss acts on z, whereas h is retained for retrieval and property prediction.
  • Contrastive pretraining: Molecular contrastive pretraining aligns two perturbed views of each molecule and separates embeddings of different molecules.The study uses 25% perturbation and a milder 15% setting on a 1M-molecule subset.
  • Scaffold supervision: Scaffold supervision encodes each BM scaffold with the same graph encoder, aligning each molecular view to its own scaffold and contrasting other scaffolds.The total loss combines molecular and scaffold objectives with scaffold weight λs = 0.1.
  • Embedding geometry: Euclidean and Lorentz variants differ only in projected-embedding similarity: cosine similarity versus negative geodesic distance after Lorentz mapping.The Lorentz transformation acts on z, while backbone embedding h remains Euclidean.
  • Evaluation: After pretraining, frozen h is evaluated through molecular retrieval and linear prediction on seven MoleculeNet tasks.Retrieval uses cosine similarity in h, while downstream predictors use fixed scaffold splits.

3 Results

Scaffold supervision reorganizes molecular neighborhoods toward identical and related scaffolds and often improves property prediction, but geometry and task determine the magnitude and direction of gains.

  • Molecular neighbourhoods: Scaffold supervision strongly increases retrieval of molecules sharing the query’s exact BM scaffold.Because exact scaffold identity is supervised directly, this confirms learning of the intended scaffold relation.
  • Molecular neighbourhoods: Excluding exact-scaffold matches, scaffold-supervised models retrieve molecules with more similar nonidentical scaffolds and complete-molecule fingerprints.The effect is positive across seeds, geometries, and augmentation settings.
  • Embedding geometry: The scaffold-supervision gain is larger under Lorentz training at both augmentation strengths, but neither geometry consistently outperforms the other.The larger Lorentz effect is primarily driven by reduced matched MolCLR baseline performance.
  • Embedding geometry: 92.0 ± 2.7% of pairs place the scaffold closer to the origin than its molecule at 15% augmentation, versus 19.3 ± 8.6% at 25%.The mean radial ordering reverses between augmentation strengths, and this ordering alone does not establish improved embedding quality.
  • Property prediction: Scaffold supervision improves the across-seed mean on six of seven Euclidean tasks at 15% augmentation and five of seven at 25%.The effect remains task-dependent: ESOL improves at both Euclidean settings, whereas QM7 worsens on average; Lorentz training also helps most tested tasks.

4 Limitations

The study provides a controlled analysis of BM scaffold supervision relative to matched MolCLR retrainings, not a comprehensive benchmark of structure-aware molecular SSL. Its scope is also constrained by the scaffold definition, dataset scale, scaffold-frequency handling, and Lorentz ablations.

  • The study compares BM scaffold supervision with matched MolCLR retrainings rather than comprehensively benchmarking all structure-aware molecular SSL approaches.It does not claim superiority over broader structure-aware methods, and matched comparisons remain future work.
  • Only Bemis–Murcko scaffolds are considered, so alternative scaffold definitions or shared structural abstractions may produce different organization.
  • Pretraining uses a 1M-molecule subset without scaffold-frequency balancing or capping, and the bounded Lorentz implementation uses one unswept rmax.The magnitude of observed effects may change at larger scale.

5 Conclusion

Scaffold supervision reshapes molecular embedding neighborhoods, while geometry affects the strength of this effect without yielding a consistent overall advantage. The study motivates Lorentz geometry as a candidate for molecular representations but treats its usefulness as empirical.

  • 5 Conclusion: Scaffold supervision brings molecules closer to both identical and related nonidentical Bemis–Murcko scaffolds, while improving property prediction on several tasks.The usefulness of the resulting organization depends on the downstream property.
  • 5 Conclusion: Lorentz geometry changes the strength of scaffold-supervision effects but does not provide a consistent advantage over Euclidean geometry.Hyperbolic geometry is considered because its volume grows exponentially with radius, making it a candidate for tree-like relationships.
  • 5 Conclusion: The total training objective combines molecular contrastive loss with scaffold-alignment loss as L = Lmol + λsLscaff.Scaffold alignment uses distinct nonempty batch scaffolds as explicit targets for molecule representations.
  • 5 Conclusion: The Lorentz model represents hyperbolic space as the upper sheet of a hyperboloid in one additional ambient dimension.Its geodesic-distance score replaces cosine similarity after lifting the scaled projection from the tangent space at the origin.
  • 5 Conclusion: The bounded Lorentz formulation limits the scaled projection before applying the exponential map, unlike the unbounded formulation.The two implementations differ only in how the scaled projection is mapped onto the hyperboloid.

B Experimental Details and Runtimes

The experiments freeze pretrained encoders for linear probing and report results across three independently pretrained encoders. Runtime measurements compare Euclidean and Lorentz models with and without scaffold alignment.

  • B Experimental Details and Runtimes: Frozen downstream evaluation averages three linear heads within each encoder, then reports mean and sample standard deviation across three independently pretrained encoders.The backbone remains frozen throughout downstream evaluation.
  • B Experimental Details and Runtimes: Pretraining runs use one NVIDIA A40 GPU, with scaffold alignment requiring more time per epoch than MolCLR for both Euclidean and Lorentz models.Representative 50-epoch runtimes are approximately 4.0 versus 6.9 hours for Euclidean models and 3.8 versus 6.8 hours for Lorentz models.
  • B Experimental Details and Runtimes: The downstream benchmark spans biological, physicochemical, and quantum-mechanical molecular properties, so scaffold information need not be equally predictive across tasks.The selected tasks are reported in Table 4.

C Complete Downstream Results

Complete downstream results use task-specific metrics and scaffold splits, with Table 5 reporting results at 15% augmentation and identifying the best and second-best models per column.

  • C Complete Downstream Results: Table 5 reports downstream results at 15% augmentation using the aggregation and metrics of Table 6, with boldface and underlining marking the best and second-best models.Classification uses ROC-AUC, while regression uses RMSE or MAE; all experiments use scaffold splits.

D Structural Evaluation

The evaluation uses a fixed 100,000-molecule cohort with exact-scaffold matches and singleton scaffolds represented separately. Retrieval metrics distinguish exact BM recovery from similarity among nonidentical scaffolds.

  • Cohort construction: The evaluation cohort contains 100,000 randomly selected molecules retained after scaffold, reconstruction, and duplicate checks.All evaluated encoders use the same cohort, without scaffold-frequency balancing or a cap.
  • Cohort construction: The cohort contains 57,939 exact BM scaffolds, including 6,701 multi-member groups and 51,238 singleton scaffolds.Queries are eligible for exact-BM evaluation only when another molecule shares their scaffold.
  • Retrieval metrics: Exact-BM retrieval uses MAP@k, while BM-Tan@k and nDCG@10 assess similarity among structurally related scaffolds.BM-Tan@k averages scaffold ECFP4 similarity among the top k, and nDCG@10 uses that similarity as relevance.

E Projection-space geometry diagnostics

The original unbounded Lorentz implementation collapsed embeddings toward a common radius, making geodesic ranking nearly equivalent to cosine ranking. The bounded formulation was adopted to restore radial variation.

  • Diagnostic: The unbounded Lorentz model drove alignment-space embeddings to the numerical radius limit, producing nearly identical distances from the origin.Negative geodesic distance became almost perfectly rank-equivalent to raw-z cosine similarity, with Spearman correlation ≈1.0.
  • Diagnostic: The bounded Lorentz variant prevents this numerical-limit collapse and restores variation in distance from the origin.The diagnostics concern z, whereas retrieval and downstream evaluation use the frozen Euclidean backbone h.

F Augmentation Diagnostics

Augmentation diagnostics characterize graph fragmentation independently of learned embeddings, while the reported downstream and retrieval tables provide evaluation references at specified augmentation settings. Fragmented views remain in the encoder input despite losing single-molecule validity.

  • Augmentation Diagnostics: For each augmentation operator and strength, the study generates 18,000 stochastic views from a fixed set of 3,000 validation molecules.Connectivity is measured using component counts, largest-component retention, and multi-component split frequency.
  • Evaluation references: The downstream-results table reports performance at 25% augmentation and marks the best and second-best models in each column.The caption does not provide the underlying task values.
  • Evaluation references: The retrieval table reports backbone-space results, separating exact-BM checks from graded metrics that exclude exact-scaffold candidates.Values are means and sample standard deviations across three independently pretrained encoders.
  • Augmentation Diagnostics: Connectivity degrades progressively as augmentation increases from 5% to 30%, making 15% and 25% milder and stronger points within one fragmentation regime.Subgraph removal has a different fragmentation profile, while feature masking preserves connectivity by construction.
  • Augmentation Diagnostics: Fragmented augmented views are passed to the encoder as disconnected graphs even when they no longer represent chemically valid single molecules.These diagnostics quantify structural fragmentation rather than filtering such views from training.

G Backbone-Space Visualization

Backbone-space visualizations compare matched models using the same molecules and t-SNE settings. Scaffold supervision produces more coherent and separated scaffold-specific groups, with scaffold embeddings closer to their molecular groups.

  • Visualization setup: The visualizations use frozen backbone embeddings at 15% augmentation, with cosine distances between L2-normalized embeddings and identical t-SNE settings across models.The comparison includes matched MolCLR and scaffold-supervised models under Euclidean and Lorentz training.
  • Qualitative organization: Scaffold supervision yields more visibly coherent scaffold-specific groups in both Euclidean and Lorentz visualizations.Several groups overlapping under MolCLR become more clearly separated after scaffold supervision.
  • Qualitative organization: In the scaffold-supervised visualizations, scaffold embeddings tend to lie closer to their associated molecular groups.Figure 6 extends the Euclidean view to 6,000 molecules and 14 frequent BM scaffolds.
Loading 2608.29886v1…