Transformer Architecture

What you will learn in this lesson:

• Why transformers process an entire sequence in parallel, unlike RNNs
• What self-attention means, specifically, as opposed to attention in general
• Why positional encodings are necessary, since self-attention alone has no sense of order
• Why multi-head attention lets different heads specialize in different relationships
• How residual connections and feedforward sublayers fit into a full transformer block

Trading Recurrence for Parallelism

An RNN's defining feature, processing one time step at a time, carrying a hidden state forward, is also its biggest computational limitation: step t cannot be computed until step t−1 finishes, which makes RNNs inherently sequential and hard to parallelize during training. Transformers remove recurrence entirely. Every position in a sequence is processed simultaneously, using self-attention to relate positions to each other directly, rather than relying on information passed step by step through a hidden state.

Self-Attention: Attention Within One Sequence

The attention lesson's query/key/value mechanism described a query (from the decoder) attending over keys and values (from the encoder), two different sequences. Self-attention is the special case where queries, keys, and values all come from the same sequence: every position computes a query and compares it against every other position's key, letting every token directly relate to every other token in a single operation.

import numpy as np

def softmax(z):
    z = z - np.max(z, axis=-1, keepdims=True)
    e = np.exp(z)
    return e / np.sum(e, axis=-1, keepdims=True)

np.random.seed(0)
n_tokens, d_model = 5, 6

X = np.random.randn(n_tokens, d_model)  # the sequence's embeddings

# Self-attention: Q, K, V are all derived from the same input X
W_q, W_k, W_v = (np.random.randn(d_model, d_model) * 0.3 for _ in range(3))
Q, K, V = X @ W_q, X @ W_k, X @ W_v

scores = (Q @ K.T) / np.sqrt(d_model)
weights = softmax(scores)
output = weights @ V

print("Self-attention weight matrix (every token attends to every token):")
print(np.round(weights, 3))

Every row sums to 1 and represents how much that token's updated representation draws from every other token in the same sequence, including itself.

Positional Encoding: Reintroducing Order

Self-attention, as just implemented, has no notion of position at all: it computes similarity between tokens based purely on their content, regardless of where they sit in the sequence. Shuffle the input tokens and shuffle the rows of the output identically, with nothing in the math itself caring about order. Since order obviously matters for language ("dog bites man" is not "man bites dog"), transformers add a positional encoding to each token's embedding before attention is applied, injecting position information directly into the input the model sees.

import numpy as np
import matplotlib.pyplot as plt

def positional_encoding(n_positions, d_model):
    position = np.arange(n_positions)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
    pe = np.zeros((n_positions, d_model))
    pe[:, 0::2] = np.sin(position * div_term)
    pe[:, 1::2] = np.cos(position * div_term)
    return pe

pe = positional_encoding(n_positions=20, d_model=16)
plt.imshow(pe.T, cmap="RdBu", aspect="auto")
plt.xlabel("Position in sequence"); plt.ylabel("Embedding dimension")
plt.title("Sinusoidal positional encoding")
plt.colorbar()
show_plot()

Each position gets a unique pattern across the embedding dimensions, built from sine and cosine waves at different frequencies. Adding this directly to a token's embedding gives the model a way to recover "where in the sequence am I" purely from the combined vector, information self-attention alone could never supply.

Multi-Head Attention: Specializing in Parallel

A single attention computation can only capture one kind of relationship pattern at a time, however the query/key projections happen to be shaped. Multi-head attention, introduced as a coding exercise in the previous lab, runs several attention computations in parallel, each with its own learned projections, so different heads can specialize: one head might track short, local syntactic relationships, another might track long-range topical relevance, all computed in the same layer and then combined.

A Full Transformer Block

A single transformer block wraps multi-head self-attention together with a feedforward sublayer, with residual connections and normalization around each:

The residual connections here are the exact same idea as ResNet's skip connections from the CNN architectures lesson, giving gradients a direct path through the block, which is just as important for training very deep stacks of transformer blocks (modern transformers commonly stack dozens of these) as it was for very deep CNNs. LayerNorm plays a role analogous to batch normalization, stabilizing the scale of activations flowing between blocks, though it normalizes across each token's own features rather than across a batch.

import numpy as np

def softmax(z):
    z = z - np.max(z, axis=-1, keepdims=True)
    e = np.exp(z)
    return e / np.sum(e, axis=-1, keepdims=True)

def layer_norm(x, eps=1e-6):
    mean = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)
    return (x - mean) / np.sqrt(var + eps)

def relu(z): return np.maximum(0, z)

np.random.seed(0)
n_tokens, d_model, d_ff = 5, 8, 16

x = np.random.randn(n_tokens, d_model)

# Self-attention sublayer
W_q, W_k, W_v = (np.random.randn(d_model, d_model) * 0.2 for _ in range(3))
Q, K, V = x @ W_q, x @ W_k, x @ W_v
attn_out = softmax((Q @ K.T) / np.sqrt(d_model)) @ V
x_after_attn = layer_norm(x + attn_out)  # residual + norm

# Feedforward sublayer
W1, b1 = np.random.randn(d_model, d_ff) * 0.2, np.zeros(d_ff)
W2, b2 = np.random.randn(d_ff, d_model) * 0.2, np.zeros(d_model)
ff_out = relu(x_after_attn @ W1 + b1) @ W2 + b2
block_output = layer_norm(x_after_attn + ff_out)  # residual + norm

print("Input shape:        ", x.shape)
print("After attention sublayer:", x_after_attn.shape)
print("Final block output: ", block_output.shape)