Skip to content

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 comicMode gives 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.