Building TrOCR from scratch: what wiring a vision encoder to a text decoder actually taught me

I spent a weekend answering one question end to end: how does a model look at a picture of text and type out the characters? I already knew the NLP Transformer — Q/K/V, multi-head attention, positional encodings — so I didn’t want a tutorial that treats the Vision Transformer as a magic box. I wanted to build the whole thing from scratch in PyTorch: the ViT encoder, the text decoder, and the cross-attention that joins them, aiming eventually at handwritten Devanagari OCR.

This is what I did, what confused me, and the handful of ideas that turned it from “a stack of papers” into code I can run.

What you’ll be able to do by the end: explain a ViT as a plain Transformer encoder, implement multi-head attention that serves self / masked / cross duty, and understand exactly how the image’s encoding gets “translated” into text.

The mental model that unlocked everything

The single sentence that made ViT stop being intimidating: a Vision Transformer is the Transformer encoder I already know, run on a sequence of image patches instead of word tokens. That’s ~80% of it. There are only two genuinely new pieces, and they sit at the two ends of the stack — patch embedding at the input, and a readout choice at the output. Everything in between (LayerNorm, attention, MLP, residuals) is identical to the NLP encoder.

The second sentence, for the OCR half: OCR is translation, and the image is the source language. An encoder reads the image once into memory; an autoregressive decoder writes the text one character at a time, looking back at that memory at every step. Swap “read French words” for “read image patches” and a translation model becomes an OCR model. The thing that lets the text decoder look back at the image is cross-attention — and that turned out to be the crux of the whole weekend.

Walking through it

Patchifying an image into tokens

A 224×224 image with patch size 16 becomes a 14×14 grid = 196 patches, each a token. The clean trick is that “cut into patches, flatten, and linearly project” is literally one convolution with kernel_size = stride = patch:

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch=16, in_ch=3, dim=768):
        super().__init__()
        self.num_patches = (img_size // patch) ** 2
        self.proj = nn.Conv2d(in_ch, dim, kernel_size=patch, stride=patch)
    def forward(self, x):            # [B,3,224,224]
        x = self.proj(x)             # [B,768,14,14]
        return x.flatten(2).transpose(1, 2)   # [B,196,768]

Multi-head self-attention over patches

I implemented this first, since it’s where a ViT actually “thinks.” The whole module is: one fused QKV projection, split into heads so each head attends in its own 64-dim subspace, scaled dot-product, merge, project out.

q, k, v = torch.split(self.qkv(x), self.dim, -1)
q = q.reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)   # [B,H,N,head_dim]
# ... same for k, v
scores = q @ k.transpose(-2, -1) / (self.head_dim ** 0.5)            # [B,H,N,N]
attn = torch.softmax(scores, dim=-1)
out = (attn @ v).transpose(1, 2).reshape(B, N, self.dim)

The one new mechanism: cross-attention

A decoder block has three sub-layers instead of two: masked self-attention (text looks at text-so-far), cross-attention (text looks at the image), and an MLP. Cross-attention is the same attention math with one twist — Query comes from the decoder (text); Key and Value come from the encoder (image memory). So I generalized my attention module to take two inputs:

def forward(self, x_q, x_kv=None, mask=None):
    if x_kv is None:
        x_kv = x_q                       # self-attention: same source
    q = self.to_q(x_q)                   # queries from text
    k, v = torch.split(self.to_kv(x_kv), self.dim, -1)   # keys/values from image
    # ... reshape to heads, scaled dot-product, optional causal mask ...
    return self.out(out), attn

The attention matrix here is T × N — text positions by image patches, not square. Each text token’s row is a softmax distribution over the image regions: “to write this character, which parts of the image do I look at?” Visualize it on a line of text and a diagonal band appears — the model learns to sweep left-to-right across the image as it emits characters. That’s reading, with nobody ever drawing character boxes. That picture is the whole reason cross-attention suits OCR.

Training: teacher forcing, then prove it overfits

Training is supervised seq2seq. You feed the decoder the target shifted right and ask it to predict the next token, in one parallel pass (the causal mask stops it peeking ahead):

dec_in = [[BOS] + ids for ids in batch]     # feed:    [BOS] a n i l
labels = [ids + [EOS] for ids in batch]      # predict:  a n i l [EOS]
loss = F.cross_entropy(logits.reshape(-1, V), labels.reshape(-1), ignore_index=PAD)

The habit I’m keeping: before touching real data or a GPU, overfit a tiny synthetic set (16 examples). If loss falls to ~0 and generate() reproduces the strings, the whole pipeline — patch embedding, cross-attention, teacher forcing, autoregressive decoding — is wired correctly. If it can’t memorize 16 examples, there’s a bug, and you find it in two minutes instead of two hours.

Gotchas I hit

“Flatten” is just serializing in a fixed order — and the fixedness is the point. I didn’t get how a 16×16×3 patch becomes a 768-vector. It’s not a computation; it walks the 3D block in some order (channel-first vs channel-last are both fine) and writes the numbers into a line. What matters is that every patch uses the same order, so the projection weights can learn a consistent meaning per position.

Patch size ≠ grid. I conflated the 16 and the 14. Sixteen is the patch edge in pixels; 14 is how many patches fit across, in patches; 196 is the sequence length. image size ÷ patch size = grid. Sanity check that nailed it: 196×768 = 224×224×3 = 150,528 — patchifying just regroups the same pixels, it adds and drops nothing.

In cross-attention, output length follows the query, not the key. My generalized attention had a subtle bug: I reshaped the output back using the key length. It worked in self-attention (query length == key length) but would have silently crashed in cross-attention, where I have 20 text queries against 196 image keys — the output must be 20 vectors (one per character), not 196. Attention always produces one output per query. Fixing that one Tn → Tq was the moment cross-attention actually clicked.

You can’t fuse the QKV projection once Q and K/V come from different places. In self-attention I used a single Linear(dim, 3*dim) on one input. Cross-attention feeds Q from text and K/V from the image — two different tensors of different lengths — so the projection has to split into a to_q and a to_kv.

Two quieter bugs in my ViT. My MLP hardcoded Linear(4*dim, dim) instead of Linear(mlp_ratio*dim, dim) — fine at the default ratio, silently wrong at any other. And I applied the final LayerNorm before the blocks; in a pre-norm Transformer it belongs after the stack, right before the CLS readout. Both are the kind of thing that trains without erroring and quietly costs you accuracy.

When this approach fits (and when it doesn’t)

Building it from scratch is the right call when the goal is understanding — every knob in HuggingFace’s VisionEncoderDecoderModel now maps to something I wrote by hand. For an actual product, you don’t train a ViT from scratch: it’s data-hungry (the original paper needed 300M images to beat CNNs), so you fine-tune a pretrained encoder and decoder. And the field is now forking — TrOCR/Donut/Nougat use encoder–decoder cross-attention, while newer models (DTrOCR, GOT-OCR2.0, general VLMs) drop the separate encoder and feed image patches straight into one decoder. For focused, single-script OCR, the cross-attention design is still the pragmatic choice.

Visit Github for the

Takeaways

  • A ViT is a Transformer encoder on patch tokens; only patch embedding and readout are new.
  • Cross-attention = Q from the decoder (text), K/V from the encoder (image); the T × N matrix is a learned alignment between characters and pixels — no segmentation needed.
  • Output length always follows the query length. This one rule explains the sneakiest cross-attention bug.
  • Self, masked, and cross attention are one mechanism with different inputs — so one well-written module serves all three.
  • Overfit 16 examples before anything else; it’s the cheapest correctness test in deep learning.
  • Weak inductive bias + lots of data beats strong priors — which is why you fine-tune pretrained weights instead of training from scratch.

Get the code

All of this — the from-scratch attention module, the ViT encoder, the text decoder, and the overfit training loop — is on GitHub:

👉 GitHub

The repo is organized so one attention module (attention.py) is the shared foundation, with encoder.py, decoder.py, and model.py building on top. Run the sanity check from inside the package with python train.py — it overfits a tiny toy set, so if the loss drops to ~0 you know the whole pipeline is wired correctly. Clone it, break it, and step through the shapes yourself; that’s where the understanding actually lands.

Going further

  • TrOCR (Li et al., 2021) — the canonical vision-encoder / text-decoder OCR model: arxiv.org/abs/2109.10282
  • An Image is Worth 16×16 Words (ViT, Dosovitskiy et al., 2020) — the patch-as-token idea and the data-hunger result.
  • Donut & Nougat — OCR-free document and academic-PDF understanding, same encoder–decoder recipe.
  • GOT-OCR2.0 / DTrOCR — the decoder-only, no-cross-attention direction, for contrast.
  • Next for me: swap the toy data for handwritten Devanagari (IIIT-HW-Dev), and decide the tokenizer at grapheme-cluster vs Unicode-code-point level — the matras and conjuncts make that the real design decision.