Implementing a Transformer From Scratch

I wanted to analyze Andrej Karpathy’s NanoGPT implementation of the Transformer in detail. I will also use the Transformer paper as a reference.

Embeddings, positions, and weights

The PyTorch nn.Embedding for the token embedding operates as a lookup table, with one row per token in the vocabulary. So the dimensionality of the matrix is vocab_size x n_embedding. This is a hyperparameter which determines the width of residual stream.

The second embedding matrix encodes position information in a learned embedding. This is different from the Transformer paper, which uses sine and cosine to compute token positions.

Adding the token embedding plus the position embedding combines them into one vector.

The output matrix used to produce the logits before softmax is tied to the token embedding table to save parameters. This is because the embedding table has most of the parameters at this scale.

LayerNorm and residual stream

Using LayerNorm to normalize FFN activations adds stability between layers during training. The Transformer vectors have a dimensionality of 768, so take the mean of each of the 768 vectors for each token, and rescale the activations so that the mean is 0 and the standard deviation is 1. Since this is done per layer, it ignores the composition within the batch.

The implementation used by GPT-2 and Karpathy uses Pre-LayerNorm, where the LayerNorm is computed last. This differs slightly from the Transformer paper using Post-LayerNorm, where LayerNorm is computed first. The difference in the residual stream in code would be something like this:

# Pre-LN
x = x + attn(ln(x))
x = x + mlp(ln(x))
x = x + attn(ln(x))

# Post-LN
x = ln(x + attn(x))
x = ln(x + mlp(x))
x = ln(x + attn(x))

Since we’re using Pre-LayerNorm, LayerNorm has to be applied a final time to get the output logits.

Causal SelfAttention

Attention starts with a fused QKV projection: c_attn is one linear layer of shape 768 x 2304, producing queries, keys, and values in a single matrix multiply, each reshaped to 12 heads of 64 dimensions so every head attends in its own subspace. This is scaled dot-product attention: the scores are QK^T / sqrt(d_k), in code q @ k.transpose(-2, -1) / sqrt(64). The scale comes from a variance argument. A dot product over d_k = 64 unit-variance dimensions has variance 64, and dividing by sqrt(64) restores it, keeping softmax from saturating. The causal mask enforces the autoregressive constraint: torch.tril sets positions above the diagonal to -inf before softmax, so position t never sees the future it is trying to predict. The fused path does the same with is_causal=True. After softmax, the weights multiply v, the heads are reshaped back to (B, T, 768), and c_proj mixes the 12 outputs together. This is W^O from the paper, and without it each head's output would stay in its own 64-dimensional slot of the residual stream.

Loss

The loss is cross-entropy between the logits and the targets, where the target at position t is the token at t+1. The causal mask makes this valid: every position predicts its next token without seeing it, so one forward pass yields B x T training examples rather than one. In code it is a single call, F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)), which flattens the batch and time dimensions and averages over all positions.

MLP

The MLP is the second half of each block. It is two linear layers with a non-linearity in between: c_fc expands 768 to 3072, a 4x widening, and c_proj projects back down to 768. The useful contrast with attention: attention is where tokens exchange information, while the MLP computes on each token independently, with no mixing across positions. Most of the model's parameters live here, since the two matrices together hold about 4.7M parameters per block against attention's 2.4M. The 4x ratio is a convention from the original Transformer paper rather than anything derived, and later models have varied it.

Forward Pass

Each Block is the two pieces already covered, wired with residual connections: x = x + self.attn(self.ln_1(x)), then x = x + self.mlp(self.ln_2(x)). The additions are the residual stream. Each sub-block reads a normalized copy of the stream, computes an update, and adds it back, so information can also flow through the block untouched. The full forward pass is: token and position embeddings are added, the result flows through 12 stacked Blocks, ln_f normalizes the stream one final time, and the tied embedding matrix projects each position's 768-dimensional vector to 50257 logits. Shapes stay (B, T, 768) the whole way through, which is what makes the blocks stackable: every Block reads and writes the same residual stream.