We are not optimizing for production here — the code is written as explicitly as possible for learning, even when a library would do it in fewer lines.
This chapter builds the original Transformer from the ground up — the same encoder–decoder stack in the diagram above, implemented in PyTorch one piece at a time. We start with hyperparameters in a single config object, then walk through each layer in order: embeddings, positional encoding, layer norm, feed-forward blocks, multi-head attention, residuals, encoder and decoder stacks, and finally the full model.
Each section below is one building block. You get a short explanation, the code for that module, and how it connects to what came before. By the end you have a runnable Transformer class, not just a diagram.
One config object for the whole model
Before any module, we define a single TransformerConfig dataclass. Every class below takes config instead of a long list of constructor arguments — so when we evolve this codebase in later chapters (GPT-2, LLaMA, …), we only add or change fields in one place.
Defaults follow the original Transformer base model: d_model=512, d_ff=2048, h=8, N=6, dropout=0.1, vocabulary size 37,000, and per-head dimension d_k=64.
dk property: dmodel // H → 512 // 8 = 64. Multi-head attention splits d_model across H heads; d_model must be divisible by H.
dmodel512 — embedding widthvocabsize37,000 — vocabulary sizeB2 — batch size for demosmax_seq_len512 — max sequence lengthdff2048 — feed-forward inner dimensionH8 — attention headsN6 — encoder/decoder blocks eachdropout0.1 — dropout probabilitynorm_eps1e-6 — LayerNorm epsilon1from dataclasses import dataclass
2import torch
3import torch.nn as nn
4import math
5
6@dataclass
7class TransformerConfig:
8 dmodel: int = 512 # embedding / model dimension
9 vocabsize: int = 37000 # vocabulary size
10 B: int = 2 # batch size (for demo runs)
11 max_seq_len: int = 512 # max sequence length (src and tgt)
12 dff: int = 2048 # feed-forward inner dimension
13 H: int = 8 # number of attention heads
14 N: int = 6 # number of encoder/decoder blocks
15 dropout: float = 0.1 # dropout probability
16 norm_eps: float = 1e-6 # epsilon for LayerNorm numerical stability
17
18 @property
19 def dk(self):
20 # per-head dimension — dmodel must be divisible by H
21 return self.dmodel // self.H
22
23
24config = TransformerConfig() # shared by every module belowMap token IDs to dense vectors
Token indices select rows of a learned embedding matrix E of shape (vocab × d_model). The model has no prior notion of word identity beyond these vectors.
Scaling by √d_model keeps the magnitude of embedding outputs comparable to the fixed sinusoidal position encodings that are added in the next step—stabilizing the combined signal early in the stack.
1class InputEmbedding(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 self.embedding = nn.Embedding(config.vocabsize, config.dmodel)
5 # scale embeddings by sqrt(dmodel) — from "Attention Is All You Need" sec 3.4
6 self.scale = math.sqrt(config.dmodel)
7
8 def forward(self, x):
9 return self.embedding(x) * self.scaleInject order with fixed sines and cosines
Attention is permutation-invariant without position information. Sinusoids at different frequencies let the model attend to relative position via linear combinations, as argued in the paper—without learning position embeddings of length max sequence from scratch (though learned embeddings are a valid alternative).
The log / exp construction avoids underflow in the 10000^{2i/d_model} divisor. register_buffer stores pe as a non-parameter tensor that still moves to GPU with the module.
1# PE(pos, 2i) = sin(pos / 10000^(2i/dmodel))
2# PE(pos, 2i+1) = cos(pos / 10000^(2i/dmodel))
3class PositionalEncoding(nn.Module):
4 def __init__(self, config) -> None:
5 super().__init__()
6 pe = torch.zeros(config.max_seq_len, config.dmodel) # (max_seq_len, dmodel)
7 pos = torch.arange(0, config.max_seq_len, dtype=torch.float).unsqueeze(1) # (max_seq_len, 1)
8 # unsqueeze(1): (max_seq_len,) → (max_seq_len, 1) so pos * divTerm → (max_seq_len, dmodel // 2)
9 divTerm = 1 / 10000 ** (
10 torch.arange(0, config.dmodel, 2, dtype=torch.float) / config.dmodel
11 ) # (dmodel // 2,)
12 pe[:, 0::2] = torch.sin(pos * divTerm) # (max_seq_len, dmodel // 2) → even cols
13 pe[:, 1::2] = torch.cos(pos * divTerm) # (max_seq_len, dmodel // 2) → odd cols
14 pe.unsqueeze_(0) # (1, max_seq_len, dmodel)
15 # unsqueeze(0): (max_seq_len, dmodel) → (1, max_seq_len, dmodel) — broadcasts to batch in forward
16 self.register_buffer("pe", pe) # (1, max_seq_len, dmodel)
17
18 def forward(self, x):
19 # x: (B, S, dmodel)
20 return x + self.pe[:, : x.shape[1], :] # (B, S, dmodel) + (1, S, dmodel) → (B, S, dmodel)Normalize along the feature axis
LayerNorm normalizes each position’s d_model-vector to zero mean and unit variance (along the last dimension). That stabilizes activations across depth.
α and bias are learnable: they start at 1 and 0 so the layer is initially an identity on normalized values, then can rescale and shift what the next sublayer sees. ε avoids dividing by zero when variance is tiny.
1# z = (x - mean) / (std + eps) * alpha + bias
2class LayerNormalization(nn.Module):
3 def __init__(self, config) -> None:
4 super().__init__()
5 self.eps = config.norm_eps
6 self.alpha = nn.Parameter(torch.ones(config.dmodel)) # (dmodel,) learnable scale
7 self.bias = nn.Parameter(torch.zeros(config.dmodel)) # (dmodel,) learnable shift
8
9 def forward(self, x):
10 # x: (B, S, dmodel)
11 mean = x.mean(dim=-1, keepdim=True) # (B, S, 1)
12 std = x.std(dim=-1, keepdim=True) # (B, S, 1)
13 return self.alpha * (x - mean) / (std + self.eps) + self.bias # (B, S, dmodel)Expand, ReLU, contract
Each position is processed independently: a linear map to a larger inner dimension d_ff (often 4× d_model), a ReLU nonlinearity, dropout, then a linear back to d_model.
This is the “MLP” inside every encoder/decoder layer; it increases expressivity without mixing sequence positions (that mixing is attention’s job).
1class FeedForward(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 self.layer1 = nn.Linear(config.dmodel, config.dff) # dmodel → dff per position
5 self.layer2 = nn.Linear(config.dff, config.dmodel) # dff → dmodel per position
6 self.dropout = nn.Dropout(config.dropout)
7
8 def forward(self, x):
9 # x: (B, S, dmodel) → (B, S, dff) → relu → dropout → (B, S, dmodel)
10 return self.layer2(self.dropout(torch.relu(self.layer1(x)))) # (B, S, dmodel)Scaled dot-product attention, h times in parallel
Linear layers produce Q, K, V; then views/splits reshape to (B, h, S, d_k) with d_k = d_model / h. Each head runs the same attention rule, outputs are concatenated along the feature axis and mixed by W_o.
The mask sets disallowed positions to −∞ (float("-inf")) before softmax (e.g. padding or future tokens). Dropout is applied on the attention weights, not on values first.
1class MultiHeadAttention(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 self.dmodel = config.dmodel
5 self.h = config.H
6 self.dk = config.dk # dmodel // H
7 self.wq = nn.Linear(config.dmodel, config.dmodel) # dmodel → dmodel
8 self.wk = nn.Linear(config.dmodel, config.dmodel) # dmodel → dmodel
9 self.wv = nn.Linear(config.dmodel, config.dmodel) # dmodel → dmodel
10 self.wo = nn.Linear(config.dmodel, config.dmodel) # dmodel → dmodel
11 self.dropout = nn.Dropout(p=config.dropout)
12
13 @staticmethod
14 def attention(q, k, v, dropout, mask):
15 # q, k, v: (B, H, S, dk)
16 _, _, _, dk = q.shape
17 # (B, H, S, dk) @ (B, H, dk, S) → (B, H, S, S)
18 attention_score = q @ k.transpose(-1, -2) / math.sqrt(dk)
19 if mask is not None:
20 # mask==0 positions get -inf so softmax ignores them
21 attention_score = attention_score.masked_fill(mask == 0, float("-inf"))
22 attention_score = attention_score.softmax(dim=-1) # (B, H, S, S)
23 return dropout(attention_score) @ v # (B, H, S, dk)
24
25 def forward(self, q, k, v, mask):
26 B, S, _ = q.shape # q: (B, S, dmodel)
27 _, Sk, _ = k.shape # k: (B, Sk, dmodel)
28 _, Sv, _ = v.shape # v: (B, Sv, dmodel)
29
30 # project and split into H heads
31 query = self.wq(q).view(B, S, self.h, self.dk).transpose(1, 2) # (B, H, S, dk)
32 key = self.wk(k).view(B, Sk, self.h, self.dk).transpose(1, 2) # (B, H, Sk, dk)
33 value = self.wv(v).view(B, Sv, self.h, self.dk).transpose(1, 2) # (B, H, Sv, dk)
34 x = MultiHeadAttention.attention(
35 query, key, value, self.dropout, mask
36 ) # (B, H, S, dk)
37 # merge heads back: (B, H, S, dk) → (B, S, H, dk) → (B, S, dmodel)
38 x = x.transpose(1, 2).contiguous().view(B, S, self.dmodel) # (B, S, dmodel)
39 return self.wo(x) # (B, S, dmodel)Sublayer, dropout, add, post-norm
This is the pre-LN pattern: normalize the input, run the sublayer (attention or FFN), apply dropout, then add the residual. The paper’s diagram is often drawn as post-LN; the math here matches the common “norm first” implementation you see in many codebases.
sublayer is a callable—so the same wrapper drives both self-attention and the position-wise FFN without duplicating wiring.
1class ResidualConnection(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 # original transformer uses post-norm: LayerNorm(x + sublayer(x))
5 self.norm = LayerNormalization(config)
6 self.dropout = nn.Dropout(p=config.dropout)
7
8 def forward(self, x, sublayer):
9 return self.norm(x + self.dropout(sublayer(x)))Self-attention and FFN, repeated N times
Each EncoderBlock wires self-attention (Q, K, V all from the same sequence) and the feed-forward stack, each wrapped in a residual connection. The lambda keeps the call shape uniform for ResidualConnection.
The Encoder applies a stack of N blocks, then a final layer norm on the sequence of vectors passed to the decoder’s cross-attention.
1class EncoderBlock(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 self.MultiHeadAttentionLayer = MultiHeadAttention(config)
5 self.FeedForwardLayer = FeedForward(config)
6 # two residual connections — one around self-attn, one around FFN
7 self.ResidualConnectionLayer = nn.ModuleList(
8 [ResidualConnection(config) for _ in range(2)]
9 )
10
11 def forward(self, x, mask):
12 # self-attention: Q=K=V=x (encoder attends to itself)
13 x = self.ResidualConnectionLayer[0](
14 x, lambda x: self.MultiHeadAttentionLayer(x, x, x, mask)
15 )
16 x = self.ResidualConnectionLayer[1](x, self.FeedForwardLayer)
17 return x
18
19class Encoder(nn.Module):
20 def __init__(self, config) -> None:
21 super().__init__()
22 self.layers = nn.ModuleList([EncoderBlock(config) for _ in range(config.N)])
23 self.norm = LayerNormalization(config)
24
25 def forward(self, x, mask):
26 for layer in self.layers:
27 x = layer(x, mask)
28 return self.norm(x)Masked self-attention, cross-attention, FFN
DecoderBlock adds a third residual sublayer: cross-attention where queries come from the decoder state and keys/values come from the encoder output. The first block still uses a causal mask on self-attention so position i cannot attend past i.
Cross-attention passes encoder_output for both K and V (the common “attend to source” pattern in seq2seq transformers).
1class DecoderBlock(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 self.MultiHeadAttentionLayer = MultiHeadAttention(config)
5 self.CrossMultiHeadAttentionLayer = MultiHeadAttention(config)
6 self.FeedForwardLayer = FeedForward(config)
7 # three residual connections — self-attn, cross-attn, FFN
8 self.ResidualConnectionLayer = nn.ModuleList(
9 [ResidualConnection(config) for _ in range(3)]
10 )
11
12 def forward(self, x, encoder_output, src_mask, tgt_mask):
13 # masked self-attention: decoder attends to its own previous tokens
14 x = self.ResidualConnectionLayer[0](
15 x, lambda x: self.MultiHeadAttentionLayer(x, x, x, tgt_mask)
16 )
17 # cross-attention: Q from decoder, K/V from encoder output
18 x = self.ResidualConnectionLayer[1](
19 x,
20 lambda x: self.CrossMultiHeadAttentionLayer(
21 x, encoder_output, encoder_output, src_mask
22 ),
23 )
24 x = self.ResidualConnectionLayer[2](x, self.FeedForwardLayer)
25 return x
26
27class Decoder(nn.Module):
28 def __init__(self, config) -> None:
29 super().__init__()
30 self.layers = nn.ModuleList([DecoderBlock(config) for _ in range(config.N)])
31 self.norm = LayerNormalization(config)
32
33 def forward(self, x, encoder_output, src_mask, tgt_mask):
34 for layer in self.layers:
35 x = layer(x, encoder_output, src_mask, tgt_mask)
36 return self.norm(x)Logits to log-probabilities
After the decoder stack, a linear map produces one score per vocabulary entry at each time step. log_softmax is numerically well-behaved and pairs naturally with the negative log-likelihood training objective used in many implementations.
1class ProjectionLayer(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 # project dmodel → vocabsize, then log_softmax for log-probabilities
5 self.layer1 = nn.Linear(config.dmodel, config.vocabsize)
6
7 def forward(self, x):
8 return torch.log_softmax(self.layer1(x), dim=-1)Assemble, stack, initialize
Transformer.__init__ builds both embedding+PE pipelines, encoder, decoder, and projection head — all from the same config defined in §00. Three methods wire the full forward path: encode → decode → project.
That is the end-to-end object you train for translation (or any task framed as source→target autoregression). With the paper defaults in §00, that is the full base Transformer at 512 / 6 / 8 / 2048.
1class Transformer(nn.Module):
2 def __init__(self, config) -> None:
3 super().__init__()
4 # src and tgt get separate embeddings + positional encodings
5 self.src_embedding = InputEmbedding(config)
6 self.tgt_embedding = InputEmbedding(config)
7 self.src_positionalEnc = PositionalEncoding(config)
8 self.tgt_positionalEnc = PositionalEncoding(config)
9 self.encoder = Encoder(config)
10 self.decoder = Decoder(config)
11 self.projection = ProjectionLayer(config)
12 self.dropout = nn.Dropout(p=config.dropout)
13
14 def encode(self, x, mask):
15 x = self.src_embedding(x)
16 x = self.src_positionalEnc(x)
17 x = self.dropout(x)
18 return self.encoder(x, mask)
19
20 def decode(self, x, encoder_output, src_mask, tgt_mask):
21 x = self.tgt_embedding(x)
22 x = self.tgt_positionalEnc(x)
23 x = self.dropout(x)
24 return self.decoder(x, encoder_output, src_mask, tgt_mask)
25
26 def project(self, x):
27 return self.projection(x)
28
29# --- smoke test ---
30transformer = Transformer(config) # config from §00
31src = torch.randint(0, config.vocabsize, (config.B, config.max_seq_len))
32tgt = torch.randint(0, config.vocabsize, (config.B, config.max_seq_len))
33enc_out = transformer.encode(src, None)
34dec_out = transformer.decode(tgt, enc_out, None, None)
35proj_out = transformer.project(dec_out)
36print("proj_out.shape:", proj_out.shape)
37total_params = sum(p.numel() for p in transformer.parameters())
38print("Total parameters:", total_params)Papers & technical sources
Primary reports and references for this chapter. Read these for full equations, training details, and official hyperparameters.
- Attention Is All You Need (Vaswani et al., 2017)Original Transformer paper — encoder–decoder, scaled dot-product attention, positional encoding.
- The Annotated Transformer (Alammar)Readable walkthrough of the same architecture with diagrams.
Comments