DINOv2 to MIST, contrastive alignment

A DINOv2 ViT-S/14-reg fine-tuned so that an image of a molecular structure diagram embeds where its molecule embeds in the frozen MIST-28M embedding space. Objective: SigLIP sigmoid pairwise loss.

This model is one half of a controlled comparison. Its sibling, dinov2-mist-molecular-depiction-predictive, is identical in every respect except the alignment objective: same molecules, same splits, same seed, same schedule, same evaluation. The comparison is the deliverable; the models are how it was obtained. Full write-up and code: https://github.com/hunter-heidenreich/molecular-depiction-alignment.

Results

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

This is the stronger of the two on real depictions, and it is the arm the pre-registered prediction expected to lose.

One seed. Orderings of this size are safe; small differences are not. The only error bars in the project are a 5-seed frozen sweep, reported in the repository.

Usage

Needs timm and torch. Nothing else, and no MIST, unless you are comparing against molecules.

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

Images are 224px RGB renders of structure diagrams, normalised as DINOv2 expects.

Comparing images to each other works with the above alone. Comparing an image to a molecule additionally needs MIST to embed that molecule, pooled last_hidden_state[:, 0, :]; mean-pooling changes what the space means.

What this is not

Not an OCSR model. Reading a molecule out of a picture is a mature field (MolScribe, MolGrapher, DECIMER, Img2Mol) and this does not compete with it. This is a controlled experiment about alignment objectives that happens to produce usable weights.

Also worth knowing: the training depictions are rendered from SMILES, so a picture carries no information the string did not already have. The in-distribution number is context; the OOD number is the result.

Checkpoint selection

Exported from epoch 24 of 24, the final epoch, which is what every number above was measured on.

The project's pre-registered selection rule picks the checkpoint with the best validation centered cosine, and for the contrastive arm that rule selected epoch 1, whose in-distribution R@1 is 0.2265 against the final epoch's 0.5938. The rule was not changed after seeing that; changing a pre-registered rule once results exist is the retrofitting the pre-registration exists to prevent. Both epochs are reported in the repository, and scripts/53_rescore.py is what scores an arbitrary saved epoch.

Licence

Research use only. See LICENSE.weights in this repository:

• Research use only
• No redistribution without permission
• No commercial use without a licensing agreement

These terms mirror MIST-28M's. The weights were trained to predict MIST's embeddings, MIST is published under exactly those conditions, and its model card does not say whether they are intended to reach a model trained on its outputs. Rather than assert a reading that might grant more than is ours to grant, this release grants no more than MIST does. It is very likely stricter than necessary, since these weights contain no MIST parameters and run without it; that is deliberate.

The weights are also a derivative of DINOv2 (Apache-2.0), whose attribution and licence text travel with them in NOTICE and LICENSE. Nothing here restricts DINOv2 itself.

Training molecules are PubChem only. WildMol-10k / MolParser data is evaluation-only and was never trained on, and nothing derived from it is in these weights.