AboutWhat I doProjectsContact
Blog·Modelado·04/07/2026·2 min

What I learned implementing a transformer from scratch

Multi-head attention, causal masking and a BPE tokenizer handwritten in PyTorch, without torch.nn.Transformer. Three lessons no course had given me, including a surprise in the ablations.

What I learned implementing a transformer from scratch

I had spent months using transformers through libraries and could explain attention from memory: queries, keys, values, softmax, done. Then I tried to write one from scratch in PyTorch — no torch.nn.Transformer, no external tokenizer — and found out how much of that explanation was memory and how much was understanding. Spoiler: I got the head reshape wrong three times.

Lesson 1: dimensions are the real exam. Multi-head attention is conceptually simple and mechanically treacherous: project, split into heads, transpose, attend, merge back. Every step is a reshape where you can fail silently — the code runs, the shapes happen to line up, and the model simply learns worse. Writing the shape tests (and the content ones: that the causal mask truly prevents looking at the future) taught me more attention than re-reading the paper.

Lesson 2: the tokenizer is half the model. Writing byte-pair encoding by hand — count pairs, merge the most frequent, repeat — demystifies why LLMs split words where they do, why numbers tokenize strangely and why swapping tokenizers breaks a trained model. Now, when a token count does not add up, I know exactly what is happening underneath.

Lesson 3: breaking things on purpose teaches more than training them. The ablation study gave the best moment of the project. Removing the positional embeddings pushed validation perplexity from 6.89 to 8.04 — expected: without positions, attention cannot tell order apart. The surprise was the other one: with a single attention head the model scored 6.76, slightly better than with several. At this model and data size, the extra heads were wasted capacity. Reading that in a paper is one thing; seeing it in your own results table is another.

What I keep. It is not the model — a toy GPT over public-domain text is useful to nobody. It is the mental model: when something does not converge, when sampling degenerates, when a fine-tune misbehaves, I now have a concrete picture of which piece might be failing and how to check it. That intuition was not in any course I took; it came from implementing every tensor and watching it fail.

Shall we connect?

Let's talk: feedback, collaboration or an opportunity.

You can write to me about a project, a technical question, to give me feedback, or about an internship or a first junior role. I always reply.

Write to me

Keep reading:

hola@jmwebsoluciones.com