Build a Transformer From Scratch in PyTorch: Attention, Encoder and Decoder

Transformers showed that sequence modeling can be built entirely on attention, with no recurrence and no convolutions. The authors of the original Transformer paper proposed an encoder-decoder design based on self-attention and showed that it trains in parallel very well and performs strongly on machine translation.

When transformers only ever show up through high-level libraries, they're easy to treat as magic. Writing one yourself is among the best ways to see what actually happens inside. What follows is a compact PyTorch version of the core parts, starting with how positions are encoded, moving through both flavors of attention and the feed-forward block, then through encoder layers, and ending with a sketch of the decoder. Useful companions are the paper itself, Harvard NLP's Annotated Transformer and the torch.nn reference.

Why write one by hand

The paper is clear, but turning it into running code involves many small decisions: tensor shapes, masking, normalization, residual connections and how modules fit together. The Annotated Transformer became famous because it converts the paper into a line-by-line educational implementation, which makes the design much easier to reason about.

Building it yourself also makes three defining ideas concrete:

1. Attention takes over the job recurrence used to do, which is relating tokens to each other.
2. Position has to be injected explicitly, since attention by itself has no sense of word order.
3. Transformer blocks are modular, which suits torch.nn.Module, ModuleList, linear layers, normalization layers and dropout very well.

A simple mental picture

The classic model has two stacks. The encoder stack applies self-attention and a feed-forward network in every layer. The decoder stack does that as well, plus masked self-attention and attention over the encoder's output.

Each sublayer is wrapped in a residual connection and layer normalization. The normalization works across the last feature dimensions and can use learnable scale and shift parameters. Seen this way, a block is less mysterious than it sounds. It's mostly linear projections, attention, normalization and residual paths.

Step 1: Embeddings and positional encoding

Transformers work on vectors, so token IDs first become dense embeddings. A positional encoding is then added so the model can tell "cat chased dog" apart from "dog chased cat". The paper used sinusoidal encodings for this.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # shape: (1, max_len, d_model)

        self.register_buffer("pe", pe)

    def forward(self, x):
        seq_len = x.size(1)
        return x + self.pe[:, :seq_len]

The input stage combines an embedding with a scaling factor:

class InputEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.scale = math.sqrt(d_model)

    def forward(self, x):
        return

Multiplying by sqrt(d_model) follows the original design and keeps embedding magnitudes in proportion with the rest of the network.

Step 2: Scaled dot-product attention

This is the core of the model. Given queries Q, keys K and values V, attention multiplies Q by the transpose of K, divides by the square root of the key dimension, applies a softmax and uses the resulting weights to mix V.

The division by sqrt(d_k) matters. Dot products grow along with the dimension, and without scaling the softmax would become far too peaked.

def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))

    attn_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, v)
    return output, attn_weights

This small function lets every token decide dynamically which other tokens deserve attention. In self-attention, queries, keys and values all come from one sequence. In encoder-decoder attention, the decoder queries the encoder's outputs.

Step 3: Multi-head attention

One big attention pattern gets replaced by several smaller ones. The inputs are projected into multiple subspaces, attention runs in each in parallel, and the results are concatenated and projected back to the model dimension.

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

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // 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 split_heads(self, x):
        batch_size, seq_len, d_model = x.size()
        x = x.view(batch_size, seq_len, self.num_heads, self.d_k)
        return x.transpose(1, 2)  # (batch, heads, seq, d_k)

    def combine_heads(self, x):
        batch_size, heads, seq_len, d_k = x.size()
        x = x.transpose(1, 2).contiguous()
        return x.view(batch_size, seq_len, heads * d_k)

    def forward(self, q, k, v, mask=None):
        q = self.split_heads(self.w_q(q))
        k = self.split_heads(self.w_k(k))
        v = self.split_heads(self.w_v(v))

        attn_output, attn_weights = scaled_dot_product_attention(q, k, v, mask)
        output = self.combine_heads(attn_output)
        return self.w_o(output), attn_weights

PyTorch's module system fits well here, because every learned projection is just a linear layer inside an nn.Module.

Step 4: Feed-forward network and layer norm

After attention, each position passes through the same feed-forward network. The original design uses two linear layers with a nonlinearity between them. The output joins the residual path and is normalized.

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x):
        return self.net(x)

PyTorch's LayerNorm normalizes over the trailing feature dimensions and can learn per-element scale and bias values, which is how normalization is used in Transformer stacks. With those pieces ready, an encoder layer is easy to assemble:

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = FeedForward(d_model, d_ff, dropout)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        attn_out, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))

        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_out))
        return x

Step 5: The encoder stack

The original model stacks several identical encoder layers. In PyTorch that's usually done with nn.ModuleList, which holds submodules and also registers them for optimization and checkpointing.

class Encoder(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, dropout=0.1):
        super().__init__()
        self.embed = InputEmbedding(vocab_size, d_model)
        self.pos = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.dropout = nn.Dropout(dropout)

    def forward(self, src, mask=None):
        x = self.embed(src)
        x = self.pos(x)
        x = self.dropout(x)

        for layer in self.layers:
            x = layer(x, mask)

        return x

At this point there's already a usable encoder for tasks such as classification or learning sequence representations. The complete encoder-decoder model then needs a decoder stack plus a final projection onto the vocabulary, an arrangement the paper describes and the Harvard implementation mirrors.

Step 6: The decoder

A complete decoder layer introduces two extra ideas:

1. Masked self-attention, which stops a token from looking at later tokens during training.
2. Cross-attention, where the decoder attends to the encoder's outputs.

In rough pseudocode, one decoder layer works like this:

# Pseudocode structure
x = norm1(x + masked_self_attention(x, x, x))
x = norm2(x + cross_attention(x, encoder_output, encoder_output))
x = norm3(x + feed_forward(x))

Once the encoder makes sense, the decoder is mostly the same thing again, with one more attention block and a causal mask.

Checking the shapes as you go

A habit that saves hours is to push a random batch through every piece the moment it's written. Feed a tensor of shape (batch, sequence, model dimension) into the positional encoding and the shape shouldn't change. Send it through multi-head attention and the output should keep that shape too, while the attention weights come back as (batch, heads, sequence, sequence). The feed-forward block widens the last dimension internally and returns it to the model size, so the output shape matches the input again.

If any of those checks fails, the problem is almost always in how heads are split or merged. Catching it at that stage is much cheaper than discovering it after a long training run that quietly produces garbage.

Another useful test is a deliberately tiny model. Two layers, a small model dimension and a toy dataset such as copying a short sequence are enough to see whether the loss actually falls. A model that can't learn to copy has a bug somewhere, and a model that can is already doing the hard part right.

Frequent bugs

When a from-scratch Transformer refuses to train, the cause is usually one of these:

• tensor shapes that go wrong while splitting or merging heads
• mask broadcasting errors, particularly in decoder self-attention
• positional encoding left out, which makes attention blind to order
• normalization placed differently from the variant being compared
• layers kept in a plain Python list instead of a ModuleList, so they never get registered

The architecture is elegant. Problems tend to hide in the plumbing.

What to take away

No single component is overwhelmingly complicated. The strength comes from the combination, which stays simple, parallelizes well and expresses a lot. The ingredients are learned embeddings, positional signals, self-attention, per-position feed-forward layers, residual shortcuts and normalization. That recipe became the base of modern sequence models.

To really understand it, go past the diagrams. Build one, print the tensor shapes, watch the attention scores shift, and break the masks on purpose before fixing them. After that, the model stops feeling like magic and starts looking like engineering.