Practice: Transformers

Implementation exercises for Transformers — building scaled dot-product attention, positional encoding, and multi-head attention from scratch, then verifying against PyTorch's built-in implementation.

🟡 Problem 1: Implement scaled dot-product attention from scratch

Task: Implement the core attention formula using plain PyTorch tensor operations (no nn.MultiheadAttention), and verify the output shape.

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

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

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)   # (batch, seq_len, seq_len)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    attn_weights = F.softmax(scores, dim=-1)
    output = attn_weights @ V
    return output, attn_weights

batch, seq_len, d_k = 2, 5, 16
Q = torch.randn(batch, seq_len, d_k)
K = torch.randn(batch, seq_len, d_k)
V = torch.randn(batch, seq_len, d_k)

output, weights = scaled_dot_product_attention(Q, K, V)
print(output.shape)     # (2, 5, 16)
print(weights.shape)    # (2, 5, 5) -- one attention weight per pair of positions
print(weights.sum(dim=-1))   # should be all 1.0 -- softmax rows sum to 1

🟡 Problem 2: Implement causal (masked) self-attention

Task: Extend Problem 1 to add a causal mask, so each position can only attend to itself and earlier positions — required for decoder-style autoregressive generation.

def create_causal_mask(seq_len):
    return torch.tril(torch.ones(seq_len, seq_len))   # lower triangular: 1s on/below diagonal

mask = create_causal_mask(seq_len)
output, weights = scaled_dot_product_attention(Q, K, V, mask=mask)

print(weights[0])   # inspect: position 0 should only attend to position 0;
                       # position 4 can attend to positions 0 through 4

Hint if stuck: torch.tril creates a matrix with 1s on and below the diagonal, 0s above — exactly the pattern needed so position \(i\) can attend to positions \(0 \dots i\) but not beyond.

🔴 Problem 3: Implement multi-head attention from scratch

Task: Extend single-head attention into multi-head attention — splitting Q/K/V into multiple heads, computing attention independently per head, then concatenating and projecting the results.

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

        self.W_q = torch.nn.Linear(d_model, d_model)
        self.W_k = torch.nn.Linear(d_model, d_model)
        self.W_v = torch.nn.Linear(d_model, d_model)
        self.W_o = torch.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)
        # shapes are now (batch, num_heads, seq_len, d_k)

        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
        weights = F.softmax(scores, dim=-1)
        attn_output = weights @ V   # (batch, num_heads, seq_len, d_k)

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

mha = MultiHeadAttention(d_model=64, num_heads=8)
x = torch.randn(2, 10, 64)
output = mha(x)
print(output.shape)   # (2, 10, 64) -- same shape as input, as expected

# Verify against PyTorch's built-in for sanity on shapes (weights will differ since untrained)
builtin_mha = torch.nn.MultiheadAttention(embed_dim=64, num_heads=8, batch_first=True)
builtin_output, _ = builtin_mha(x, x, x)
print(builtin_output.shape)   # should match: (2, 10, 64)

Hint if stuck: The trickiest part is the reshape/transpose to split into heads and back — draw out the tensor shape at each line on paper if you get confused; the key insight is that d_model gets split evenly across num_heads, each with dimension d_k = d_model / num_heads.

🟡 Problem 4: Implement sinusoidal positional encoding

Task: Implement the original Transformer paper's sinusoidal positional encoding and visualize it as a heatmap.

\[ PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d}}\right), \quad PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d}}\right) \]

def positional_encoding(seq_len, d_model):
    pe = torch.zeros(seq_len, d_model)
    position = torch.arange(0, seq_len).unsqueeze(1).float()
    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)
    return pe

pe = positional_encoding(seq_len=50, d_model=64)
print(pe.shape)   # (50, 64)

import matplotlib.pyplot as plt
plt.imshow(pe.numpy(), cmap='RdBu', aspect='auto')
plt.xlabel('Embedding dimension')
plt.ylabel('Position in sequence')
plt.title('Sinusoidal Positional Encoding')
plt.show()

Related DL Notes

• Practice: RNN & LSTM
• Practice: PyTorch
• Transformer Interview Questions