Transformer Mechanics Scratch Implementation

Building the Engine

Using a high-level API is one thing; understanding the machinery under the hood is another. For a senior engineer, mastering the internal tensor flow of a Transformer is non-negotiable. We will build a decoder-only Transformer from scratch in PyTorch, focusing on the architecture that powers models like GPT.

This process involves constructing the core components piece by piece: the attention mechanism, the multi-head structure that gives it power, the feed-forward network, and the residual connections that hold it all together. We will track tensor dimensions meticulously at every step, as this is the key to truly understanding the data's journey through the model.

The Core Mechanism: Scaled Dot-Product Attention

At its heart, attention is a way for a model to weigh the importance of different parts of the input sequence. It operates on three inputs: a Query ( $Q$ ), a Key ( $K$ ), and a Value ( $V$ ). The query represents the current position's request for information. The keys from all other positions are compared to this query to calculate a similarity score. These scores are then used to create a weighted sum of the values, effectively highlighting the most relevant information for the current query.

This entire process is captured in a single, elegant formula.

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

For a decoder, we must ensure that the model can't cheat by looking ahead at future tokens during training. We enforce this with causal masking, a mechanism that sets the attention scores for future positions to negative infinity before the softmax step. This effectively makes their probability zero.

import torch
import torch.nn as nn
import math

def scaled_dot_product_attention(q, k, v, mask=None):
    # q, k, v shape: (batch_size, num_heads, seq_len, d_head)
    d_k = q.size(-1)
    
    # Matmul for scores: (..., seq_len, d_head) @ (..., d_head, seq_len) -> (..., seq_len, seq_len)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
    
    # Apply mask if provided
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9) # -1e9 is equivalent to -infinity
        
    # Softmax to get attention weights
    attention_weights = F.softmax(scores, dim=-1)
    
    # Matmul weights with values
    # (..., seq_len, seq_len) @ (..., seq_len, d_head) -> (..., seq_len, d_head)
    output = torch.matmul(attention_weights, v)
    
    return output, attention_weights

Multiple Perspectives: Multi-Head Attention

Instead of performing a single attention calculation, (MHA) splits the model's embedding space into multiple smaller "heads." Each head performs attention independently, allowing the model to focus on different types of information simultaneously. For instance, one head might track syntactic relationships while another focuses on semantic connections.

The implementation involves projecting the input embedding into Q, K, and V for each head, running scaled dot-product attention in parallel, concatenating the results, and finally passing them through a final linear layer.

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_head = 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):
        # Reshape from (batch_size, seq_len, d_model) to (batch_size, num_heads, seq_len, d_head)
        x = x.view(batch_size, -1, self.num_heads, self.d_head)
        return x.transpose(1, 2)
        
    def forward(self, q, k, v, mask):
        batch_size = q.size(0)
        
        # 1. Project Q, K, V
        q = self.W_q(q)
        k = self.W_k(k)
        v = self.W_v(v)
        
        # 2. Split into heads
        q = self.split_heads(q, batch_size)
        k = self.split_heads(k, batch_size)
        v = self.split_heads(v, batch_size)
        
        # 3. Apply scaled dot-product attention
        attention_output, _ = scaled_dot_product_attention(q, k, v, mask)
        
        # 4. Concatenate heads and project
        attention_output = attention_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        output = self.W_o(attention_output)
        
        return output

Assembling the Transformer Block

A full Transformer block consists of more than just attention. Each block contains two main sub-layers: a Multi-Head Attention mechanism and a simple, position-wise feed-forward network. Crucially, each of these sub-layers has a around it, followed by (LayerNorm). The residual connection (or skip connection) adds the input of the sub-layer to its output, which helps prevent the vanishing gradient problem in deep networks. LayerNorm then normalizes the output to stabilize training.

The placement of LayerNorm is a key design choice. Pre-LN applies normalization before the sub-layer, while Post-LN applies it after the residual connection. Pre-LN generally leads to more stable training for very deep models, which is why it's common in modern architectures. We will implement Pre-LN.

class PositionWiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.linear2(self.relu(self.linear1(x)))

class DecoderBlock(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.feed_forward = PositionWiseFeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask):
        # Pre-LN implementation
        
        # Self-attention sub-layer
        norm_x = self.norm1(x)
        attn_output = self.self_attn(norm_x, norm_x, norm_x, mask)
        x = x + self.dropout(attn_output) # Residual connection
        
        # Feed-forward sub-layer
        norm_x = self.norm2(x)
        ff_output = self.feed_forward(norm_x)
        x = x + self.dropout(ff_output) # Residual connection
        
        return x

With the core block complete, we can stack them to create the full decoder. We also need to handle the initial token embeddings and add positional information, since the self-attention mechanism itself has no inherent sense of sequence order.

Applied AI Engineering for Frontier Labs