The Transformer Architecture: Self-Attention Under the Hood
Before the introduction of the Transformer architecture in 2017, Recurrent Neural Networks (RNNs) and Long Short-Term Memory networks (LSTMs) ruled sequential data processing. However, they had a fundamental limitation: they process tokens sequentially, one step at a time. This sequence-dependent processing makes parallelization impossible, limiting training speed on modern GPUs.
Transformers solved this with the Self-Attention mechanism. By looking at all words in a sequence simultaneously and calculating how much attention to pay to each, the Transformer achieved state-of-the-art results while enabling massive parallelization.
Transformer Architecture Breakdown: 3D block diagram of Transformer encoder-decoder architecture
1. The Core Idea: Query, Key, and Value
To understand self-attention, think of it like searching a database:
- A Query (
Q) represents the token currently being processed (what we are looking for). - A Key (
K) represents all the tokens in the sequence (what we compare the query against). - A Value (
V) represents the actual content of the tokens (what we retrieve).
For each input token vector x, we project it into Query, Key, and Value vectors using weight matrices W^Q, W^K, and W^V:
q = x W^Q, \quad k = x W^K, \quad v = x W^V
2. Scaled Dot-Product Attention
The similarity between a Query and a Key is calculated using their dot product. Higher values indicate that the tokens are highly related. The mathematical formulation of Scaled Dot-Product Attention is:
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V
Where:
Q, K, Vare the query, key, and value matrices.d_kis the dimension of the key vectors.\sqrt{d_k}is the scaling factor.
Why do we scale by \sqrt{d_k}?
For large values of d_k, the dot product grows large in magnitude. This pushes the softmax function into regions with extremely small gradients (vanishing gradients). Dividing by \sqrt{d_k} stabilizes training by keeping the values in a range where softmax has active gradients.
Q Matrix K^T Matrix Attention Map V Matrix Final Output
[ q1 ] x [ k1 k2 ] ---> [ s11 s12 ] x [ v1 ] ---> [ o1 ]
[ q2 ] [ s21 s22 ] [ v2 ] [ o2 ]
3. Multi-Head Attention
Instead of performing self-attention once, the Transformer splits the queries, keys, and values into multiple subspaces, allowing the model to attend to information from different representation subspaces at different positions simultaneously.
The formulas for Multi-Head Attention are:
\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O
\text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V)
Where the projection matrices are:
W_i^Q \in \mathbb{R}^{d_{model} \times d_k}W_i^K \in \mathbb{R}^{d_{model} \times d_k}W_i^V \in \mathbb{R}^{d_{model} \times d_v}W^O \in \mathbb{R}^{h d_v \times d_{model}}
4. PyTorch Implementation of Multi-Head Attention
Let's write a complete, raw implementation of Multi-Head Attention in PyTorch using tensor reshaping operations.
import torch
import torch.nn as nn
import math
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_heads
# Projection matrices
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 forward(self, q, k, v, mask=None):
batch_size = q.size(0)
# 1. Project inputs and split into heads
# Shape changes: (batch, seq_len, d_model) -> (batch, seq_len, num_heads, d_k) -> transpose to (batch, num_heads, seq_len, d_k)
Q = self.w_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
K = self.w_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
V = self.w_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
# 2. Compute Scaled Dot-Product Attention scores
# (batch, num_heads, seq_len, d_k) x (batch, num_heads, d_k, seq_len) -> (batch, num_heads, seq_len, seq_len)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
if mask is not None:
# Mask out invalid padding or future tokens (causal masking)
scores = scores.masked_fill(mask == 0, -1e9)
attention_weights = torch.softmax(scores, dim=-1)
# 3. Multiply attention weights by values
# (batch, num_heads, seq_len, seq_len) x (batch, num_heads, seq_len, d_k) -> (batch, num_heads, seq_len, d_k)
context = torch.matmul(attention_weights, V)
# 4. Concatenate heads back and apply output projection
# Transpose back: (batch, num_heads, seq_len, d_k) -> (batch, seq_len, num_heads, d_k)
context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
return self.w_o(context)
# Testing the implementation
batch_size, seq_len, d_model, num_heads = 2, 8, 64, 4
qkv_tensor = torch.randn(batch_size, seq_len, d_model)
mha = MultiHeadAttention(d_model=d_model, num_heads=num_heads)
out = mha(qkv_tensor, qkv_tensor, qkv_tensor)
print("Output tensor shape:", out.shape) # Expected: [2, 8, 64]5. Positional Encoding
Because the Self-Attention mechanism computes attention across all tokens in parallel, it is completely invariant to sequence order. In other words, the sequence "the cat ate the fish" and "the fish ate the cat" would result in identical self-attention representations.
To restore sequence order, we inject Positional Encodings into the input token embeddings. These are vectors containing unique periodic signals based on sine and cosine functions:
PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)
PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)
This mathematical formulation gives the model a unique, geometry-based positional coordinate for every token.
Conclusion
The Transformer architecture's shift away from recurrence to self-attention marked the start of the modern LLM era. By projecting representations into Queries, Keys, and Values, computing parallelizable dot-product weights, and layering multi-head tracking, it handles sequence data with unprecedented scale and speed.