Transformer from scratch in PyTorch

In this tutorial, we implement the transformer from a conceptual description to an actual code implementation.

How does a transformer architecture work?

Creating embeddings

A transformer operates over a sequence of embeddings, each of dimension , often with positional information already added. A good tokenization should maximize the unconditional entropy of tokens in a sequence. The reason for this can be seen in the extreme: imagine a tokenization where, 70% of the time, the unconditionally correct token to predict was some special token ; there would be far less signal pushing your model to learn the conditional distribution rather than fit the unconditional distribution.

Transforming embeddings

The transformer is fundamentally a residual network atop the input embeddings. The residual is computed by performing the following flow:

Below explains the blocks in more detail.

LayerNorm

LayerNorm normalizes each embedding vector independently so that the statistics across its dimensions are well behaved. Specifically, it transforms the vector so that its values have a mean of 0 and a standard deviation of 1.

Multi-Head Self-Attention

Self-attention works by linearly projecting each input embedding (of dimension ) in the sequence into query , key , and value spaces. The attention layer computes how similar each query is to each key (this is a quadratic operation, which is why standard attention is quadratic), producing a square “similarity” matrix. This similarity matrix is then masked (e.g., using causal masking in LLMs), and a softmax distribution is computed over these similarities for each embedding (i.e., over the second dimension of the attention matrix). This resulting attention distribution determines the weighting of how values from the value projection are combined to form the attention output.

The above describes how single-head attention works. For multi-head attention with distinct heads, each head projects the input embedding from dimension down to and performs attention in space. The results of the attention heads are then concatenated into a vector of dimension and mixed with a learned linear output projection before being added back to the residual stream.

Position-wise MLP

The position-wise MLP is a hidden layer MLP (typically with a hidden size of ) run on each embedding vector independently.

Implementing a Transformer in practice

Below is a reference implementation of the transformer in PyTorch:

import torch

import torch.nn as nn

class MyTransformer(nn.Module):

    def __init__(self, D_model : int, num_heads : int):
        super().__init__()
        self.D_model = D_model
        self.num_heads = num_heads
        self.D_head = D_model // num_heads
        assert D_model % num_heads == 0, f"D_model ({D_model}) needs to be evenly divisible by the number of heads ({num_heads})"

        # LayerNorm
        self.attention_layer_norm = nn.LayerNorm(D_model)
        self.mlp_layer_norm = nn.LayerNorm(D_model)

        # Attention stream projections
        self.kqv_projection = nn.Linear(D_model, 3 * D_model)
        self.output_mixer_projection = nn.Linear(D_model, D_model)

        # Embedding MLP projections
        self.embedding_mlp = nn.Sequential(
            nn.Linear(D_model, 4 * D_model), nn.GELU(), nn.Linear(4 * D_model, D_model)
        )

    def _attention_residual_stream(self, X : torch.Tensor, mask : torch.Tensor) -> torch.Tensor:
        B, S, _ = X.shape

        # (batch_size, seq_len, D_model) -> (batch_size, seq_len, 3 * D_model)
        kqv = self.kqv_projection(X)

        # Split into chunks to run the multihead projection
        # (batch_size, seq_len, 3 * D_model) -> (batch_size, seq_len, 3 , num_heads, D_head)
        kqv = kqv.reshape(B, S, 3, self.num_heads, self.D_head)

        # Move 3 to the front so we can disentangle to K, Q, V, and make it so every head is S x D_head
        # (batch_size, seq_len, 3 , num_heads, D_head) -> (3, batch_size, num_heads, seq_length, D_head)
        kqv = kqv.permute((2, 0, 3, 1, 4))

        # Disentangle K, Q, V
        # Each matrix is (batch_size, num_heads, seq_length, D_head)
        K, Q, V = kqv.unbind(dim=0)

        # Compute attention scores
        # (batch_size, num_heads, seq_length, D_head) @ (batch_size, num_heads, D_head, seq_length) -> (batch_size, num_heads, seq_length, seq_length)
        scores = Q @ K.transpose(-2, -1) / (self.D_head ** 0.5)

        # Apply the mask, which should be of shape (seq_length, seq_length)
        scores = scores.masked_fill(~mask, float("-inf"))

        # (batch_size, num_heads, seq_length, seq_length)
        softmaxed_weights = torch.softmax(scores, dim=-1)

        # (batch_size, num_heads, seq_length, seq_length) @ (batch_size, num_heads, seq_length, D_head) -> (batch_size, num_heads, seq_length, D_head)
        weighted_values = softmaxed_weights @ V

        # (batch_size, num_heads, seq_length, D_head) -> (batch_size, seq_length, num_heads, D_head)
        weighted_values = weighted_values.transpose(1, 2)

        # Convert head outputs into expected shape (batch_size, seq_length, D_model)
        weighted_values = weighted_values.reshape((B, S, self.D_model))

        return self.output_mixer_projection(weighted_values)


    def _mlp_residual_stream(self, X: torch.Tensor):
        B, S, _ = X.shape

        # Flatten everything into a giant batch dimension
        flattened_X = X.reshape((B * S, self.D_model))
        embedded_flattened_X = self.embedding_mlp(flattened_X)

        return embedded_flattened_X.reshape((B, S, self.D_model))

    def forward(self, X_input : torch.Tensor):
        # X shape: (batch_size, seq_len, D_model)
        B, S, _ = X_input.shape

        # Create a mask for the attention mechanism
        mask = torch.tril(torch.ones(S, S, dtype=torch.bool, device=X_input.device)).unsqueeze(0).unsqueeze(0)  # (1, 1, seq_len, seq_len)

        # Residual Block 1
        X_attention_residual = self._attention_residual_stream(self.attention_layer_norm(X_input), mask)
        X_post_attention = X_input + X_attention_residual

        # Residual Block 2
        X_mlp_residual = self._mlp_residual_stream(self.mlp_layer_norm(X_post_attention))
        return X_post_attention + X_mlp_residual