This is the transformer architecture explained in plain language: the design that powers GPT, BERT, and every modern AI system. You will implement self-attention, multi-head attention, and positional encoding from scratch in Python. The Transformer is the single most important architecture in deep learning, and once it clicks you will see it everywhere.

“Attention is all you need.”

Ashish Vaswani, Google Brain

Last Updated: July 2026 | Tested on: Python 3.14.6, PyTorch 2.12.1, NumPy 2.4.6 | Difficulty: Expert | Reading Time: 13 minutes

Recurrent Neural Networks (RNNs) read a sentence one word at a time, passing a hidden state from left to right. That is like reading a book through a tiny window that shows one word, where you have to remember everything you saw so far. It is slow, and by the time you reach the end you have half forgotten the start. The Transformer, introduced in 2017, threw that window away.

Instead of recurrence it uses attention, a mechanism that lets every word in a sentence look directly at every other word in one shot. A word 50 positions away is just as easy to reach as the word right next to it. Because all the words are processed together rather than in a chain, training runs in parallel on a GPU (Graphics Processing Unit) and the model learns much longer connections between words.

The core idea is self-attention. Picture a room full of people, and each person is holding up a sign. For each word the network builds three small vectors: a Query (what am I looking for?), a Key (what do I contain?), and a Value (what do I offer?). To decide how much one word should care about another, you match that word’s Query against the other word’s Key with a dot product, then run all the matches through softmax so they turn into percentages that add up to 1. A high score means the two words are relevant to each other.

The new representation of each word is then a weighted blend of everyone’s Values, mixed according to those percentages. Multi-head attention just runs this whole matching game several times at once, so the model can track several kinds of relationships in parallel.

Here is what we cover:

• Self-attention: Query, Key, Value from scratch
• Multi-head attention and why multiple heads matter
• Positional encoding, giving the model a sense of word order
• The full Transformer encoder block
• Why Transformers replaced RNNs for nearly everything

Prerequisites

The Transformer Architecture

The diagram shows the Transformer’s encoder and decoder side by side. The encoder reads the whole input at once through self-attention layers, so every token can look at every other token. The decoder writes the output one token at a time using masked attention, which means it can only peek at tokens it has already produced, never future ones. The trick that makes both halves strong is multi-head attention, which runs several attention patterns in parallel to catch different relationships at once. One head might follow grammar while another follows which word refers to which. This is why the Transformer replaced RNNs for most language tasks: it handles all positions together instead of crawling through them one by one.

Self-Attention from Scratch

Think of it like searching a library: your Query is the topic you are after, each book’s Key is the label on its spine, and its Value is the actual content inside. You compare your request against every spine label, then pull out the contents of the books that match best. The second diagram zooms into the self-attention mechanism for one token: it builds Query, Key, and Value vectors from learned weights, scores its Query against every Key with a dot product, runs those scores through softmax to get attention weights, and blends the Values using those weights.

This Query-Key-Value (QKV) flow is how each token decides which other tokens to pay attention to. Dividing the scores by the square root of the dimension keeps them from blowing up, which would otherwise make softmax too confident and stall learning. Getting this calculation straight is what lets you read attention heatmaps later and figure out why a Transformer behaves the way it does.

📄 self_attention.py: implementing self-attention from scratch

import torch
import torch.nn.functional as F
import numpy as np

def self_attention(Q, K, V):
    """Scaled dot-product self-attention."""
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1) / np.sqrt(d_k)
    weights = F.softmax(scores, dim=-1)
    output = weights @ V
    return output, weights

# Viraj demonstrates with a 4-word sentence
# Each word is represented as a 6-dimensional embedding
torch.manual_seed(42)
seq_len, d_model = 4, 6
words = ["The", "cat", "sat", "mat"]

# Token embeddings (normally learned)
X = torch.randn(1, seq_len, d_model)

# Projection matrices (normally learned)
W_Q = torch.randn(d_model, d_model) * 0.1
W_K = torch.randn(d_model, d_model) * 0.1
W_V = torch.randn(d_model, d_model) * 0.1

Q = X @ W_Q  # Queries: "What am I looking for?"
K = X @ W_K  # Keys: "What do I contain?"
V = X @ W_V  # Values: "What do I offer?"

output, attention_weights = self_attention(Q, K, V)

print(f"Input shape:     {X.shape}")
print(f"Q, K, V shape:   {Q.shape}")
print(f"Output shape:    {output.shape}")

print(f"\nAttention weights (each row sums to 1):")
w = attention_weights.squeeze().detach().numpy()
for i, word in enumerate(words):
    row = " ".join(f"{w[i][j]:.3f}" for j in range(len(words)))
    print(f"  {word:>4} attends to: [{row}]")

▶ Output

Input shape:     torch.Size([1, 4, 6])
Q, K, V shape:   torch.Size([1, 4, 6])
Output shape:    torch.Size([1, 4, 6])

Attention weights (each row sums to 1):
   The attends to: [0.236 0.261 0.251 0.253]
   cat attends to: [0.238 0.244 0.267 0.251]
   sat attends to: [0.227 0.264 0.250 0.259]
   mat attends to: [0.246 0.243 0.267 0.244]

What happened here: Self-attention lets every word look at every other word. The attention weights all hover around 0.25 because this network has not been trained yet, so it has no idea yet which words belong together. Once trained, “cat” would lean hard toward “sat” (subject and verb) and “mat” (where it sat), while “The” would mostly point at “cat” (the noun it describes). Dividing by sqrt(d_k) keeps the raw dot products from getting huge, which would otherwise push softmax into a flat, saturated zone where it stops learning. Your exact numbers may shift slightly on a different PyTorch build, since they come from random untrained weights.

Multi-Head Attention

One attention head can only track one kind of relationship at a time. That is like asking a single person to watch grammar, meaning, and word order all at once. They will miss things. Multi-head attention hires a small team instead. It runs N attention operations in parallel (typically 8 or 12), each with its own learned projections. One head might learn grammar (subject and verb), another learns meaning (dog and pet), and another learns position (which words sit next to each other). At the end their outputs are concatenated and projected back to the original size, so the rest of the network sees one combined view.

📄 multi_head_attention.py: multi-head attention in PyTorch

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.num_heads = num_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch, seq_len, d_model = x.shape

        Q = self.W_q(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2)

        scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)
        attn = torch.softmax(scores, dim=-1)
        context = attn @ V

        context = context.transpose(1, 2).contiguous().view(batch, seq_len, d_model)
        return self.W_o(context)

# Aditi tests multi-head attention
mha = MultiHeadAttention(d_model=64, num_heads=8)
x = torch.randn(2, 10, 64)  # 2 sentences, 10 tokens, 64-dim embeddings
output = mha(x)

print(f"Input:  {x.shape}")
print(f"Output: {output.shape}")
print(f"Heads:  {mha.num_heads}, d_k per head: {mha.d_k}")
print(f"Params: {sum(p.numel() for p in mha.parameters()):,}")

▶ Output

Input:  torch.Size([2, 10, 64])
Output: torch.Size([2, 10, 64])
Heads:  8, d_k per head: 8
Params: 16,640

What happened here: The 64-dimensional embedding is split into 8 heads of 8 dimensions each. Think of it like 8 readers each given a thin slice of the meaning, where one watches grammar, one watches topic, one watches position. Each head runs its own attention on its slice and learns a different kind of relationship. The 8 results are stitched back into 64 dimensions and passed through a final projection layer W_o. This is the core building block of every Transformer. Modern models such as GPT-5.5, Claude Opus 4.8, BERT, and Llama all use exactly this mechanism, just with far larger dimensions (current at the time of writing; models change fast, so check the provider docs).

Positional Encoding: Teaching Order

📄 positional_encoding.py: adding position information to embeddings

import torch
import torch.nn.functional as F
import numpy as np

def positional_encoding(max_len, d_model):
    """Sinusoidal positional encoding from 'Attention Is All You Need'."""
    pe = np.zeros((max_len, d_model))
    position = np.arange(max_len)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))

    pe[:, 0::2] = np.sin(position * div_term)  # Even dimensions
    pe[:, 1::2] = np.cos(position * div_term)  # Odd dimensions
    return torch.tensor(pe, dtype=torch.float32)

# Anvay creates positional encodings
pe = positional_encoding(max_len=20, d_model=8)
print("Positional encoding shape:", pe.shape)
print("\nFirst 5 positions (8 dimensions each):")
for i in range(5):
    vals = " ".join(f"{pe[i][j]:>6.3f}" for j in range(8))
    print(f"  Position {i}: [{vals}]")

print("\nKey property: nearby positions have similar encodings,")
print("distant positions have different encodings.")
print(f"  Similarity(pos 0, pos 1): {F.cosine_similarity(pe[0:1], pe[1:2]).item():.4f}")
print(f"  Similarity(pos 0, pos 10): {F.cosine_similarity(pe[0:1], pe[10:11]).item():.4f}")

▶ Output

Positional encoding shape: torch.Size([20, 8])

First 5 positions (8 dimensions each):
  Position 0: [ 0.000  1.000  0.000  1.000  0.000  1.000  0.000  1.000]
  Position 1: [ 0.841  0.540  0.100  0.995  0.010  1.000  0.001  1.000]
  Position 2: [ 0.909 -0.416  0.199  0.980  0.020  1.000  0.002  1.000]
  Position 3: [ 0.141 -0.990  0.296  0.955  0.030  1.000  0.003  1.000]
  Position 4: [-0.757 -0.654  0.389  0.921  0.040  0.999  0.004  1.000]

Key property: nearby positions have similar encodings,
distant positions have different encodings.
  Similarity(pos 0, pos 1): 0.8838
  Similarity(pos 0, pos 10): 0.4240

What happened here: On its own, self-attention is blind to order. It treats the input as a bag of words, not a sequence, so “dog bites man” and “man bites dog” look identical to it. Positional encoding fixes that by adding sinusoidal waves of different frequencies to each embedding dimension, a bit like giving every seat in a stadium a unique row-and-seat pattern. Adjacent positions 0 and 1 come out highly similar (0.88), while positions 0 and 10 are clearly further apart (0.42), and the gap keeps growing with distance.

Because the waves are smooth and repeating, the model can learn relative offsets: the step from position 3 to 5 carries the same shape of change as the step from 100 to 102.

Common Mistakes

Interview Corner

Q: Why did Transformers replace RNNs?

Three reasons. First, parallelism: self-attention processes every position at the same time, while an RNN has to walk through the sequence step by step. On a GPU that is roughly a 10x to 100x training speedup. Second, long-range dependencies: attention connects any two positions directly, an O(1) path, whereas in an RNN information has to survive O(n) hops before two distant words can interact. Third, scalability: Transformers keep improving as you add billions of parameters and trillions of training tokens, while RNNs hit diminishing returns much sooner.

Practice Exercises

1. Implement a complete Transformer encoder block (self-attention + feed-forward + layer norm + residual).
2. Implement masked (causal) self-attention for autoregressive generation.
3. Visualize attention weights for a real sentence using a pre-trained model from HuggingFace.
4. Compare the parameter count of a 6-layer Transformer with a 6-layer LSTM (Long Short-Term Memory) of equal hidden size.

More in this series:

• Deep Learning Capstone: Image Classifier on Hugging Face
• NLP: Text Classification with HuggingFace Transformers
• How LLMs Are Trained: Pretraining, SFT, and RLHF Explained

Frequently Asked Questions

Interview Questions on the Transformer Architecture

Interviewers rarely ask for definitions. They ask what happens in situations like these.

Q: Your Transformer produces the same output for “dog bites man” and “man bites dog”. What did you most likely forget?

Positional encoding. Self-attention on its own is order-blind: it treats the input as a bag of tokens, so two sentences with the same words in a different order look identical to it. The fix is to add positional encodings (sinusoidal, as in the original paper, or learned) to the token embeddings before the first attention layer, which injects a sense of sequence.

Q: You train on 8,000-token sequences and the GPU runs out of memory almost immediately. What do you check first?

Attention memory scales as O(n^2) with sequence length, so going from 512 to 8,000 tokens blows up the score matrix roughly 250 times. First confirm the sequence length and batch size are what you expect. Then reduce the effective footprint: shrink the batch, use gradient checkpointing, switch to a memory-efficient kernel like FlashAttention, or chunk the input. The quadratic term, not the model depth, is usually the culprit.

Q: Why are the attention scores divided by the square root of d_k?

As the head dimension d_k grows, the dot product of Query and Key vectors grows in magnitude too, because you are summing over more terms. Large scores push softmax into a saturated region where one weight is near 1 and the rest near 0, and the gradients there are tiny. Dividing by sqrt(d_k) keeps the scores in a sane range so softmax stays smooth and the model keeps learning.

Q: What is the difference between encoder self-attention and decoder self-attention?

Encoder self-attention is bidirectional: every token can look at every other token, past and future, because the whole input is known up front. Decoder self-attention is masked (causal): a token can only attend to itself and earlier tokens, never later ones. That mask is what lets the decoder generate text one token at a time without cheating by peeking at the answer it has not written yet.

Q: What do residual connections and layer normalization do in a Transformer block?

Each sub-layer adds its output back to its input (the residual, or skip, connection), which gives gradients a short path straight through a deep stack of 6 to 96 layers and prevents them from vanishing. Layer normalization then rescales each token’s vector to a stable mean and variance, which keeps the numbers well-behaved and makes training far less sensitive to the learning rate. Skip these and deep Transformers become very hard to train.

Q: What is cross-attention and where does it appear?

Cross-attention (also called encoder-decoder attention) sits inside each decoder block. Unlike self-attention, its Queries come from the decoder while its Keys and Values come from the encoder’s output. This is how a translation or summarization model lets the text it is generating look back at the source sentence. Encoder-only models like BERT have no cross-attention, and decoder-only models like GPT drop it too and rely on masked self-attention alone.

What’s Next?

You now understand the Transformer architecture from the inside out: self-attention with Query, Key, and Value vectors, why the scores are scaled by sqrt(d_k), how multi-head attention tracks several relationships at once, and how positional encoding gives an order-blind model a sense of sequence. That single stack of ideas is what powers almost every modern AI system you use. Next, in the BERT and GPT tutorial, we see how BERT uses the encoder for understanding tasks and GPT uses the decoder for generation tasks: two sides of the same Transformer coin.

Want to see where this fits in the bigger picture? Head back to the Python + AI/ML tutorial series home for the full roadmap from Python basics to production AI.

Reference: the complete, always-current details live in Hugging Face documentation.

Related Posts

Previous: The Attention Mechanism From Scratch in PyTorch

Next: Natural Language Processing (NLP): BERT, GPT, and Modern Language Models

Series Home: Python + AI/ML Tutorial Series