Structural Hierarchy and Latent Geometry in Self-Supervised Molecular Representation Learning

Abstract

Chemical structure admits useful hierarchical abstractions: atoms assemble into rings, rings into scaffolds, and scaffolds give rise to families of related compounds that share a common core. Molecular self-supervised contrastive learning typically ignores this structure, forming positive pairs through random augmentations, while embedding molecules in Euclidean space. We study whether explicitly encoding a molecule’s Bemis–Murcko scaffold, and using it to supervise the molecular embedding, changes what the model learns, and whether this effect depends on the embedding geometry, comparing Euclidean and Lorentz contrastive objectives. Pre-trained on 750,000 molecules drawn from a one-million-molecule subset of PubChem and evaluated on MoleculeNet, the resulting model improves downstream property-prediction performance over a matched non-scaffold baseline on most benchmarks, while also learning a more scaffold-structured representation space in which molecules embed closer to their underlying scaffolds and to other molecules built on related cores. Training with the Lorentz objective sharpens this scaffold organization only in combination with the scaffold signal, and the effect does not translate into a consistent downstream advantage over Euclidean training.