Architecture¶
The task and the supervision available¶
OCSR maps a raster depiction of a molecule to a machine-readable structure. Two output formats dominate:
- Graph decoding (MolGrapher, MolScribe's graph head): predict atoms with image coordinates, then bonds between them. Stronger at the top end, and its errors are localizable. It needs atom-level coordinate supervision.
- Sequence decoding (DECIMER, Img2Mol, MolScribe's character head): emit SMILES autoregressively. Weaker supervision signal, but it only needs (image, SMILES) pairs.
MolMini decodes sequences. That is not a claim that sequences are better: it is a consequence of the supervision this project actually has. Training images come from RDKit rendering, which yields an exact SMILES per image and nothing else. Building coordinate supervision would mean extracting atom positions from the drawer, which is possible but is a different project.
Encoder: convolutions then attention¶
384x384x1
-> conv 7x7 stride 4 -> 96x96x64
-> conv 3x3 stride 2 + 3 ConvNeXt blocks -> 48x48x128
-> conv 3x3 stride 2 + 3 ConvNeXt blocks -> 24x24x256
-> linear -> 576 tokens x 384
-> 4 transformer blocks with axial 2D RoPE
Attention at pixel resolution would dominate the FLOP budget and buy nothing:
the features that matter early (strokes, junctions, character glyphs) are local,
which is exactly what convolutions are for. Attention becomes necessary only
once tokens have to be related across the image -- a ring closure digit 1
appearing twice, or an R group referenced elsewhere in the drawing -- and by
then there are only 576 of them.
Resolution is not a free parameter. Atom labels, charges and subscripts are
small glyphs; at 224 px a [N+] inside a dense fused-ring system stops being
legible at all. 384 px is the smallest input at which the tokens near a crowded
heteroatom still carry a readable glyph.
Decoder¶
Six pre-norm blocks: causal self-attention with 1D RoPE, cross-attention to the
576 encoder tokens, SwiGLU feed-forward. Embeddings are tied to the output head.
The vocabulary is atom-level (Br, [C@@H], %10 are single tokens), which
keeps sequences about 40% shorter than character-level and removes a whole class
of malformed-bracket failures the model would otherwise have to learn to avoid.
Component choices¶
| Choice | Instead of | Why |
|---|---|---|
| RMSNorm | LayerNorm | one statistic instead of two; no measured quality cost |
| SwiGLU | GELU MLP | better loss per parameter at equal FLOPs |
| RoPE (2D encoder, 1D decoder) | learned absolute positions | relative by construction; the encoder gets true 2D positions instead of a flattened raster order |
| QK-norm | none | attention logits stay bounded at small width, which is what allows the higher LR |
| Tied embeddings | separate head | at 88 vocabulary entries this is small, but free |
| Depth-scaled residual init | default init | keeps activation variance flat with depth, so warmup can be short |
| WSD schedule | cosine | flat peak means any mid-run checkpoint is usable and the step budget can change without invalidating the schedule |
| Weight EMA | last iterate | consistently better under a flat LR, at the cost of one parameter copy |
Why on-the-fly data¶
Rendering one depiction costs ~2 ms of CPU; one training step costs far more GPU than that per image. A frozen image set would therefore buy nothing and cost generalization: the model would see each image many times and memorize rendering artifacts. Generating fresh means the model essentially never sees the same pixels twice, and the augmentation distribution -- not a fixed dataset -- is what it learns to be invariant to.
The corresponding hazard is that the augmentation distribution is the specification. Anything absent from it (a depiction convention, a degradation mode) is out of distribution at test time no matter how much compute is spent.
Known limits¶
- Hand-drawn depictions. RDKit
comicModegives wobbly strokes, but it does not reproduce human line breaks, inconsistent bond lengths or misplaced labels. Real hand-drawn input is out of distribution and is reported separately for that reason. - Markush structures and R groups. Not in the corpus and not rendered.
- Reaction schemes, multi-molecule pages. Out of scope; the model expects one molecule per image.
- Very large molecules. The declared scope caps heavy atoms; peptides and polymers fall outside it.