DINOv2 to MIST, predictive 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: smooth-L1 regression onto the frozen target, no negatives (the JEPA move).
This model is one half of a controlled comparison. Its sibling,
dinov2-mist-molecular-depiction-contrastive,
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.
| R@1 | R@5 | |
|---|---|---|
| in-distribution (synthetic renders) | 0.4963 | 0.8045 |
| out-of-distribution (real depictions) | 0.2847 | 0.5779 |
| degradation to real | -43% |
This arm holds up better out of distribution under a frozen backbone. Once the backbone is trainable the ordering reverses, which is the finding.
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.