INTRODUCTION: THE ARCHITECTURE THAT CHANGED EVERYTHING
In 2017, a team at Google published a paper with an audacious title: "Attention Is All You Need." This wasn't just academic bravado. The Transformer architecture they introduced fundamentally changed how we build AI systems. Today, every major language model from GPT-4 to Claude to Gemini builds on this foundation. If you've used any modern AI tool, you've interacted with a Transformer.
As a software engineer, you might have heard that Transformers are "just matrix multiplications" or that they "use attention mechanisms." While technically true, these descriptions miss the elegant mathematical reasoning that makes Transformers work. This article will take you on a journey from the problems that motivated Transformers to the cutting-edge optimizations used in production systems today. By the end, you'll understand not just what the formulas are, but why they must be exactly as they are, and you'll have the knowledge to implement a Transformer from scratch.
We'll follow a single red thread throughout: how do we build a system that can understand the relationships between elements in a sequence, regardless of how far apart those elements are, and do so efficiently enough to train on billions of examples? Every formula, every architectural choice, stems from answering this question.
THE PROBLEM: WHY RECURRENT NETWORKS FAILED US
Before Transformers, the dominant approach to sequence processing was the Recurrent Neural Network (RNN). To understand why Transformers exist, we need to understand why RNNs, despite their elegance, couldn't scale to the problems we wanted to solve.
Imagine you're building a translation system. You receive the English sentence "The cat sat on the mat" and need to produce the French translation. An RNN processes this sequentially. It reads "The," updates its internal state, then reads "cat," updates its state again, and so on. The idea is that by the time it finishes reading the sentence, its hidden state contains a compressed representation of everything it has seen.
Here's how an RNN updates its hidden state at each time step:
h_t = tanh(W_hh * h_{t-1} + W_xh * x_t + b_h)
In this formula, h_t represents the hidden state at time step t, x_t is the input at time t, W_hh is a weight matrix that transforms the previous hidden state, W_xh transforms the current input, and b_h is a bias term. The tanh function squashes the result to keep values bounded.
The problem becomes apparent when you think about what happens over many time steps. Suppose you're translating a long document, and a critical piece of information appears in the first sentence, but you need it to correctly translate the hundredth sentence. That information must survive being multiplied by W_hh one hundred times. In practice, this leads to two catastrophic problems.
First, the vanishing gradient problem. When training neural networks, we compute gradients that tell us how to adjust weights. These gradients must flow backward through time. If the gradient at step 100 needs to influence the weights that processed step 1, it must be multiplied by the derivative of the tanh function one hundred times. Since tanh's derivative is always less than one, the gradient shrinks exponentially. By the time it reaches the early steps, it's essentially zero, and those weights never learn.
Second, even if we solve vanishing gradients with techniques like LSTMs or GRUs, there's a more fundamental problem: sequential processing is slow. Modern GPUs excel at parallel computation. They can multiply massive matrices incredibly quickly. But RNNs force us to process sequences one step at a time because each step depends on the previous one. You cannot compute h_100 until you've computed h_99, which requires h_98, and so on back to h_1. This sequential dependency makes training painfully slow on the massive datasets we need for modern AI.
The Transformer architecture solves both problems with a radical idea: what if we could look at all positions in the sequence simultaneously and let the model learn which positions are relevant to each other? This is where attention comes in.
THE CORE INSIGHT: ATTENTION AS A DIFFERENTIABLE LOOKUP TABLE
SCALED DOT-PRODUCT ATTENTION: THE MATHEMATICAL FOUNDATION
Now we're ready to derive the actual attention formula used in Transformers. We'll build it step by step, understanding why each component exists.
Our goal is to create a mechanism where, for each position in a sequence, we can compute a representation that incorporates information from all other positions based on learned relevance. We need three things:
First, a way to express "what I'm looking for" at each position. We call this the Query (Q). Think of it as a search query in a database.
Second, a way to express "what I have to offer" at each position. We call this the Key (K). Think of it as the index in a database that we match against.
Third, the actual information we want to retrieve. We call this the Value (V). Think of it as the data stored in the database.
* For better understanding these important concepts you‘ll find more explanations in the ADDENDUM of this article.
Why separate Keys and Values? Because what we use to determine relevance (the Key) might be different from what we actually want to retrieve (the Value). For example, when translating "The cat sat," the word "cat" might be relevant because it's a noun (that's what the Key captures), but what we actually want to retrieve is its full semantic meaning (that's what the Value captures).
We create Q, K, and V by multiplying our input by learned weight matrices:
Q = X * W_Q
K = X * W_K
V = X * W_V
Here, X is our input matrix where each row is a word embedding, and W_Q, W_K, W_V are learnable weight matrices. These matrices are learned during training to extract the right kind of information for queries, keys, and values respectively.
Now, to compute attention, we need to measure how well each query matches each key. The natural choice is the dot product because it measures similarity: if two vectors point in the same direction, their dot product is large; if they're orthogonal, it's zero; if they point in opposite directions, it's negative.
We compute all query-key similarities at once by multiplying Q and K transposed:
scores = Q * K^T
This gives us a matrix where entry (i,j) tells us how much position i (the query) should attend to position j (the key). The dimensions work out as follows: if we have n positions and each query/key is d-dimensional, then Q is n by d, K^T is d by n, and their product is n by n.
Here's where we encounter our first problem. As the dimension d gets larger, the dot products get larger in magnitude. To see why, consider that a dot product is a sum of d terms. If each term has variance sigma squared, the sum has variance d times sigma squared by the properties of variance. This means the dot product grows with the square root of d.
Why is this a problem? Because we're going to apply softmax to these scores. Softmax with very large inputs becomes extremely peaked. To see this, consider:
softmax([10, 9, 8]) = [0.665, 0.245, 0.090]
softmax([100, 90, 80]) = [0.9999, 0.0001, 0.0000]
When the inputs to softmax are very large, it essentially becomes a hard selection, picking only the largest value and ignoring everything else. This kills gradients during training because the derivative of softmax becomes nearly zero everywhere except at the maximum.
The solution is to scale the dot products by the square root of the dimension:
scaled_scores = (Q * K^T) / sqrt(d_k)
This scaling ensures that regardless of the dimension, the variance of the scores stays roughly constant. The square root specifically counteracts the square root growth we identified earlier.
Now we apply softmax to get attention weights:
attention_weights = softmax(scaled_scores)
Finally, we use these weights to take a weighted average of the values:
output = attention_weights * V
Putting it all together, we get the scaled dot-product attention formula:
Attention(Q, K, V) = softmax((Q * K^T) / sqrt(d_k)) * V
Let's implement this in code to make it concrete:
import numpy as np
def scaled_dot_product_attention(Q, K, V):
"""
Compute scaled dot-product attention.
Args:
Q: Query matrix of shape (n, d_k) where n is sequence length
K: Key matrix of shape (n, d_k)
V: Value matrix of shape (n, d_v)
Returns:
output: Attention output of shape (n, d_v)
attention_weights: Attention weight matrix of shape (n, n)
"""
# Get the dimension of keys for scaling
d_k = Q.shape[-1]
# Compute attention scores by taking dot product of queries and keys
# Shape: (n, d_k) @ (d_k, n) = (n, n)
scores = np.matmul(Q, K.transpose(-2, -1))
# Scale scores by square root of key dimension
# This prevents softmax from becoming too peaked
scaled_scores = scores / np.sqrt(d_k)
# Apply softmax to get attention weights
# Each row sums to 1, representing a probability distribution
attention_weights = softmax(scaled_scores)
# Compute weighted sum of values
# Shape: (n, n) @ (n, d_v) = (n, d_v)
output = np.matmul(attention_weights, V)
return output, attention_weights
def softmax(x):
"""
Compute softmax values for each row of matrix x.
Numerically stable implementation that subtracts max value.
"""
# Subtract max for numerical stability
# This prevents overflow from exp of large numbers
exp_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
This is the core of the Transformer. Everything else builds on this foundation. Notice how the formula emerged from first principles: we wanted a differentiable lookup mechanism, we chose dot products for similarity, we scaled to prevent saturation, and we used softmax to get normalized weights.
MULTI-HEAD ATTENTION: LEARNING MULTIPLE PERSPECTIVES
We now have a working attention mechanism, but there's a limitation. A single attention operation can only learn one type of relationship between positions. In language, we might want to capture many different types of relationships simultaneously. For example, we might want one attention head to focus on syntactic relationships (which words modify which), another on semantic relationships (which words have similar meanings), and another on positional relationships (which words are nearby).
This is the motivation for multi-head attention. Instead of computing attention once, we compute it multiple times in parallel with different learned projections. Each "head" can learn to attend to different aspects of the input.
Here's how it works mathematically. Instead of having single weight matrices W_Q, W_K, and W_V, we have h sets of them, where h is the number of heads:
For head i:
Q_i = X * W_Q^i
K_i = X * W_K^i
V_i = X * W_V^i
Each head computes its own attention:
head_i = Attention(Q_i, K_i, V_i)
Now we have h different attention outputs, each capturing different relationships. We concatenate them and project back to the original dimension:
MultiHead(Q, K, V) = Concat(head_1, head_2, ..., head_h) * W_O
The concatenation combines all the different perspectives, and W_O is a learned weight matrix that mixes them together in useful ways.
There's an important implementation detail here. If we have d_model as our model dimension and h heads, we typically make each head's dimension d_model / h. This way, the total computational cost stays the same as single-head attention, but we get the benefit of multiple perspectives.
Let's implement multi-head attention:
class MultiHeadAttention:
"""
Multi-head attention mechanism.
Allows the model to jointly attend to information from different
representation subspaces at different positions.
"""
def __init__(self, d_model, num_heads):
"""
Initialize multi-head attention.
Args:
d_model: Dimension of the model (must be divisible by num_heads)
num_heads: Number of attention heads
"""
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 # Dimension per head
# Weight matrices for all heads combined
# We use a single large matrix and split it, which is more efficient
self.W_Q = np.random.randn(d_model, d_model) * 0.01
self.W_K = np.random.randn(d_model, d_model) * 0.01
self.W_V = np.random.randn(d_model, d_model) * 0.01
self.W_O = np.random.randn(d_model, d_model) * 0.01
def split_heads(self, x):
"""
Split the last dimension into (num_heads, d_k).
Reshape from (batch_size, seq_len, d_model) to
(batch_size, num_heads, seq_len, d_k)
"""
batch_size, seq_len, _ = x.shape
# Reshape to (batch_size, seq_len, num_heads, d_k)
x = x.reshape(batch_size, seq_len, self.num_heads, self.d_k)
# Transpose to (batch_size, num_heads, seq_len, d_k)
return x.transpose(0, 2, 1, 3)
def combine_heads(self, x):
"""
Inverse of split_heads.
Reshape from (batch_size, num_heads, seq_len, d_k) to
(batch_size, seq_len, d_model)
"""
batch_size, _, seq_len, _ = x.shape
# Transpose to (batch_size, seq_len, num_heads, d_k)
x = x.transpose(0, 2, 1, 3)
# Reshape to (batch_size, seq_len, d_model)
return x.reshape(batch_size, seq_len, self.d_model)
def forward(self, X):
"""
Compute multi-head attention.
Args:
X: Input tensor of shape (batch_size, seq_len, d_model)
Returns:
output: Attention output of shape (batch_size, seq_len, d_model)
"""
batch_size = X.shape[0]
# Linear projections for all heads at once
Q = np.matmul(X, self.W_Q) # (batch_size, seq_len, d_model)
K = np.matmul(X, self.W_K)
V = np.matmul(X, self.W_V)
# Split into multiple heads
Q = self.split_heads(Q) # (batch_size, num_heads, seq_len, d_k)
K = self.split_heads(K)
V = self.split_heads(V)
# Compute attention for all heads in parallel
# Scores shape: (batch_size, num_heads, seq_len, seq_len)
scores = np.matmul(Q, K.transpose(0, 1, 3, 2))
scores = scores / np.sqrt(self.d_k)
# Apply softmax
attention_weights = softmax(scores)
# Apply attention weights to values
# Output shape: (batch_size, num_heads, seq_len, d_k)
attention_output = np.matmul(attention_weights, V)
# Combine heads
# Shape: (batch_size, seq_len, d_model)
combined = self.combine_heads(attention_output)
# Final linear projection
output = np.matmul(combined, self.W_O)
return output
The beauty of multi-head attention is that it's a simple extension of single-head attention, but it dramatically increases the model's capacity to learn complex relationships. In practice, Transformers typically use eight or sixteen heads, allowing them to capture many different types of patterns simultaneously.
POSITIONAL ENCODING: TEACHING THE MODEL ABOUT ORDER
We now face a critical problem. Our attention mechanism is completely position-agnostic. If you shuffle the words in a sentence, you get exactly the same attention output. To see why, notice that the attention formula only depends on the content of Q, K, and V, not on where they appear in the sequence.
This is actually a feature of attention: it allows the model to look at relationships regardless of distance. But it's also a bug: word order matters in language. "Dog bites man" means something very different from "Man bites dog."
We need to inject positional information into our model. The question is: how? We could simply add a learned position embedding for each position, but this has a problem. If we train on sequences up to length 1000, what happens when we encounter a sequence of length 1001 at test time? We have no learned embedding for position 1001.
The Transformer paper introduced a clever solution: sinusoidal positional encodings. These are deterministic functions of position that have several desirable properties. The encoding for position pos and dimension i is:
PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
Let's unpack why this formula is designed this way. First, notice that we use different frequencies for different dimensions. The denominator 10000^(2i/d_model) grows exponentially with dimension i. This means early dimensions oscillate quickly (high frequency), while later dimensions oscillate slowly (low frequency).
Why use multiple frequencies? Because this allows the model to easily learn to attend by relative position. Suppose you want to know if two positions are exactly three steps apart. With multiple frequencies, there will be some frequency where the phase difference between positions that are three steps apart is consistent, regardless of absolute position.
The use of both sine and cosine for alternating dimensions is also deliberate. For any position pos, we can express the encoding at position pos + k (for any offset k) as a linear function of the encoding at position pos. This is because of the trigonometric identities:
sin(a + b) = sin(a)cos(b) + cos(a)sin(b)
cos(a + b) = cos(a)cos(b) - sin(a)sin(b)
This means the model can learn to attend to relative positions through linear transformations, which neural networks are very good at.
Let's implement positional encoding:
def get_positional_encoding(seq_len, d_model):
"""
Generate sinusoidal positional encodings.
Args:
seq_len: Length of the sequence
d_model: Dimension of the model
Returns:
pos_encoding: Positional encoding matrix of shape (seq_len, d_model)
"""
# Create a matrix to hold positional encodings
pos_encoding = np.zeros((seq_len, d_model))
# Create position indices: [0, 1, 2, ..., seq_len-1]
position = np.arange(seq_len).reshape(-1, 1)
# Create dimension indices: [0, 2, 4, ..., d_model-2]
# We use every other dimension for sine and cosine
div_term = np.exp(np.arange(0, d_model, 2) *
-(np.log(10000.0) / d_model))
# Apply sine to even indices
pos_encoding[:, 0::2] = np.sin(position * div_term)
# Apply cosine to odd indices
pos_encoding[:, 1::2] = np.cos(position * div_term)
return pos_encoding
To use positional encoding, we simply add it to our input embeddings:
# Get word embeddings from embedding layer
word_embeddings = embedding_layer(input_tokens) # Shape: (batch, seq_len, d_model)
# Get positional encodings
pos_encodings = get_positional_encoding(seq_len, d_model) # Shape: (seq_len, d_model)
# Add positional information to word embeddings
# Broadcasting handles the batch dimension automatically
input_with_position = word_embeddings + pos_encodings
The addition is element-wise. Now each position has a unique signature based on its position, and the model can learn to use this information to understand word order.
Modern variants of Transformers sometimes use learned positional embeddings or relative positional encodings, but the sinusoidal approach remains popular because it generalizes to any sequence length without retraining.
FEED-FORWARD NETWORKS: ADDING COMPUTATIONAL DEPTH
Attention allows positions to communicate and share information, but it's fundamentally a linear operation followed by a softmax. Even with multiple heads, we're just doing weighted averages. To give the model more computational power, we need non-linearity.
This is where the feed-forward network comes in. After attention, we apply a simple two-layer neural network to each position independently. The formula is:
FFN(x) = max(0, x * W_1 + b_1) * W_2 + b_2
Let's break this down. The first layer expands the dimension from d_model to a larger dimension d_ff (typically 4 times larger). This expansion gives the network more capacity to compute complex functions. We apply a ReLU activation (the max(0, ...) part), which introduces non-linearity. Then we project back down to d_model.
Why expand and then contract? This is a standard pattern in neural networks called a bottleneck. The expansion allows the network to compute in a higher-dimensional space where it might be easier to separate and transform features, and the contraction forces it to compress the useful information back to the original dimension.
The key insight is that this network is applied independently to each position. Unlike attention, which mixes information across positions, the feed-forward network processes each position separately. This division of labor is important: attention handles communication between positions, and the feed-forward network handles computation within each position.
Here's the implementation:
class FeedForwardNetwork:
"""
Position-wise feed-forward network.
Applies the same two-layer network to each position independently.
"""
def __init__(self, d_model, d_ff):
"""
Initialize feed-forward network.
Args:
d_model: Dimension of the model
d_ff: Dimension of the hidden layer (typically 4 * d_model)
"""
self.d_model = d_model
self.d_ff = d_ff
# First layer expands dimension
self.W_1 = np.random.randn(d_model, d_ff) * np.sqrt(2.0 / d_model)
self.b_1 = np.zeros(d_ff)
# Second layer contracts back to original dimension
self.W_2 = np.random.randn(d_ff, d_model) * np.sqrt(2.0 / d_ff)
self.b_2 = np.zeros(d_model)
def forward(self, x):
"""
Apply feed-forward network.
Args:
x: Input tensor of shape (batch_size, seq_len, d_model)
Returns:
output: Transformed tensor of same shape as input
"""
# First layer with ReLU activation
# Shape: (batch, seq_len, d_model) -> (batch, seq_len, d_ff)
hidden = np.matmul(x, self.W_1) + self.b_1
hidden = np.maximum(0, hidden) # ReLU activation
# Second layer
# Shape: (batch, seq_len, d_ff) -> (batch, seq_len, d_model)
output = np.matmul(hidden, self.W_2) + self.b_2
return output
The ReLU activation is crucial. Without it, the entire network would be a linear transformation, which could be collapsed into a single matrix multiplication. The non-linearity allows the network to learn complex, non-linear transformations of the input.
In modern Transformers, you'll sometimes see other activation functions like GELU (Gaussian Error Linear Unit), which is a smooth approximation of ReLU. The choice of activation function can affect training dynamics and final performance, but the basic principle remains the same: we need non-linearity to increase model capacity.
LAYER NORMALIZATION: STABILIZING TRAINING
As we stack multiple Transformer layers, we encounter a training stability problem. The outputs of each layer can have wildly different scales, which makes training difficult. Gradients can explode or vanish, and the model can become very sensitive to initialization.
Layer normalization solves this by normalizing the inputs to each layer to have mean zero and variance one. For a vector x, layer normalization computes:
LayerNorm(x) = gamma * ((x - mu) / sqrt(sigma^2 + epsilon)) + beta
Let's understand each component. First, we compute the mean mu and variance sigma^2 across the features:
mu = (1/d) * sum(x_i)
sigma^2 = (1/d) * sum((x_i - mu)^2)
We subtract the mean and divide by the standard deviation. This normalizes the distribution to have mean zero and variance one. The epsilon (typically 1e-6) is added for numerical stability to prevent division by zero.
But here's the clever part: we then apply a learned affine transformation with parameters gamma and beta. These are learned during training and allow the model to undo the normalization if needed. Why would we want to undo normalization? Because sometimes the optimal distribution isn't zero mean and unit variance. By making gamma and beta learnable, we give the model the flexibility to learn the best distribution for each layer.
The key difference between layer normalization and batch normalization (used in CNNs) is what we normalize over. Batch normalization normalizes across the batch dimension, which creates dependencies between examples in a batch. Layer normalization normalizes across the feature dimension for each example independently. This makes it more suitable for sequence models where sequence lengths vary.
Here's the implementation:
class LayerNormalization:
"""
Layer normalization.
Normalizes inputs across the feature dimension for each example.
"""
def __init__(self, d_model, epsilon=1e-6):
"""
Initialize layer normalization.
Args:
d_model: Dimension of the model
epsilon: Small constant for numerical stability
"""
self.epsilon = epsilon
# Learnable scale parameter (initialized to 1)
self.gamma = np.ones(d_model)
# Learnable shift parameter (initialized to 0)
self.beta = np.zeros(d_model)
def forward(self, x):
"""
Apply layer normalization.
Args:
x: Input tensor of shape (batch_size, seq_len, d_model)
Returns:
normalized: Normalized tensor of same shape as input
"""
# Compute mean and variance across the feature dimension
# keepdims=True preserves dimensions for broadcasting
mean = np.mean(x, axis=-1, keepdims=True)
variance = np.var(x, axis=-1, keepdims=True)
# Normalize to zero mean and unit variance
x_normalized = (x - mean) / np.sqrt(variance + self.epsilon)
# Apply learned affine transformation
# gamma and beta are broadcast across batch and sequence dimensions
output = self.gamma * x_normalized + self.beta
return output
Layer normalization is typically applied before or after each sub-layer in the Transformer. The original paper applied it after (post-norm), but modern implementations often apply it before (pre-norm) because this tends to stabilize training for very deep networks.
RESIDUAL CONNECTIONS: ENABLING DEEP NETWORKS
Even with layer normalization, training very deep networks is challenging. As we stack more layers, gradients must flow through all of them during backpropagation. Each layer transformation can distort the gradient, making it harder for early layers to learn.
Residual connections, introduced in ResNet for computer vision, provide an elegant solution. Instead of learning a transformation F(x), we learn a residual function F(x) that we add to the input:
output = x + F(x)
This simple addition has profound effects. First, it creates a direct path for gradients to flow backward. During backpropagation, the gradient of the addition operation is just one, so gradients can flow unchanged through the residual connection. This helps prevent vanishing gradients in deep networks.
Second, it makes optimization easier. Learning F(x) = 0 (the identity function) is much easier than learning an arbitrary transformation that happens to be close to identity. With residual connections, the network can start from the identity and learn small refinements, rather than having to learn the entire transformation from scratch.
In Transformers, we apply residual connections around both the attention and feed-forward sub-layers:
# After multi-head attention
x = x + MultiHeadAttention(x)
x = LayerNorm(x)
# After feed-forward network
x = x + FeedForwardNetwork(x)
x = LayerNorm(x)
The combination of residual connections and layer normalization is often called "Add & Norm" in the Transformer literature. Let's implement a complete Transformer layer:
class TransformerEncoderLayer:
"""
A single Transformer encoder layer.
Consists of multi-head attention and feed-forward network,
each with residual connections and layer normalization.
"""
def __init__(self, d_model, num_heads, d_ff, dropout_rate=0.1):
"""
Initialize Transformer encoder layer.
Args:
d_model: Dimension of the model
num_heads: Number of attention heads
d_ff: Dimension of feed-forward hidden layer
dropout_rate: Dropout probability for regularization
"""
self.attention = MultiHeadAttention(d_model, num_heads)
self.feed_forward = FeedForwardNetwork(d_model, d_ff)
self.norm1 = LayerNormalization(d_model)
self.norm2 = LayerNormalization(d_model)
self.dropout_rate = dropout_rate
def dropout(self, x, training=True):
"""
Apply dropout for regularization.
Randomly sets elements to zero during training.
"""
if not training:
return x
# Create a mask of random values
mask = np.random.binomial(1, 1 - self.dropout_rate, x.shape)
# Scale by 1/(1-p) to maintain expected value
return x * mask / (1 - self.dropout_rate)
def forward(self, x, training=True):
"""
Process input through the Transformer layer.
Args:
x: Input tensor of shape (batch_size, seq_len, d_model)
training: Whether we're in training mode (affects dropout)
Returns:
output: Transformed tensor of same shape as input
"""
# Multi-head attention with residual connection and normalization
# Pre-norm variant: normalize before the sub-layer
attention_output = self.attention.forward(self.norm1.forward(x))
attention_output = self.dropout(attention_output, training)
x = x + attention_output # Residual connection
# Feed-forward network with residual connection and normalization
ff_output = self.feed_forward.forward(self.norm2.forward(x))
ff_output = self.dropout(ff_output, training)
x = x + ff_output # Residual connection
return x
Notice that we've also added dropout, which randomly sets some activations to zero during training. This is a regularization technique that prevents overfitting by forcing the network to learn redundant representations.
THE COMPLETE TRANSFORMER ARCHITECTURE: PUTTING IT ALL TOGETHER
We now have all the building blocks. Let's assemble them into a complete Transformer. The original Transformer was designed for sequence-to-sequence tasks like translation, so it has an encoder and a decoder. We'll focus on the encoder, which is what's used in models like BERT, and briefly discuss the decoder.
The encoder consists of a stack of identical layers, each containing multi-head attention and a feed-forward network with residual connections and layer normalization. The input goes through an embedding layer, gets positional encodings added, and then flows through all the encoder layers.
Here's the complete encoder:
class TransformerEncoder:
"""
Complete Transformer encoder.
Stacks multiple encoder layers and handles input embedding.
"""
def __init__(self, vocab_size, d_model, num_layers, num_heads,
d_ff, max_seq_len, dropout_rate=0.1):
"""
Initialize Transformer encoder.
Args:
vocab_size: Size of the vocabulary
d_model: Dimension of the model
num_layers: Number of encoder layers to stack
num_heads: Number of attention heads in each layer
d_ff: Dimension of feed-forward hidden layer
max_seq_len: Maximum sequence length (for positional encoding)
dropout_rate: Dropout probability
"""
self.d_model = d_model
self.num_layers = num_layers
# Embedding layer converts token IDs to dense vectors
# We initialize with small random values
self.embedding = np.random.randn(vocab_size, d_model) * 0.01
# Positional encoding is fixed (not learned)
self.pos_encoding = get_positional_encoding(max_seq_len, d_model)
# Stack of encoder layers
self.layers = [
TransformerEncoderLayer(d_model, num_heads, d_ff, dropout_rate)
for _ in range(num_layers)
]
# Final layer normalization
self.final_norm = LayerNormalization(d_model)
self.dropout_rate = dropout_rate
def embed_tokens(self, token_ids):
"""
Convert token IDs to embeddings.
Args:
token_ids: Integer array of shape (batch_size, seq_len)
Returns:
embeddings: Dense vectors of shape (batch_size, seq_len, d_model)
"""
# Look up embeddings for each token
embeddings = self.embedding[token_ids]
# Scale embeddings by sqrt(d_model) as in the original paper
# This prevents the positional encoding from dominating
embeddings = embeddings * np.sqrt(self.d_model)
return embeddings
def add_positional_encoding(self, embeddings):
"""
Add positional encodings to embeddings.
Args:
embeddings: Token embeddings of shape (batch, seq_len, d_model)
Returns:
embeddings_with_pos: Embeddings with positional info added
"""
seq_len = embeddings.shape[1]
# Add positional encoding (broadcasting handles batch dimension)
return embeddings + self.pos_encoding[:seq_len, :]
def forward(self, token_ids, training=True):
"""
Process input tokens through the entire encoder.
Args:
token_ids: Input token IDs of shape (batch_size, seq_len)
training: Whether we're in training mode
Returns:
output: Encoded representations of shape (batch, seq_len, d_model)
"""
# Convert tokens to embeddings
x = self.embed_tokens(token_ids)
# Add positional information
x = self.add_positional_encoding(x)
# Apply dropout to embeddings
if training:
mask = np.random.binomial(1, 1 - self.dropout_rate, x.shape)
x = x * mask / (1 - self.dropout_rate)
# Pass through all encoder layers
for layer in self.layers:
x = layer.forward(x, training)
# Final normalization
x = self.final_norm.forward(x)
return x
The decoder is similar to the encoder but with one crucial addition: masked self-attention. In the decoder, when generating the output sequence, we can only attend to positions that have already been generated. This is enforced by adding a mask to the attention scores before the softmax:
# Create a mask that prevents attending to future positions
mask = np.triu(np.ones((seq_len, seq_len)) * -1e9, k=1)
# Add mask to attention scores before softmax
scores = scores + mask
The mask sets future positions to negative infinity, so after softmax they become zero. This ensures the model can only use past context when generating each token.
The decoder also has a second attention mechanism called cross-attention, where queries come from the decoder but keys and values come from the encoder output. This allows the decoder to attend to the input sequence while generating the output.
TRAINING THE TRANSFORMER: LOSS FUNCTIONS AND OPTIMIZATION
Now that we have a complete architecture, how do we train it? The training process depends on the task, but let's consider the most common case: language modeling, where we predict the next token given previous tokens.
For language modeling, we use cross-entropy loss. Given the model's predicted probability distribution over the vocabulary and the true next token, the loss is:
Loss = -log(P(true_token | context))
Why negative log probability? Because we want to maximize the probability of the correct token, which is equivalent to minimizing the negative log probability. The log converts products of probabilities (for a sequence) into sums, which are easier to work with numerically.
For a batch of examples, we average the loss:
Total_Loss = -(1/N) * sum(log(P(true_token_i | context_i)))
Here's how we compute this in practice:
def compute_language_modeling_loss(logits, target_tokens):
"""
Compute cross-entropy loss for language modeling.
Args:
logits: Model outputs of shape (batch_size, seq_len, vocab_size)
These are unnormalized scores for each token
target_tokens: True next tokens of shape (batch_size, seq_len)
Integer indices into vocabulary
Returns:
loss: Scalar loss value
"""
batch_size, seq_len, vocab_size = logits.shape
# Reshape for easier processing
logits_flat = logits.reshape(-1, vocab_size) # (batch*seq_len, vocab)
targets_flat = target_tokens.reshape(-1) # (batch*seq_len,)
# Compute softmax probabilities
# Subtract max for numerical stability
logits_shifted = logits_flat - np.max(logits_flat, axis=1, keepdims=True)
exp_logits = np.exp(logits_shifted)
probs = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)
# Get probability of correct token for each position
# We use advanced indexing to select the probability of the target token
batch_indices = np.arange(batch_size * seq_len)
correct_probs = probs[batch_indices, targets_flat]
# Compute negative log likelihood
# Add small epsilon to prevent log(0)
loss = -np.mean(np.log(correct_probs + 1e-10))
return loss
To optimize the model, we use gradient descent. The original Transformer paper used the Adam optimizer with a specific learning rate schedule. The learning rate increases linearly for the first warmup_steps, then decreases proportionally to the inverse square root of the step number:
learning_rate = (d_model^(-0.5)) * min(step^(-0.5), step * warmup_steps^(-1.5))
This schedule has a warmup phase where the learning rate gradually increases, which helps stabilize training in the early stages when the model's parameters are random. After warmup, the learning rate gradually decreases, allowing the model to fine-tune its parameters.
Why this specific schedule? The warmup prevents the model from making large updates based on the initial random gradients, which could push it into a bad region of parameter space. The decay allows the model to make smaller and smaller adjustments as it converges, preventing oscillation around the optimum.
Here's a simplified training loop:
def train_transformer(model, train_data, num_epochs, warmup_steps):
"""
Train a Transformer model.
Args:
model: TransformerEncoder instance
train_data: List of (input_tokens, target_tokens) pairs
num_epochs: Number of passes through the training data
warmup_steps: Number of steps for learning rate warmup
"""
step = 0
d_model = model.d_model
for epoch in range(num_epochs):
epoch_loss = 0
for input_tokens, target_tokens in train_data:
step += 1
# Compute learning rate with warmup
lr = (d_model ** -0.5) * min(
step ** -0.5,
step * (warmup_steps ** -1.5)
)
# Forward pass
output = model.forward(input_tokens, training=True)
# Compute loss
# We need to add a linear layer to project to vocabulary size
logits = np.matmul(output, model.embedding.T) # Weight sharing
loss = compute_language_modeling_loss(logits, target_tokens)
epoch_loss += loss
# Backward pass (compute gradients)
# In practice, this would use automatic differentiation
gradients = compute_gradients(loss, model)
# Update parameters using Adam optimizer
update_parameters_adam(model, gradients, lr, step)
if step % 100 == 0:
print(f"Step {step}, Loss: {loss:.4f}, LR: {lr:.6f}")
avg_loss = epoch_loss / len(train_data)
print(f"Epoch {epoch + 1}, Average Loss: {avg_loss:.4f}")
In practice, we'd use a framework like PyTorch or TensorFlow that handles gradient computation automatically. The key insight is that the entire Transformer is differentiable, so we can use backpropagation to compute gradients and update all parameters end-to-end.
ADVANCED TOPIC: MIXTURE OF EXPERTS
As we scale Transformers to billions or trillions of parameters, we face a computational challenge. Larger models are more capable, but they're also much slower and more expensive to run. Mixture of Experts (MoE) offers a solution: we can have a huge model but only use a small part of it for each input.
The core idea is to replace the feed-forward network in each Transformer layer with multiple expert networks and a gating mechanism that decides which experts to use for each token. Instead of:
output = FeedForward(x)
We have:
gate_scores = Softmax(x * W_gate)
output = sum(gate_scores[i] * Expert_i(x) for i in top_k(gate_scores))
Let's break down what's happening. First, we compute gate scores for each expert using a learned weight matrix W_gate. These scores tell us how relevant each expert is for the current input. We apply softmax to convert them to probabilities.
Then, instead of using all experts, we select only the top k experts with the highest gate scores. This is called sparse routing. For each selected expert, we compute its output and weight it by the gate score. The final output is the weighted sum of the selected experts' outputs.
Why does this work? Different experts can specialize in different types of inputs. For example, in a language model, one expert might specialize in technical text, another in conversational text, and another in formal writing. The gating network learns to route each input to the most appropriate experts.
The key benefit is computational efficiency. If we have 128 experts but only use the top 2 for each token, we do roughly the same computation as a regular Transformer with 2 feed-forward networks, but we have the capacity of 128 networks. This allows us to scale to much larger models.
Here's an implementation:
class MixtureOfExpertsLayer:
"""
Mixture of Experts layer.
Routes each token to a subset of expert networks based on learned gating.
"""
def __init__(self, d_model, d_ff, num_experts, top_k):
"""
Initialize MoE layer.
Args:
d_model: Model dimension
d_ff: Expert hidden dimension
num_experts: Total number of expert networks
top_k: Number of experts to use per token
"""
self.d_model = d_model
self.num_experts = num_experts
self.top_k = top_k
# Gating network that decides which experts to use
self.W_gate = np.random.randn(d_model, num_experts) * 0.01
# Create multiple expert networks
self.experts = [
FeedForwardNetwork(d_model, d_ff)
for _ in range(num_experts)
]
def forward(self, x):
"""
Route inputs through mixture of experts.
Args:
x: Input of shape (batch_size, seq_len, d_model)
Returns:
output: Processed input of same shape
"""
batch_size, seq_len, d_model = x.shape
# Flatten batch and sequence dimensions for processing
x_flat = x.reshape(-1, d_model) # (batch*seq_len, d_model)
# Compute gate scores for all experts
gate_logits = np.matmul(x_flat, self.W_gate) # (batch*seq_len, num_experts)
gate_scores = softmax(gate_logits)
# Select top-k experts for each token
# Get indices of top-k experts
top_k_indices = np.argsort(gate_scores, axis=1)[:, -self.top_k:]
# Get the gate scores for selected experts
batch_indices = np.arange(x_flat.shape[0])[:, None]
top_k_gates = gate_scores[batch_indices, top_k_indices]
# Normalize gate scores of selected experts to sum to 1
top_k_gates = top_k_gates / np.sum(top_k_gates, axis=1, keepdims=True)
# Compute output as weighted combination of expert outputs
output_flat = np.zeros_like(x_flat)
for i in range(self.top_k):
# Get the expert index for each token
expert_indices = top_k_indices[:, i]
# Process each token with its selected expert
for token_idx in range(x_flat.shape[0]):
expert_idx = expert_indices[token_idx]
expert_output = self.experts[expert_idx].forward(
x_flat[token_idx:token_idx+1]
)
gate_weight = top_k_gates[token_idx, i]
output_flat[token_idx] += gate_weight * expert_output[0]
# Reshape back to original dimensions
output = output_flat.reshape(batch_size, seq_len, d_model)
return output
There's a subtle challenge with training MoE models: load balancing. If the gating network always routes to the same few experts, most experts never get trained and we lose the benefits of having many experts. To address this, we add an auxiliary loss that encourages balanced routing:
load_balance_loss = num_experts * sum(f_i * P_i)
Here, f_i is the fraction of tokens routed to expert i, and P_i is the average gate probability for expert i. This loss is minimized when all experts are used equally. We add this to the main loss with a small weight to encourage balanced routing without overriding the main objective.
Modern MoE models like GPT-4 and Gemini use this architecture to achieve massive scale while keeping inference costs manageable. The key insight is that not every parameter needs to be active for every input, allowing us to build much larger models than would otherwise be feasible.
ADVANCED TOPIC: REASONING MODELS AND CHAIN-OF-THOUGHT
Recent advances in language models have focused on improving their reasoning capabilities. The key insight is that complex reasoning often requires intermediate steps. When humans solve a difficult problem, we don't jump directly to the answer; we work through it step by step.
Chain-of-Thought (CoT) prompting encourages models to generate these intermediate steps. Instead of asking "What is 47 times 23?" and expecting the answer directly, we prompt the model to show its work:
Question: What is 47 times 23?
Let's solve this step by step:
47 * 23 = 47 * 20 + 47 * 3
= 940 + 141
= 1081
This simple change dramatically improves performance on reasoning tasks. But why does it work mathematically? The key is that Transformers have limited computational depth per token. Each layer can only perform a fixed amount of computation. By generating intermediate steps, we're effectively giving the model more computational steps to solve the problem.
Think of it this way: if the model has 24 layers and generates 10 tokens of reasoning, it gets 24 times 10 equals 240 layers worth of computation, versus just 24 layers if it tries to answer directly. This is why longer chains of thought often lead to better answers.
Recent models like OpenAI's o1 take this further with learned reasoning. Instead of relying on prompting, they're trained to generate reasoning steps automatically. The training process involves:
First, we generate many possible reasoning chains for each problem using techniques like tree search or sampling. Some chains lead to correct answers, others to incorrect ones.
Second, we use reinforcement learning to train the model to prefer reasoning chains that lead to correct answers. The reward signal is binary: did the final answer match the ground truth?
The key mathematical framework is policy gradient reinforcement learning. We treat the model as a policy that generates a sequence of reasoning tokens. The objective is to maximize expected reward:
J(theta) = E[R(y) | x, theta]
Here, theta represents the model parameters, x is the input problem, y is the generated reasoning chain and answer, and R(y) is the reward (1 for correct, 0 for incorrect). We compute the gradient:
gradient J(theta) = E[R(y) * gradient log P(y | x, theta)]
This tells us how to adjust the parameters to increase the probability of high-reward reasoning chains. In practice, we estimate this expectation by sampling multiple reasoning chains and computing their rewards.
The challenge is that most reasoning chains are incorrect, so the reward signal is sparse. To address this, we use several techniques:
We use value functions to estimate the expected reward of partial reasoning chains, allowing us to guide the search toward promising directions.
We use curriculum learning, starting with easier problems and gradually increasing difficulty.
We use process rewards, giving credit for correct intermediate steps even if the final answer is wrong. This provides denser feedback.
Here's a simplified implementation of the training loop:
def train_reasoning_model(model, problems, num_iterations):
"""
Train a model to generate reasoning chains using reinforcement learning.
Args:
model: Transformer model
problems: List of (problem, correct_answer) pairs
num_iterations: Number of training iterations
"""
for iteration in range(num_iterations):
total_reward = 0
for problem, correct_answer in problems:
# Sample multiple reasoning chains
num_samples = 4
chains = []
rewards = []
for _ in range(num_samples):
# Generate a reasoning chain
reasoning_chain = model.generate(
problem,
max_length=200,
temperature=0.8 # Some randomness for exploration
)
# Extract the final answer from the chain
final_answer = extract_answer(reasoning_chain)
# Compute reward
reward = 1.0 if final_answer == correct_answer else 0.0
chains.append(reasoning_chain)
rewards.append(reward)
total_reward += reward
# Compute baseline (average reward) for variance reduction
baseline = np.mean(rewards)
# Update model to increase probability of high-reward chains
for chain, reward in zip(chains, rewards):
# Compute advantage (reward minus baseline)
advantage = reward - baseline
# Compute log probability of the chain
log_prob = model.compute_log_probability(problem, chain)
# Policy gradient: increase log prob proportional to advantage
loss = -advantage * log_prob # Compute gradients and update parameters
gradients = compute_gradients(loss, model)
update_parameters(model, gradients, learning_rate=1e-5)
avg_reward = total_reward / (len(problems) * num_samples)
print(f"Iteration {iteration}, Average Reward: {avg_reward:.3f}")
The key insight is that we're not just training the model to predict the next token based on supervised data. We're training it to explore different reasoning strategies and reinforce the ones that lead to correct answers. This allows the model to discover reasoning patterns that might not be present in the training data.
One important consideration is that longer reasoning chains cost more at inference time. Each generated token requires a forward pass through the entire model. This creates a trade-off: longer reasoning improves accuracy but increases latency and cost. In practice, we might use adaptive computation, where the model learns to generate longer chains only for difficult problems.
ADVANCED TOPIC: CREATING HIGH-QUALITY TRAINING DATA
The quality of a Transformer model is fundamentally limited by the quality of its training data. As the saying goes in machine learning: garbage in, garbage out. Creating high-quality training data at scale is one of the most important and challenging aspects of building modern AI systems.
For language models, we typically start with web-scale text data. But raw web text has many problems. It contains factual errors, toxic content, spam, duplicate content, and low-quality writing. The data creation pipeline involves several stages of filtering and processing.
The first stage is deduplication. Web crawls contain massive amounts of duplicate or near-duplicate content. Training on duplicates wastes computation and can cause the model to memorize specific examples rather than learning general patterns. We use techniques like MinHash to efficiently find near-duplicates:
def compute_minhash_signature(text, num_hashes=128):
"""
Compute MinHash signature for approximate duplicate detection.
Args:
text: Input text string
num_hashes: Number of hash functions to use
Returns:
signature: Array of minimum hash values
"""
# Tokenize text into shingles (overlapping n-grams)
# We use character-level 5-grams for robustness
shingles = set()
for i in range(len(text) - 4):
shingle = text[i:i+5]
shingles.add(shingle)
# Initialize signature with infinity
signature = np.full(num_hashes, np.inf)
# For each shingle, compute multiple hash values
for shingle in shingles:
# Convert shingle to bytes for hashing
shingle_bytes = shingle.encode('utf-8')
for i in range(num_hashes):
# Use different hash functions by adding salt
hash_value = hash((shingle_bytes, i))
# Keep minimum hash value for each hash function
signature[i] = min(signature[i], hash_value)
return signature
def estimate_jaccard_similarity(sig1, sig2):
"""
Estimate Jaccard similarity from MinHash signatures.
The fraction of matching hash values approximates the Jaccard
similarity of the original shingle sets.
"""
matches = np.sum(sig1 == sig2)
return matches / len(sig1)
The mathematical foundation of MinHash is elegant. The Jaccard similarity between two sets A and B is defined as:
J(A, B) = |A intersection B| / |A union B|
MinHash provides an unbiased estimate of this similarity. The probability that the minimum hash value is the same for two sets equals their Jaccard similarity. By using multiple hash functions, we can estimate the similarity with high accuracy.
After deduplication, we apply quality filtering. This is where things get interesting from a machine learning perspective. We can't manually review billions of documents, so we need automated quality assessment. One approach is to train a classifier to predict quality:
def train_quality_classifier(high_quality_examples, low_quality_examples):
"""
Train a classifier to distinguish high-quality from low-quality text.
Args:
high_quality_examples: List of high-quality documents
low_quality_examples: List of low-quality documents
Returns:
classifier: Trained model that predicts quality scores
"""
# Extract features from text
def extract_features(text):
features = {}
# Length features
features['num_words'] = len(text.split())
features['num_sentences'] = text.count('.') + text.count('!') + text.count('?')
features['avg_word_length'] = np.mean([len(w) for w in text.split()])
# Vocabulary diversity (type-token ratio)
words = text.lower().split()
features['vocabulary_diversity'] = len(set(words)) / max(len(words), 1)
# Punctuation and formatting
features['punctuation_ratio'] = sum(c in '.,!?;:' for c in text) / max(len(text), 1)
features['capitalization_ratio'] = sum(c.isupper() for c in text) / max(len(text), 1)
# Readability (simplified Flesch-Kincaid)
avg_sentence_length = features['num_words'] / max(features['num_sentences'], 1)
features['readability'] = avg_sentence_length * features['avg_word_length']
return np.array(list(features.values()))
# Prepare training data
X_train = []
y_train = []
for text in high_quality_examples:
X_train.append(extract_features(text))
y_train.append(1) # High quality
for text in low_quality_examples:
X_train.append(extract_features(text))
y_train.append(0) # Low quality
X_train = np.array(X_train)
y_train = np.array(y_train)
# Train a simple logistic regression classifier
# In practice, we might use a more sophisticated model
classifier = train_logistic_regression(X_train, y_train)
return classifier
The challenge is obtaining labeled examples of high and low quality text. One approach is to use proxy signals. For example, we might consider text from reputable sources like Wikipedia or academic papers as high quality, and text from spam sites as low quality. We can also use engagement signals: documents that users spend more time reading might be higher quality.
Another crucial aspect of data creation is diversity. If all our training data comes from similar sources, the model will have blind spots. We want data covering many topics, writing styles, and perspectives. We can measure diversity using topic modeling:
def measure_topic_diversity(documents, num_topics=100):
"""
Measure the diversity of topics in a document collection.
Uses Latent Dirichlet Allocation (LDA) to discover topics
and measures how evenly documents are distributed across topics.
"""
# Create vocabulary and document-term matrix
vocabulary = build_vocabulary(documents)
doc_term_matrix = create_doc_term_matrix(documents, vocabulary)
# Run LDA to discover topics
# Each topic is a distribution over words
# Each document is a distribution over topics
topic_distributions = run_lda(doc_term_matrix, num_topics)
# Compute entropy of topic distribution
# High entropy means documents are spread across many topics
topic_counts = np.sum(topic_distributions > 0.1, axis=0)
topic_probs = topic_counts / np.sum(topic_counts)
# Shannon entropy measures diversity
entropy = -np.sum(topic_probs * np.log(topic_probs + 1e-10))
max_entropy = np.log(num_topics)
# Normalize to 0-1 range
diversity_score = entropy / max_entropy
return diversity_score
For instruction-tuning data, which teaches models to follow instructions, the creation process is even more sophisticated. We need examples of instructions paired with high-quality responses. There are several approaches:
Human annotation is the gold standard but expensive. We hire human annotators to write instructions and responses, or to rate model-generated responses. The challenge is ensuring consistency across annotators. We use detailed guidelines and regular calibration sessions.
Synthetic data generation uses existing models to create training data for new models. This seems circular, but it works when done carefully. For example, we might use a large model to generate diverse instructions, then have humans write responses, or vice versa. The key is that human involvement ensures quality at critical points.
Distillation involves using a large, powerful model to create training data for a smaller model. We generate many responses from the large model, filter for quality, and use them to train the smaller model. The mathematical framework is knowledge distillation:
def distillation_loss(student_logits, teacher_logits, temperature):
"""
Compute distillation loss that transfers knowledge from teacher to student.
Args:
student_logits: Unnormalized predictions from student model
teacher_logits: Unnormalized predictions from teacher model
temperature: Softmax temperature for smoothing distributions
Returns:
loss: Distillation loss value
"""
# Apply temperature scaling to soften the distributions
# Higher temperature makes distributions more uniform
student_probs = softmax(student_logits / temperature)
teacher_probs = softmax(teacher_logits / temperature)
# Compute KL divergence between distributions
# This measures how different the student's predictions are from the teacher's
kl_div = np.sum(teacher_probs * np.log(teacher_probs / (student_probs + 1e-10)))
# Scale by temperature squared (from the derivative of the softmax)
loss = kl_div * (temperature ** 2)
return loss
The temperature parameter is crucial. With temperature equals one, we get the normal softmax. With higher temperatures, the distribution becomes more uniform, revealing more information about the teacher's uncertainty. The student learns not just the most likely answer but the teacher's full probability distribution over answers.
For reinforcement learning from human feedback (RLHF), which is used to align models with human preferences, we need preference data. Humans compare multiple model outputs and indicate which is better. This creates a dataset of comparisons rather than absolute labels.
We train a reward model to predict human preferences:
def train_reward_model(comparisons):
"""
Train a model to predict which response humans prefer.
Args:
comparisons: List of (prompt, response_a, response_b, preference) tuples
where preference is 0 if A is preferred, 1 if B is preferred
"""
# The reward model takes a prompt and response and outputs a scalar score
reward_model = create_reward_model()
for prompt, response_a, response_b, preference in comparisons:
# Get reward scores for both responses
score_a = reward_model.forward(prompt, response_a)
score_b = reward_model.forward(prompt, response_b)
# Compute probability that A is preferred using Bradley-Terry model
# This assumes preferences follow a logistic distribution
prob_a_preferred = 1 / (1 + np.exp(score_b - score_a))
# Compute loss based on actual preference
if preference == 0: # A is preferred
loss = -np.log(prob_a_preferred + 1e-10)
else: # B is preferred
loss = -np.log(1 - prob_a_preferred + 1e-10)
# Update reward model parameters
gradients = compute_gradients(loss, reward_model)
update_parameters(reward_model, gradients)
return reward_model
The Bradley-Terry model is a classical model from psychometrics that assumes if item A has quality score r_A and item B has quality score r_B, the probability that A is preferred is:
P(A preferred over B) = exp(r_A) / (exp(r_A) + exp(r_B))
= 1 / (1 + exp(r_B - r_A))
This gives us a principled way to convert pairwise comparisons into scalar reward scores that we can use for reinforcement learning.
The final stage of data creation is often the most important: human review and iteration. We sample from our filtered dataset, manually review it, identify remaining problems, and update our filtering criteria. This iterative process gradually improves data quality. The key insight is that data creation is not a one-time process but an ongoing effort that directly impacts model quality.
PUTTING IT ALL TOGETHER: IMPLEMENTING A COMPLETE TRANSFORMER
Now that we understand all the components, let's implement a complete, working Transformer that you could actually train. We'll create a simplified but functional version that demonstrates all the key concepts:
class Transformer:
"""
Complete Transformer model for language modeling.
Combines all components: embeddings, positional encoding,
encoder layers, and output projection.
"""
def __init__(self, vocab_size, d_model=512, num_layers=6,
num_heads=8, d_ff=2048, max_seq_len=512,
dropout_rate=0.1):
"""
Initialize complete Transformer.
Args:
vocab_size: Size of vocabulary
d_model: Model dimension (must be divisible by num_heads)
num_layers: Number of Transformer layers
num_heads: Number of attention heads per layer
d_ff: Feed-forward hidden dimension
max_seq_len: Maximum sequence length
dropout_rate: Dropout probability
"""
self.vocab_size = vocab_size
self.d_model = d_model
self.max_seq_len = max_seq_len
# Token embedding layer
# Maps token IDs to dense vectors
self.token_embedding = np.random.randn(vocab_size, d_model) * 0.02
# Positional encoding (fixed, not learned)
self.positional_encoding = get_positional_encoding(max_seq_len, d_model)
# Stack of Transformer encoder layers
self.encoder_layers = [
TransformerEncoderLayer(d_model, num_heads, d_ff, dropout_rate)
for _ in range(num_layers)
]
# Final layer normalization
self.output_norm = LayerNormalization(d_model)
# Output projection to vocabulary
# We use weight tying: share weights with embedding layer
# This reduces parameters and often improves performance
self.output_projection = self.token_embedding.T
self.dropout_rate = dropout_rate
def encode(self, token_ids, training=True):
"""
Encode input tokens to contextualized representations.
Args:
token_ids: Integer array of shape (batch_size, seq_len)
training: Whether in training mode (affects dropout)
Returns:
encoded: Representations of shape (batch_size, seq_len, d_model)
"""
batch_size, seq_len = token_ids.shape
# Embed tokens
# Look up embedding for each token ID
embedded = self.token_embedding[token_ids]
# Scale embeddings by sqrt(d_model)
# This prevents positional encodings from dominating
embedded = embedded * np.sqrt(self.d_model)
# Add positional encodings
# Broadcasting handles batch dimension
pos_enc = self.positional_encoding[:seq_len, :]
x = embedded + pos_enc
# Apply dropout to embeddings
if training:
mask = np.random.binomial(1, 1 - self.dropout_rate, x.shape)
x = x * mask / (1 - self.dropout_rate)
# Pass through encoder layers
for layer in self.encoder_layers:
x = layer.forward(x, training)
# Final normalization
x = self.output_norm.forward(x)
return x
def forward(self, token_ids, training=True):
"""
Full forward pass: encode and project to vocabulary.
Args:
token_ids: Input token IDs of shape (batch_size, seq_len)
training: Whether in training mode
Returns:
logits: Unnormalized scores over vocabulary
Shape: (batch_size, seq_len, vocab_size)
"""
# Encode input
encoded = self.encode(token_ids, training)
# Project to vocabulary
# Shape: (batch, seq_len, d_model) @ (d_model, vocab)
# = (batch, seq_len, vocab)
logits = np.matmul(encoded, self.output_projection)
return logits
def generate(self, prompt_tokens, max_new_tokens=50, temperature=1.0):
"""
Generate text autoregressively given a prompt.
Args:
prompt_tokens: Initial token IDs of shape (1, prompt_len)
max_new_tokens: Maximum number of tokens to generate
temperature: Sampling temperature (higher = more random)
Returns:
generated_tokens: Complete sequence including prompt
"""
# Start with the prompt
current_tokens = prompt_tokens.copy()
for _ in range(max_new_tokens):
# Get predictions for next token
# We only use the last position's logits
logits = self.forward(current_tokens, training=False)
next_token_logits = logits[0, -1, :] # Last position
# Apply temperature scaling
# Higher temperature makes distribution more uniform
scaled_logits = next_token_logits / temperature
# Convert to probabilities
probs = softmax(scaled_logits.reshape(1, -1))[0]
# Sample next token
next_token = np.random.choice(self.vocab_size, p=probs)
# Append to sequence
current_tokens = np.concatenate([
current_tokens,
np.array([[next_token]])
], axis=1)
# Stop if we exceed max length
if current_tokens.shape[1] >= self.max_seq_len:
break
return current_tokens
def compute_loss(self, input_tokens, target_tokens):
"""
Compute cross-entropy loss for language modeling.
Args:
input_tokens: Input sequence of shape (batch_size, seq_len)
target_tokens: Target sequence (shifted by 1) of same shape
Returns:
loss: Scalar loss value
"""
# Forward pass
logits = self.forward(input_tokens, training=True)
# Compute cross-entropy loss
batch_size, seq_len, vocab_size = logits.shape
# Reshape for easier computation
logits_flat = logits.reshape(-1, vocab_size)
targets_flat = target_tokens.reshape(-1)
# Compute softmax probabilities
logits_shifted = logits_flat - np.max(logits_flat, axis=1, keepdims=True)
exp_logits = np.exp(logits_shifted)
probs = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)
# Get probability of correct token
batch_indices = np.arange(batch_size * seq_len)
correct_probs = probs[batch_indices, targets_flat]
# Compute negative log likelihood
loss = -np.mean(np.log(correct_probs + 1e-10))
return loss
This implementation includes all the key components we've discussed: token embeddings, positional encodings, multi-head attention, feed-forward networks, layer normalization, and residual connections. You could train this model on text data using the training loop we outlined earlier.
To use this model in practice, you would:
First, prepare your data by tokenizing text into integer IDs. You would use a tokenizer like Byte-Pair Encoding (BPE) that breaks text into subword units, balancing vocabulary size with the ability to represent any text.
Second, create training batches where each example is a sequence of tokens. For language modeling, the input is tokens zero through n minus one, and the target is tokens one through n (shifted by one position).
Third, run the training loop, computing loss and gradients for each batch, and updating parameters using an optimizer like Adam.
Fourth, evaluate on a held-out validation set to monitor for overfitting. If validation loss stops improving while training loss continues to decrease, you're overfitting and should stop training or increase regularization.
Fifth, use the trained model for generation by providing a prompt and sampling tokens autoregressively until you reach a stopping condition.
OPTIMIZATION TECHNIQUES FOR PRODUCTION SYSTEMS
When deploying Transformers in production, we need to optimize for speed and memory efficiency. Several techniques are commonly used:
Flash Attention is an algorithm that computes attention more efficiently by being aware of GPU memory hierarchy. Standard attention computes the full attention matrix, which requires quadratic memory in sequence length. Flash Attention computes attention in blocks, keeping intermediate results in fast SRAM rather than slow HBM memory.
The key insight is that we don't need to materialize the full attention matrix. We can compute the output in chunks:
def flash_attention(Q, K, V, block_size=64):
"""
Memory-efficient attention using block-wise computation.
Args:
Q, K, V: Query, key, value matrices
block_size: Size of blocks for chunked computation
Returns:
output: Attention output
"""
seq_len, d = Q.shape
num_blocks = (seq_len + block_size - 1) // block_size
# Initialize output and normalization statistics
output = np.zeros_like(Q)
row_max = np.full(seq_len, -np.inf)
row_sum = np.zeros(seq_len)
# Process in blocks
for i in range(num_blocks):
# Get query block
q_start = i * block_size
q_end = min((i + 1) * block_size, seq_len)
Q_block = Q[q_start:q_end]
for j in range(num_blocks):
# Get key-value block
k_start = j * block_size
k_end = min((j + 1) * block_size, seq_len)
K_block = K[k_start:k_end]
V_block = V[k_start:k_end]
# Compute attention scores for this block
scores = np.matmul(Q_block, K_block.T) / np.sqrt(d)
# Update running max for numerical stability
block_max = np.max(scores, axis=1)
new_max = np.maximum(row_max[q_start:q_end], block_max)
# Compute exponentials with updated max
exp_scores = np.exp(scores - new_max[:, None])
# Update running sum
correction = np.exp(row_max[q_start:q_end] - new_max)
row_sum[q_start:q_end] = row_sum[q_start:q_end] * correction + np.sum(exp_scores, axis=1)
# Update output
output[q_start:q_end] = output[q_start:q_end] * correction[:, None] + np.matmul(exp_scores, V_block)
# Update max
row_max[q_start:q_end] = new_max
# Normalize output
output[q_start:q_end] = output[q_start:q_end] / row_sum[q_start:q_end, None]
return output
This block-wise computation reduces memory usage from O(n squared) to O(n), making it possible to process much longer sequences.
Quantization reduces the precision of weights and activations to save memory and increase speed. Instead of using 32-bit floating point numbers, we might use 8-bit integers. The challenge is maintaining accuracy with reduced precision.
The basic idea is to map floating point values to integers using a scale and zero point:
quantized_value = round(float_value / scale) + zero_point
To dequantize:
float_value = (quantized_value - zero_point) * scale
The scale and zero point are chosen to minimize quantization error. For a range of values from min_val to max_val:
scale = (max_val - min_val) / (2^bits - 1)
zero_point = round(-min_val / scale)
Here's an implementation:
def quantize_tensor(tensor, num_bits=8):
"""
Quantize floating point tensor to lower precision.
Args:
tensor: Float array to quantize
num_bits: Number of bits for quantized representation
Returns:
quantized: Integer array
scale: Scaling factor for dequantization
zero_point: Zero point for dequantization
"""
# Compute range of values
min_val = np.min(tensor)
max_val = np.max(tensor)
# Compute scale and zero point
qmin = 0
qmax = 2 ** num_bits - 1
scale = (max_val - min_val) / (qmax - qmin)
# Handle edge case where all values are the same
if scale == 0:
scale = 1.0
zero_point = qmin - round(min_val / scale)
zero_point = np.clip(zero_point, qmin, qmax)
# Quantize
quantized = np.round(tensor / scale + zero_point)
quantized = np.clip(quantized, qmin, qmax).astype(np.int8)
return quantized, scale, zero_point
def dequantize_tensor(quantized, scale, zero_point):
"""
Convert quantized tensor back to floating point.
"""
return (quantized.astype(np.float32) - zero_point) * scale
Modern quantization techniques like GPTQ and AWQ use more sophisticated methods that minimize the impact on model quality. They might use different scales for different parts of the network or calibrate on actual data to find optimal quantization parameters.
Kernel fusion combines multiple operations into a single GPU kernel to reduce memory bandwidth. For example, instead of computing layer normalization as separate operations (subtract mean, divide by std, multiply by gamma, add beta), we can fuse them into a single kernel that reads the input once and writes the output once.
The mathematical operations are the same, but the memory access pattern is much more efficient. This is especially important on modern GPUs where memory bandwidth is often the bottleneck.
CONCLUSION: THE TRANSFORMER REVOLUTION CONTINUES
We've journeyed from the fundamental problems that motivated Transformers through the mathematical details of every component to advanced topics like mixture of experts and reasoning models. The key insights that make Transformers work are:
Attention provides a differentiable mechanism for looking up relevant information anywhere in a sequence, solving the long-range dependency problem that plagued RNNs.
Scaling the dot product by the square root of dimension prevents saturation of the softmax function, enabling stable training.
Multi-head attention allows the model to learn multiple types of relationships simultaneously, dramatically increasing capacity without increasing computational cost proportionally.
Positional encodings inject information about token position into an otherwise position-agnostic architecture, using sinusoidal functions that generalize to any sequence length.
Residual connections and layer normalization enable training of very deep networks by providing gradient flow and stabilizing activations.
The feed-forward networks add computational depth and non-linearity, allowing the model to compute complex functions of its inputs.
These components combine into an architecture that is both powerful and efficient. The Transformer's ability to process sequences in parallel makes it dramatically faster to train than RNNs, while its attention mechanism gives it the capacity to capture long-range dependencies.
The mathematics behind Transformers is not arbitrary. Every formula, every architectural choice, emerged from solving concrete problems. The scaled dot-product attention formula exists because we needed a differentiable lookup mechanism that doesn't saturate. Multi-head attention exists because we need to capture multiple types of relationships. Positional encoding exists because attention is position-agnostic.
Understanding these mathematical foundations empowers you as a software engineer to not just use Transformers but to modify and improve them. You can now reason about why certain design choices were made and how you might adapt the architecture for your specific needs.
The field continues to evolve rapidly. New optimizations like Flash Attention make models faster and more memory-efficient. New training techniques like RLHF align models with human preferences. New architectures like mixture of experts scale to trillions of parameters. But all of these build on the fundamental Transformer architecture we've explored.
As you implement and experiment with Transformers, remember that the mathematics serves a purpose. Each formula solves a specific problem. By understanding the causal chain from problem to solution, you gain the insight needed to push the boundaries of what's possible with these remarkable models.
ADDENDUM - THE LIBRARY ANALOGY: UNDERSTANDING QUERIES, KEYS, AND VALUES
Imagine you're in a library searching for information. This is the perfect analogy for understanding how attention works with Queries, Keys, and Values.
You walk into the library with a specific question in mind: "I need information about climate change." This question is your Query. It represents what you're looking for, what information you need right now.
Each book in the library has an index card describing its contents: "This book covers: environmental science, global warming, carbon emissions." These descriptions are the Keys. They represent what each book has to offer, how it can be discovered or matched against your needs.
The actual content inside each book - the chapters, paragraphs, facts, and knowledge - that's the Value. This is what you actually want to retrieve and read once you've found a relevant book.
Here's the crucial insight: the Key (the index card description) might be different from the Value (the actual content). The index card might say "environmental science" but the book contains detailed climate models, historical temperature data, and policy recommendations. You use the Key to decide if the book is relevant, but you retrieve the Value to get the actual information.
When you search the library, you:
First, compare your Query against all the Keys to find which books are most relevant to your question. This comparison gives you relevance scores.
Second, you don't just pick the single most relevant book. Instead, you might take information from multiple books, weighted by how relevant each one is. If a book on climate science is highly relevant, you take a lot from it. If a book on general environmental topics is somewhat relevant, you take a little from it.
Third, you combine all this information (the Values) according to the relevance weights to form your answer.
This is exactly what attention does mathematically.
A CONCRETE SENTENCE EXAMPLE
Let's work through a real example with an actual sentence to see Q, K, and V in action. Consider this sentence:
"The animal didn't cross the street because it was too tired."
We want to figure out what "it" refers to. Let's focus on computing the representation for the word "it" using attention.
First, we need to understand that each word starts with an embedding - a vector of numbers that captures its meaning. Let's say each word is represented by a simple two-dimensional vector for illustration:
animal: [0.8, 0.1] # High on "living thing" dimension
street: [0.2, 0.9] # High on "location" dimension
it: [0.5, 0.5] # Neutral, needs context
tired: [0.9, 0.2] # High on "living thing" dimension
Now, when we process the word "it," we need to look back at the other words to understand what "it" refers to. This is where Q, K, and V come in.
The Query for "it" represents: "What am I looking for to understand this pronoun?" We create the Query by multiplying the "it" embedding by a learned weight matrix W_Q:
Q_it = embedding_it * W_Q
Let's say W_Q has learned to extract features that help identify what pronouns refer to. After this multiplication, we might get:
Q_it = [0.7, 0.3] # "I'm looking for a living thing that could be tired"
Now we create Keys for all words in the sentence. The Key represents: "What kind of information do I offer?" We multiply each word's embedding by W_K:
K_animal = embedding_animal * W_K = [0.9, 0.2] # "I offer info about a living creature"
K_street = embedding_street * W_K = [0.1, 0.8] # "I offer info about a location"
K_it = embedding_it * W_K = [0.5, 0.5] # "I offer ambiguous info"
K_tired = embedding_tired * W_K = [0.8, 0.3] # "I offer info about a state of being"
Now comes the matching step. We compute how well the Query matches each Key using dot products:
score_animal = dot(Q_it, K_animal) = 0.7*0.9 + 0.3*0.2 = 0.63 + 0.06 = 0.69
score_street = dot(Q_it, K_street) = 0.7*0.1 + 0.3*0.8 = 0.07 + 0.24 = 0.31
score_it = dot(Q_it, K_it) = 0.7*0.5 + 0.3*0.5 = 0.35 + 0.15 = 0.50
score_tired = dot(Q_it, K_tired) = 0.7*0.8 + 0.3*0.3 = 0.56 + 0.09 = 0.65
These scores tell us: "animal" is highly relevant (0.69), "tired" is also quite relevant (0.65), "street" is less relevant (0.31), and "it" itself is moderately relevant (0.50).
We apply softmax to convert these scores into weights that sum to one:
weight_animal = 0.35 # Highest weight
weight_street = 0.15 # Lowest weight
weight_it = 0.20 # Medium weight
weight_tired = 0.30 # High weight
Now we create Values. The Value represents: "What actual information do I contain?" We multiply each word's embedding by W_V:
V_animal = embedding_animal * W_V = [0.85, 0.15] # Rich semantic info about the animal
V_street = embedding_street * W_V = [0.25, 0.88] # Rich semantic info about the street
V_it = embedding_it * W_V = [0.48, 0.52] # Ambiguous semantic info
V_tired = embedding_tired * W_V = [0.82, 0.25] # Info about the tired state
Finally, we compute the attention output as a weighted sum of the Values:
output_it = 0.35 * V_animal + 0.15 * V_street + 0.20 * V_it + 0.30 * V_tired
= 0.35 * [0.85, 0.15] + 0.15 * [0.25, 0.88] + 0.20 * [0.48, 0.52] + 0.30 * [0.82, 0.25]
= [0.298, 0.053] + [0.038, 0.132] + [0.096, 0.104] + [0.246, 0.075]
= [0.678, 0.364]
Notice that the output for "it" now looks much more like "animal" and "tired" than like "street." The attention mechanism has successfully pulled information from the relevant words and created a new representation for "it" that incorporates the understanding that it refers to the animal.
WHY SEPARATE Q, K, AND V?
You might wonder: why do we need three different transformations? Why not just use the word embeddings directly?
The separation serves a crucial purpose. Let me illustrate with another example:
Consider the sentence: "The chef who made the pasta was Italian."
When we're processing "was," we need to figure out what the subject is. The Query from "was" might be asking: "Who is performing this state of being?"
The word "chef" has a Key that says: "I am a noun, I can be a subject of a sentence."
But the word "pasta" also has a Key that says: "I am also a noun, I could be a subject."
Both "chef" and "pasta" are nouns, so their Keys might be similar when we're looking for subjects. But their Values are very different:
The Value of "chef" contains: semantic information about a person, a profession, someone who cooks.
The Value of "pasta" contains: semantic information about food, something that is cooked.
The attention mechanism uses the Keys to determine that both "chef" and "pasta" are grammatically relevant (they're both nouns), but it might give "chef" a higher weight because it's the main noun of the subject phrase. Then it retrieves the Values, which contain the actual semantic content.
This separation allows the model to make decisions about relevance (using Q and K) that are independent from the information being retrieved (V). The Keys can be specialized for matching and routing, while the Values can be specialized for carrying semantic content.
A DATABASE ANALOGY
Another way to think about this is through database operations:
Imagine a database of employees with records like:
Employee 1: Index: "Software Engineer, Python, 5 years experience" (This is the Key) Full Record: {name: "Alice", skills: ["Python", "ML", "Docker"], projects: [...]} (This is the Value)Employee 2: Index: "Product Manager, 3 years experience" (This is the Key) Full Record: {name: "Bob", skills: ["Strategy", "Communication"], projects: [...]} (This is the Value)
But what you actually retrieve are the Values (the full records) with all the detailed information about each employee. The Key helped you find the right records, but the Value is what you actually use.
THE MATHEMATICAL REASON FOR SEPARATION
From a mathematical perspective, having separate Q, K, and V matrices gives the model more flexibility. If we used the same transformation for all three, we'd be computing:
attention = softmax(X * W * W^T * X^T) * X * W
attention = softmax(X * W_Q * W_K^T * X^T) * X * W_V
A TRANSLATION EXAMPLE
Let's consider a translation task to see another perspective. Suppose we're translating "The cat sat" to French.
When translating "sat," our Query might be: "I need to know what verb form to use, which requires knowing the subject's number and gender."
The Key for "cat" says: "I am a singular noun, potentially a subject."
The Value for "cat" contains: "This is 'cat' in French, it's masculine singular."
The attention mechanism uses the Query-Key matching to determine that "cat" is relevant to translating "sat" (because it's the subject). Then it retrieves the Value which contains the actual information needed: that the subject is masculine singular, so we should use "s'est assis" rather than "s'est assise" or "se sont assis."
CODE EXAMPLE SHOWING THE DIFFERENCE
Let's write code that makes the distinction crystal clear:
import numpy as np
def demonstrate_qkv_difference():
"""
Show concretely how Q, K, V serve different purposes.
"""
# Sentence: "The cat chased the mouse"
# We'll compute attention for the word "chased"
# Word embeddings (simplified to 4 dimensions)
embeddings = {
'the_1': np.array([0.1, 0.2, 0.3, 0.4]),
'cat': np.array([0.9, 0.1, 0.8, 0.2]), # Animal features
'chased': np.array([0.5, 0.5, 0.5, 0.5]),
'the_2': np.array([0.1, 0.2, 0.3, 0.4]),
'mouse': np.array([0.8, 0.2, 0.7, 0.3]) # Animal features
}
# Different weight matrices for different purposes
# W_Q: Learns to extract "what I'm looking for"
# For a verb, this might extract features about needing a subject
W_Q = np.array([
[1.0, 0.0, 0.5, 0.0], # Emphasize subject-related features
[0.0, 0.0, 0.0, 0.0],
[0.5, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 0.0]
])
# W_K: Learns to extract "what I can offer for matching"
# This might extract features about being a noun (potential subject)
W_K = np.array([
[1.0, 0.0, 0.8, 0.0], # Emphasize noun-related features
[0.0, 0.1, 0.0, 0.0],
[0.8, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 0.1]
])
# W_V: Learns to extract "what information I carry"
# This might extract rich semantic content
W_V = np.array([
[0.9, 0.1, 0.7, 0.2], # Preserve semantic meaning
[0.1, 0.9, 0.2, 0.7],
[0.7, 0.2, 0.9, 0.1],
[0.2, 0.7, 0.1, 0.9]
])
# Compute Q, K, V for each word
words = ['the_1', 'cat', 'chased', 'the_2', 'mouse']
print("Computing attention for 'chased':\n")
# Query: What is "chased" looking for?
Q_chased = embeddings['chased'] @ W_Q
print(f"Query from 'chased': {Q_chased}")
print(" (Represents: 'I need information about my subject')\n")
# Keys: What can each word offer for matching?
print("Keys (what each word offers for matching):")
keys = {}
for word in words:
keys[word] = embeddings[word] @ W_K
print(f" {word:8s}: {keys[word]} ", end="")
if word in ['cat', 'mouse']:
print("(I'm a noun, could be a subject)")
elif word in ['the_1', 'the_2']:
print("(I'm a determiner, not a subject)")
else:
print("(I'm the verb)")
# Compute attention scores
print("\nAttention scores (Query · Key):")
scores = {}
for word in words:
scores[word] = np.dot(Q_chased, keys[word])
print(f" {word:8s}: {scores[word]:.3f}")
# Apply softmax to get weights
score_values = np.array(list(scores.values()))
exp_scores = np.exp(score_values - np.max(score_values))
weights = exp_scores / np.sum(exp_scores)
print("\nAttention weights (after softmax):")
for i, word in enumerate(words):
print(f" {word:8s}: {weights[i]:.3f}")
# Values: What information does each word actually contain?
print("\nValues (actual information to retrieve):")
values = {}
for word in words:
values[word] = embeddings[word] @ W_V
print(f" {word:8s}: {values[word]}")
# Compute final output
output = np.zeros(4)
for i, word in enumerate(words):
output += weights[i] * values[word]
print(f"\nFinal attention output for 'chased': {output}")
print("\nThis output is a weighted combination of all Values,")
print("where 'cat' and 'mouse' contribute most because their Keys")
print("matched the Query best, but we retrieve their Values (semantic content).")
# Show the key insight
print("\n" + "="*60)
print("KEY INSIGHT:")
print("="*60)
print("• Query (from 'chased'): 'What subject do I need?'")
print("• Keys (from all words): 'Am I a potential subject?'")
print("• Matching: 'cat' and 'mouse' score high (they're nouns)")
print("• Values: Rich semantic content about what cat/mouse mean")
print("• Output: Semantic information weighted by relevance")
print("\nThe Keys determined WHO to pay attention to.")
print("The Values determined WHAT information to retrieve.")
print("These can be different!")
# Run the demonstration
demonstrate_qkv_difference()
THE FUNDAMENTAL PRINCIPLE
The core principle is this: attention is a content-based routing mechanism.
The Query asks: "What do I need?"
The Keys answer: "Here's how I can help you find it."
The Values provide: "Here's what I actually contain."
By separating these three roles, the model can learn sophisticated patterns like "when looking for a subject, match on grammatical features (Q-K matching), but retrieve semantic features (V)." Or "when resolving a pronoun, match on entity types (Q-K), but retrieve full entity information (V)."
This separation is what makes attention so powerful and flexible. It's not just looking up similar words; it's performing learned, task-specific routing of information based on context.
