DEV Community

AheadMint
AheadMint

Posted on Originally published at aheadmint.com

Transformer Architecture in 5 Questions Developers Actually Ask (with 12 Lines of PyTorch)

Most Transformer explainers start with the famous diagram from "Attention Is All You Need" and lose half their readers at the first box. As a developer, I found it easier to learn the architecture by answering the five questions I kept asking myself, and by writing the core operation in code.

1. What problem was the Transformer actually solving?

RNNs and LSTMs read a sentence one token at a time, like a message passed down a line of people. That makes long-range relationships fragile and training hard to parallelise. The Transformer replaced the chain with attention: every token can look directly at every other token in one step. Think of it as moving from a relay race to a meeting room where everyone can hear everyone.

2. What are Q, K and V, really?

Each token is projected into three vectors by learned weight matrices:

  • Query: what this token is looking for
  • Key: what this token offers
  • Value: the information it passes on if it turns out to be relevant

Compare queries with keys, turn the scores into weights, and use the weights to mix values. That is the whole idea.

3. What does that look like in code?

Here is scaled dot-product attention, the heart of every Transformer block:

import torch
import torch.nn.functional as F

def attention(q, k, v, mask=None):
    # q, k, v: (batch, seq_len, d_k)
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / d_k ** 0.5   # compare + scale
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = F.softmax(scores, dim=-1)              # normalise
    return weights @ v                               # collect values
Enter fullscreen mode Exit fullscreen mode

The / d_k ** 0.5 line is easy to skip, but it stops large dot products from pushing softmax into near one-hot outputs, which would kill gradients.

4. Why multiple heads?

One attention pattern can't capture everything in "Aman told Ravi that he would lead the project": who "he" is, which verb goes with which noun, what "lead" applies to. Multi-head attention runs several attention operations in parallel on different learned projections and concatenates the results. Each head is like a different camera angle on the same match.

5. How does a model like GPT generate text with this?

Two more pieces make it work:

  • Positional encoding: attention by itself ignores order, so "dog bites man" and "man bites dog" would look identical without it.
  • Causal masking: in a decoder, the mask argument above blocks each token from attending to future tokens. The model predicts one token, appends it, and repeats.

Add a feed-forward layer, residual connections and layer normalisation around each sublayer, stack the block N times, and you have the architecture behind modern LLMs.

A quick map of the model families

Family Example Best at
Encoder-only BERT Understanding: classification, search
Decoder-only GPT Generation
Encoder-decoder T5, original Transformer Translation, summarisation

Go deeper

I wrote a longer, beginner-friendly walkthrough that covers every block (encoder, decoder, cross-attention, residuals, masking) with everyday analogies: Transformer Architecture Explained on AheadMint.

If you've got a clearer mental model for Q/K/V, share it in the comments. I'm always collecting better analogies.

Top comments (0)