Building a Transformer in Pytorch

Building a Transformer from scratch in PyTorch is a rite of passage for modern AI practitioners. While libraries like Hugging Face are great for production, understanding the “guts” involves implementing the architecture introduced in the Attention is All You Need paper.

To keep this digestible, we’ll break the construction into its core components: Scaled Dot-Product Attention, Multi-Head Attention, and the Transformer Block.

---

1. Scaled Dot-Product Attention

This is the heart of the model. It calculates how much focus one word should have on another by comparing a Query (Q) and a Key (K), then applying that weight to the Value (V).

Mathematically:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Python

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

def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    # Compute attention scores
    scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, v), weights

---

2. Multi-Head Attention (MHA)

Instead of performing a single attention function, we split the embedding dimension into “heads.” This allows the model to simultaneously attend to information from different representation subspaces (e.g., one head focuses on grammar, another on the subject-verb relationship).

Python

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_model = d_model
        assert d_model % num_heads == 0
        
        self.d_k = d_model // num_heads
        
        # Linear layers for Q, K, V projections
        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.out_linear = nn.Linear(d_model, d_model)

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)
        
        # 1. Linear projections and split into heads
        q = self.q_linear(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = self.k_linear(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = self.v_linear(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 2. Apply scaled dot-product attention
        x, _ = scaled_dot_product_attention(q, k, v, mask)
        
        # 3. Concatenate heads and put back through final linear layer
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        return self.out_linear(x)

---

3. The Transformer Block (Encoder Layer)

A standard Transformer layer consists of the MHA we just built, followed by a Feed-Forward Network (FFN). Crucially, we use Residual Connections and Layer Normalization to keep the gradients stable.

Python

class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        self.attention = MultiHeadAttention(d_model, num_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.ReLU(),
            nn.Linear(4 * d_model, d_model)
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Sublayer 1: Attention + Residual
        attn_out = self.attention(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))
        
        # Sublayer 2: Feed Forward + Residual
        ff_out = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_out))
        return x

---

4. Putting it Together: The Encoder

A full Transformer Encoder is simply a stack of these blocks, preceded by an Embedding layer and Positional Encoding (since Transformers have no inherent sense of word order).

Python

class TransformerEncoder(nn.Module):
    def __init__(self, vocab_size, d_model, num_layers, num_heads, max_len=512):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        # Simplified positional encoding for brevity
        self.pos_embedding = nn.Embedding(max_len, d_model)
        
        self.layers = nn.ModuleList([
            TransformerBlock(d_model, num_heads) for _ in range(num_layers)
        ])

    def forward(self, x, mask=None):
        batch_size, seq_len = x.size()
        positions = torch.arange(0, seq_len).expand(batch_size, seq_len).to(x.device)
        
        out = self.embedding(x) + self.pos_embedding(positions)
        
        for layer in self.layers:
            out = layer(out, mask)
        return out

---

Key Takeaways

• Permutation Invariance: Transformers don’t “see” order unless you add Positional Encodings.
• Complexity: The self-attention mechanism has a computational complexity of $O(n^2 \cdot d)$, where $n$ is the sequence length. This is why long documents are harder for Transformers to process.
• The “Why”: Unlike RNNs, Transformers allow for massive parallelization during training because the entire sequence is processed at once.