Paper Overview
The paper introduces the Transformer architecture, which has become the foundation of modern natural language processing. It drops recurrence and convolution entirely and draws every dependency between tokens through self-attention. Because each position can attend to every other position in a single step, the model parallelises across a sequence instead of walking it left to right, and long-range dependencies stop being something the network has to carry across many time steps.
Self-Attention and Multi-Head Attention
At the core is scaled dot-product attention: every token emits a query, a key, and a value, and the output for a position is a weighted average of all the values, with weights set by how well that position's query matches the others' keys. A single pass gives the model one view of how tokens relate. Multi-head attention runs several of these in parallel over different learned projections, giving the trained model different representation subspaces whose outputs are concatenated and projected together. The illustrative head patterns below are authored schematics, not learned syntax or coreference from a trained checkpoint.
Architecture Details
Encoder
A stack of N=6 identical layers, each pairing a multi-head self-attention block with a position-wise feed-forward network. Every sub-layer applies dropout to its output, adds a residual connection, then applies layer normalization. This original post-norm order matters: later pre-norm variants move normalization before the sublayer. The feed-forward network uses ReLU.
Decoder
Also six identical layers, but with two changes. Its self-attention is masked so a position attends up to and including its current decoder input. For teacher forcing, those inputs are shifted right by one token: BOS starts the sequence, and position i consumes the previous target rather than the target it predicts. The mask and the shift work together to prevent target leakage. Above that sits a cross-attention block, where decoder queries attend over the encoder's output, conditioning the generated sequence on the source.
Positional Encoding
Attention is order-agnostic on its own, so the Transformer adds sinusoidal positional encodings to the token embeddings — injecting absolute and relative position without reintroducing recurrence.
Implementation
The architecture is complete only when the attention heads feed residuals, normalization, the encoder memory and a shifted decoder. The following nine operations replace the constructor-only fragment with an executable one-layer encoder–decoder. Source and target lengths deliberately differ, so cross-attention cannot accidentally reuse decoder K/V.
Authored offline CPU reference for the original paper architecture: post-norm residuals, ReLU, scaled shared embeddings and interleaved sinusoidal positions. Synthetic IDs and random reduced weights expose the complete tensor path; they do not translate or demonstrate learned syntax, coreference, pretrained semantics, native GPU performance or benchmark reproduction. Equations below explain the cited operations.
The original Tensor2Tensor repository contains several later recipes. At pinned revision 15ea434070b3, transformer_base_v2 explicitly changes to pre-norm; this walkthrough follows the paper’s post-norm order instead. Its literal even/odd positional formulas also come from §3.5, rather than silently inheriting a library timing layout. This is an authored paper reference, not a port of every Tensor2Tensor default.
1. Scale embeddings and add sinusoids
Original paper
The source and shifted target use the same embedding matrix.
Math and tensor shapes
D is d_model. The even/odd coordinates follow the paper exactly, including an unpaired final sine for odd D. Dropout is applied after summing the scaled embeddings and positions (§5.4). Source and target use one shared vocabulary; the same learned matrix is reused for the output readout. Sinusoids supply position, not learned linguistic meaning.
token IDs [B, N] → embeddings [B, N, D]
Python implementation
def sinusoidal_positions(length, d_model, *, dtype=torch.float32):
"""Return [N, D] with paper section 3.5 even/odd coordinates."""
positive_integer(length, "length")
positive_integer(d_model, "d_model")
if not dtype.is_floating_point:
raise ValueError("positional dtype must be floating")
positions = torch.arange(length, dtype=dtype)[:, None]
even = torch.arange(0, d_model, 2, dtype=dtype)
angles = positions / (10000.0 ** (even / d_model))
result = torch.empty(length, d_model, dtype=dtype)
result[:, 0::2] = torch.sin(angles)
# For odd D, the last sine has no corresponding cosine coordinate.
result[:, 1::2] = torch.cos(angles[:, :d_model // 2])
return result
def embed_positions(tokens, embedding, dropout):
"""Scale token embeddings, add positions, then apply residual dropout."""
x = embedding(tokens) * math.sqrt(embedding.embedding_dim)
positions = sinusoidal_positions(x.shape[1], x.shape[2], dtype=x.dtype)
return dropout(x + positions[None])2. Project queries and a separate K/V source
Original paper
Self-attention shares the input; cross-attention uses encoder memory.
Math and tensor shapes
Three learned matrices map the model width to concatenated per-head features. In self-attention Xq and Xm are the same tensor. In encoder–decoder attention, decoder states provide Q while the encoder output supplies both K and V. The reduced reference uses bias-free attention projections as in the paper’s matrix notation. Batch and feature widths are checked explicitly.
Xq [B, Nq, D], Xm [B, Nk, D] → Q, K, V; Nq need not equal Nk
Python implementation
def project_qkv(queries, key_values, q_projection, k_projection, v_projection):
"""Project queries [B,Nq,D] and a separate K/V source [B,Nk,D]."""
validate_features(queries, "queries")
validate_features(key_values, "key_values")
if queries.shape[0] != key_values.shape[0]:
raise ValueError("query and key/value batch sizes must match")
if queries.dtype != key_values.dtype:
raise ValueError("query and key/value dtypes must match")
if (queries.shape[-1] != q_projection.in_features
or key_values.shape[-1] != k_projection.in_features
or key_values.shape[-1] != v_projection.in_features):
raise ValueError("feature width must match the projection inputs")
return (q_projection(queries), k_projection(key_values),
v_projection(key_values))3. Split the width into parallel heads
Original paper
Each head has its own projected representation subspace.
Math and tensor shapes
The reshape groups contiguous projected coordinates into H heads; the transpose puts heads before sequence position for batched multiplication. D must be divisible by H. The paper’s base model uses D=512, H=8 and Dh=64; smoke reduces these to D=8, H=2 and Dh=4. Random reference heads have no learned syntax or coreference specialization.
[B, N, D] → reshape [B, N, H, Dh] → transpose [B, H, N, Dh]
Python implementation
def split_heads(x, num_heads):
"""Map [B,N,D] to [B,H,N,Dh] without mixing token positions."""
validate_features(x, "projected features")
positive_integer(num_heads, "num_heads")
batch, length, width = x.shape
if width % num_heads:
raise ValueError("d_model must be divisible by num_heads")
shaped = x.reshape(batch, length, num_heads, width // num_heads)
return shaped.transpose(1, 2)4. Scale each head’s query–key products
Original paper
The denominator is the key width within a head.
Math and tensor shapes
Scale by sqrt(Dh), not sqrt(D). Under the paper’s assumption of independent zero-mean, unit-variance query/key coordinates, the dot-product variance grows with Dh; dividing by sqrt(Dh) controls that growth. This argument is conditional, not a guarantee that every trained score distribution has unit variance. Softmax will operate over the final key axis.
[B, H, Nq, Dh] × [B, H, Dh, Nk] → scores [B, H, Nq, Nk]
Python implementation
def scaled_scores(q, k):
"""Compute [B,H,Nq,Nk] scores with per-head, not model-width, scale."""
if q.ndim != 4 or k.ndim != 4 or min(q.shape + k.shape) < 1:
raise ValueError("Q and K must be nonempty [B,H,N,Dh]")
if q.shape[:2] != k.shape[:2] or q.shape[-1] != k.shape[-1]:
raise ValueError("Q and K batch, head count and head width must match")
scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
if not torch.isfinite(scores).all():
raise ValueError("attention scores must be finite before masking")
return scores5. Block padded keys and future inputs before softmax
Original paper
Decoder self-attention includes its current input position.
Math and tensor shapes
True in key_padding means blocked. Encoder and cross-attention block padded source keys; decoder self-attention combines shifted-input padding with a lower triangle including the diagonal. The triangle is only accepted for equal-length self-attention. Masking keys does not force padded query outputs to zero: those outputs are excluded from the demonstration loss. Fully masked rows raise ValueError before softmax; returning NaN or inventing a uniform row would violate this API contract.
key padding [B, Nk], causal [Nq, Nk] → broadcast mask [B, 1, Nq, Nk]
Python implementation
def apply_mask(scores, key_padding=None, *, causal=False):
"""Block padded keys and future positions; reject fully masked rows."""
if scores.ndim != 4 or min(scores.shape) < 1:
raise ValueError("scores must be nonempty [B,H,Nq,Nk]")
if not scores.is_floating_point() or not torch.isfinite(scores).all():
raise ValueError("scores must be finite floating values")
batch, _, query_length, key_length = scores.shape
allowed = torch.ones(batch, 1, query_length, key_length, dtype=torch.bool)
if key_padding is not None:
if (not isinstance(key_padding, torch.Tensor)
or key_padding.dtype != torch.bool
or key_padding.shape != (batch, key_length)):
raise ValueError(
"key_padding must be boolean [B,Nk], True = blocked")
if key_padding.device.type != "cpu":
raise ValueError("key_padding must be a CPU tensor")
allowed = allowed & ~key_padding[:, None, None, :]
if causal:
if query_length != key_length:
raise ValueError("causal mask is for equal-length self-attention")
# Lower triangle includes the diagonal, hence the current input token.
triangle = torch.ones(
query_length, key_length, dtype=torch.bool).tril()
allowed = allowed & triangle
if not allowed.any(dim=-1).all():
raise ValueError("fully masked attention row has no valid softmax")
return scores.masked_fill(~allowed, -torch.inf)6. Mix values and project the concatenated heads
Original paper
Each head produces a weighted value vector per query.
Math and tensor shapes
The stable library softmax subtracts a row maximum internally. Negative-infinity entries receive exactly zero probability. Multiplying by V sums over keys, and transposing back before reshaping preserves the head coordinate order. W_O mixes all heads into model-width output. Tests compare the whole vectorized result against an independent per-head, per-query calculation, including Nq=3 and Nk=5.
weights [B, H, Nq, Nk] → head outputs [B, H, Nq, Dh] → [B, Nq, D]
Python implementation
def attend_and_merge(masked_scores, values, output_projection):
"""Softmax, mix V, concatenate heads, then apply learned W_O."""
if (masked_scores.ndim != 4 or values.ndim != 4
or min(masked_scores.shape + values.shape) < 1
or masked_scores.shape[:2] != values.shape[:2]
or masked_scores.shape[-1] != values.shape[-2]
or values.shape[1] * values.shape[-1]
!= output_projection.in_features):
raise ValueError("score, value and output-projection shapes must match")
if (torch.isnan(masked_scores).any() or torch.isposinf(masked_scores).any()
or not torch.isfinite(masked_scores).any(dim=-1).all()
or not torch.isfinite(values).all()):
raise ValueError(
"attention needs finite values and an unmasked key per row")
weights = torch.softmax(masked_scores, dim=-1)
context = weights @ values
batch, heads, length, head_width = context.shape
merged = context.transpose(1, 2).contiguous().reshape(
batch, length, heads * head_width)
return output_projection(merged), weights7. Complete the original post-norm encoder layer
Original paper
Normalization follows each residual addition.
Math and tensor shapes
Attention sees the unnormalized layer input. Each residual sum is normalized over the D coordinates at each token, then the position-wise FFN uses ReLU. Residual dropout is before addition, as §5.4 specifies. LayerNorm epsilon=1e-6 is an explicit numerical choice; the paper does not specify it. A fixed-input oracle checks this order and the population-variance normalization. There is no pre-norm, GELU, RoPE or extra final normalization.
attention [B, Ns, D]; FFN [B, Ns, D] → [B, Ns, Dff] → [B, Ns, D]
Python implementation
class EncoderLayer(nn.Module):
"""Original post-norm self-attention and position-wise ReLU FFN."""
def __init__(self, d_model, heads, d_ff, dropout=0.1):
super().__init__()
positive_integer(d_ff, "d_ff")
self.attention = MultiHeadAttention(d_model, heads)
self.ffn = nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(),
nn.Linear(d_ff, d_model))
self.norm1 = nn.LayerNorm(d_model, eps=1e-6)
self.norm2 = nn.LayerNorm(d_model, eps=1e-6)
self.dropout = nn.Dropout(dropout)
def forward(self, x, source_padding=None):
"""LayerNorm is after addition, never before the attention or FFN."""
attention, weights = self.attention(x, x, source_padding)
x = self.norm1(x + self.dropout(attention))
x = self.norm2(x + self.dropout(self.ffn(x)))
return x, weights8. Add masked self-attention and encoder cross-attention
Original paper
The middle sublayer reads the encoder output, not decoder K/V.
Math and tensor shapes
First, masked self-attention sees only shifted decoder inputs up to and including its position. Next, decoder queries read encoder memory M with the source-padding mask. Finally, the same ReLU FFN pattern runs at every target position. All three residual additions precede LayerNorm. The oracle verifies these three operations with source length 3 and target length 2; behavioral tests perturb future inputs and padded memory.
decoder [B, Nt, D], memory [B, Ns, D] → cross weights [B, H, Nt, Ns]
Python implementation
class DecoderLayer(nn.Module):
"""Masked self-attention, encoder K/V cross-attention, then ReLU FFN."""
def __init__(self, d_model, heads, d_ff, dropout=0.1):
super().__init__()
positive_integer(d_ff, "d_ff")
self.self_attention = MultiHeadAttention(d_model, heads)
self.cross_attention = MultiHeadAttention(d_model, heads)
self.ffn = nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(),
nn.Linear(d_ff, d_model))
self.norm1 = nn.LayerNorm(d_model, eps=1e-6)
self.norm2 = nn.LayerNorm(d_model, eps=1e-6)
self.norm3 = nn.LayerNorm(d_model, eps=1e-6)
self.dropout = nn.Dropout(dropout)
def forward(self, x, memory, target_padding=None, source_padding=None):
"""Memory length Nk may differ from decoder query length Nq."""
attended, self_weights = self.self_attention(
x, x, target_padding, causal=True)
x = self.norm1(x + self.dropout(attended))
attended, cross_weights = self.cross_attention(
x, memory, source_padding)
x = self.norm2(x + self.dropout(attended))
x = self.norm3(x + self.dropout(self.ffn(x)))
return x, self_weights, cross_weights9. Run the shifted encoder–decoder tensor path
Original paper
Offset target embeddings by one position before decoding.
Math and tensor shapes
The one-layer reference reserves PAD=0 and BOS=1 in a shared 13-token synthetic vocabulary. Its forward API accepts target labels and shifts them internally: predicting target i never consumes target i. The embedding matrix E is shared by source, target and readout (§3.4). Smoke computes mean cross-entropy over nonpadding labels, backpropagates through every parameter and takes one illustrative SGD step. Dropout is configured at 0.1 but disabled in eval mode for deterministic checks; gradients still run. The paper’s six-layer depth, width512, FFN2048, training schedule, label smoothing and beam search are outside this reduced mechanics example.
source [2, 5], target [2, 3] → encoder [2, 5, 8], decoder [2, 3, 8], logits [2, 3, 13]
Python implementation
def shift_targets(targets, bos_id=1):
"""Input at position i is BOS or target i-1, never target i."""
bos = torch.full_like(targets[:, :1], bos_id)
return torch.cat((bos, targets[:, :-1]), dim=1)
class ReducedTransformer(nn.Module):
"""One encoder/decoder layer, D=8, H=2, FFN=16, vocabulary=13 by default."""
def __init__(self, vocab_size=13, d_model=8, heads=2, d_ff=16, dropout=0.1):
super().__init__()
positive_integer(vocab_size, "vocab_size")
if vocab_size < 3:
raise ValueError(
"vocabulary needs PAD=0, BOS=1 and a content token")
self.embedding = nn.Embedding(vocab_size, d_model)
self.embedding_dropout = nn.Dropout(dropout)
self.encoder = EncoderLayer(d_model, heads, d_ff, dropout)
self.decoder = DecoderLayer(d_model, heads, d_ff, dropout)
def forward(self, source, targets):
"""Return states/logits and masks' probabilities for teacher forcing."""
validate_tokens(source, self.embedding.num_embeddings, "source")
validate_tokens(targets, self.embedding.num_embeddings, "targets")
if source.shape[0] != targets.shape[0]:
raise ValueError("source and target batch sizes must match")
decoder_input = shift_targets(targets)
source_padding = source == 0
target_padding = decoder_input == 0
memory, encoder_weights = self.encoder(
embed_positions(source, self.embedding, self.embedding_dropout),
source_padding)
decoded, self_weights, cross_weights = self.decoder(
embed_positions(
decoder_input, self.embedding, self.embedding_dropout),
memory, target_padding, source_padding)
# One shared matrix for source embedding, target embedding and readout.
logits = F.linear(decoded, self.embedding.weight)
return {
"memory": memory, "decoded": decoded, "logits": logits,
"decoder_input": decoder_input, "encoder_weights": encoder_weights,
"self_weights": self_weights, "cross_weights": cross_weights,
}Actual smoke JSON reports source SHA-256, seed, reduced dimensions, all state shapes, finite loss and parameter gradient norms, one SGD update, exactly zero future-target and padded-source interference, and zero probability on masked source keys. Fully masked rows are rejected. This is mechanics evidence from random weights.
transformer.pyrequirements-transformer.txt
python -m pip install -r requirements-transformer.txt python transformer.py --smoke --seed 7
With that tensor path explicit, the attention machinery carries forward into BERT's bidirectional encoder and the memory and execution questions explored in FlashAttention. The reference verifies the original operations; the models below show where training and scale took them.
Practical Impact
The Transformer architecture has revolutionized NLP and beyond:
-
Foundation for BERT, GPT, and other models
- Enabled pre-training on massive text corpora
- Led to state-of-the-art results across NLP tasks
-
Cross-domain Applications
- Computer Vision (ViT)
- Speech Recognition
- Protein Structure Prediction (AlphaFold)
Critical Analysis
Strengths
- Parallel processing capability
- Better handling of long-range dependencies
- Scalability to large datasets
Limitations
- Quadratic memory complexity with sequence length
- Requires large amounts of training data
- Computationally intensive training
Personal Notes
In my experience implementing Transformers, the key challenges include:
- Managing attention matrix memory for long sequences
- Proper initialization of positional encodings
- Balancing the number of attention heads