Molecular Depiction Alignment: Contrastive vs Predictive

Overview

Two ways to align an image of a molecular structure diagram into the embedding space of a frozen chemistry language model, compared on one data and evaluation stack. A DINOv2 ViT-S/14-reg vision tower, a frozen MIST-28M text tower, and a 197,120-parameter linear map between them.

Arm A optimizes the SigLIP sigmoid pairwise loss. Arm B regresses the frozen target directly in latent space with no negatives. Everything else is held identical, and an assertion checks that before any comparison table prints.

The comparison is the deliverable. The models are how it was obtained.

Results

Retrieval R@1 against 64-molecule galleries of nearest-Tanimoto distractors, chance 0.0156. Out-of-distribution is WildMol-10k, real depictions extracted from patents and papers, never trained on.

With the backbone frozen, the ceiling is the representation. Fifty times more data moved a closed-form ridge ceiling by +0.020 against +0.088 predicted by a log-linear law fitted at small scale, and 54.3% of DINOv2’s feature variance on these images is depiction style rather than molecular identity.

Unfreezing moves retrieval 9.6x in-distribution and 11.1x out of distribution, and the pre-registered prediction reverses: the predictive arm transfers better under a frozen backbone and the contrastive arm wins outright once the backbone can move.

Features

• Pre-registered predictions, committed before any training run and quoted verbatim in the results with how each came out. One was wrong outright and one held then reversed.
• Four harness gates run before any headline number is believed: an oracle retrieval control, MIST against Tanimoto similarity, a random-embedding control, and a trivial image-statistics floor.
• Seven linear-probe points across five MoleculeNet datasets, on Bemis-Murcko scaffold splits committed as build artifacts so a probe number is comparable across clones.
• An out-of-distribution slice of real literature depictions, deduplicated against every training set by InChIKey with a positive control confirming the check detects injected leakage.
• Open weights and data, with an export path that verifies the published weights are the same function of pixels as the checkpoint they came from.

Usage

The released weights need timm and torch, and nothing from this repository:

import timm, torch

backbone = timm.create_model(
    "vit_small_patch14_reg4_dinov2.lvd142m", pretrained=False, num_classes=0, img_size=224
)
backbone.load_state_dict(torch.load("backbone.pt", weights_only=True))
head = torch.nn.Linear(384, 512)
head.load_state_dict(torch.load("projection_head.pt", weights_only=True))
backbone.eval(); head.eval()

with torch.no_grad():
    embedding = head(backbone(pixels))   # (batch, 512), in MIST space

Comparing images to each other works with the above alone. Comparing an image to a molecule additionally needs MIST to embed that molecule.

Weights are on Hugging Face as contrastive and predictive, and the 20K depiction set is published so the comparison can be re-run without the rendering environment.

Retrospective

Scoped to a couple of weekends, and complete rather than paused. There is no roadmap.

The parts worth keeping are methodological. A metric defect scored a working model at the pre-registered clean-negative floor with every exit criterion passing, and it was caught by asking what a perfect model would score rather than by a test. The arm-neutral checkpoint selection rule, chosen because it is the one scalar both arms are measured by, selected epoch 1 of 24 for the contrastive arm at a cost of 2.6x on in-distribution retrieval. That was reported rather than fixed, since changing a pre-registered rule after seeing results is what pre-registration exists to prevent.

Every v2 number is one seed. The only repeats in the project are a five-seed frozen sweep, and the retrieval orderings clear that noise floor by roughly two orders of magnitude while individual probe cells do not.

The code is Apache-2.0. The released weights carry research-only terms mirroring MIST’s, since they were trained to predict its embeddings.