Transformer Architecture Q&A

Course Content

Transformer Architecture Q&A

6 sections · 60 lessons

Define weight tying in GPT-2 and why it’s 124M params vs ~163M raw?

Where GPT-2 small's parameters liveTokenembedding: 38.6MPositionembedding: 0.8M12 blocks x7.09M: 85.1MTotal,tied: 124.4MSeparate LMhead: +38.6Mtopbottom163.0M only if the output matrix is counted as a second copy.
The 39M gap between the two famous numbers is one matrix used twice, not a smaller model.

What you need to know

The count, verified

Python
import torch.nn as nnV, T_max, d, L = 50257, 1024, 768, 12wte, wpe, ln_f = nn.Embedding(V, d), nn.Embedding(T_max, d), nn.LayerNorm(d)# nn.TransformerEncoderLayer has the same weights as a GPT-2 block:# fused QKV + bias, W_o + bias, 768→3072→768 MLP + biases, two LayerNormsblock = nn.TransformerEncoderLayer(d, nhead=12, dim_feedforward=4 * d, norm_first=True)count = lambda m: sum(p.numel() for p in m.parameters())tied = count(wte) + count(wpe) + L * count(block) + count(ln_f)print(f"{count(block):,} per block")                 # 7,087,872print(f"{tied:,} with a tied LM head")               # 124,439,808print(f"{tied + V * d:,} with a separate LM head")   # 163,037,184
Text
token embedding   50,257 × 768   = 38.6Mpositions          1,024 × 768   =  0.8M12 blocks × 7.09M               = 85.1M   (about 12·d² each, plus biases and norms)final norm                       ≈  0.0Mtotal (tied)                     = 124.4M+ separate LM head               = 163.0M

Why it works

  • Symmetry of the job. The input embedding says "this token's meaning is this vector"; the output head says "this vector means this token is likely". Sharing forces the two to agree.
  • More training signal. Rare tokens appear rarely as inputs, but the output softmax updates every row at every step, so tied embeddings of rare tokens are trained more.
  • Evidence. Press and Wolf (2017), "Using the Output Embedding to Improve Language Models", showed that tying lowers perplexity and saves parameters in language models.

When models untie

Untying gives the model more freedom: input and output roles can differ. At large sizes the extra V × d matrix is a small share, so it is affordable:

ModelVocab × dTied?Embedding share of total
GPT-2 small50,257 × 768yes~31%
Llama 3.2 1B128,256 × 2,048yes~21% (counted once)
Llama-3-8B128,256 × 4,096no~13% (two matrices)

Small models with large vocabularies tie, because an untied head would be a big fraction of the model.

A real-life example

A team builds a 135M-parameter model for on-device next-word suggestions in a Hindi-English keyboard app. They use a 64,000-token vocabulary with d = 768, so each embedding matrix is 49M parameters. Untied, the head alone would add 49M parameters — about 36% more to download and hold in phone memory.

They tie the weights. The app's download stays small, and in their tests the tied model's quality matches the untied one, because at this size the two roles benefit from sharing. For a later 8B server model, the same team leaves them untied, since the extra matrix is a small share there.

Follow-up questions to expect

  • "Is anything removed at inference?" — No. Both names point to one tensor; the computation is the same, only memory is saved.
  • "Does tying constrain the model?" — Yes, input and output embeddings must be the same; large models untie to relax that.
  • "How do you tie in PyTorch?" — Assign the same parameter: lm_head.weight = wte.weight, so gradients from both uses add into one tensor.