NLP Foundations
Weight Tying of Input and Output Embeddings
Why most language models score every candidate next token using the transpose of the exact same matrix that looked up the input tokens, cutting parameters and improving perplexity in one move.
intermediate · 7 min read
A transformer needs two matrices shaped by the vocabulary: one to turn an input token id into a vector (the embedding matrix E, see input-embeddings-and-lookup), and one to turn the final hidden state back into a score for every vocabulary entry (the unembedding matrix U, feeding softmax-logits). The naive design trains these separately. Almost nobody does that.
The trick
Weight tying sets U = E^T: the unembedding matrix is literally the transpose of the input embedding matrix, not a separately learned one. A single set of vocab_size x d_model parameters does double duty, once as a lookup table, once (transposed) as a projection to logits. Press and Wolf, 2016 showed this is not just a parameter-saving hack, it measurably improves perplexity, and it became close to standard practice, used in GPT-2 among many others.
The intuition: if two tokens are similar enough in meaning to warrant similar input vectors, they should also be similar enough in meaning to warrant similar output probabilities in matching contexts. Tying enforces that the model uses one consistent geometry for "what this token means" rather than learning two independently, which turns out to be an easier optimisation problem and removes an entire vocab_size x d_model block of redundant parameters, no small saving when the vocabulary is 50,000 to 250,000 tokens.
What it buys and what it costs
The parameter saving is largest relative to total model size in smaller models; in a model with tens of billions of parameters, one embedding table is a rounding error, and some large-scale training recipes untie the matrices anyway, having found that letting input and output geometries diverge slightly does not hurt and occasionally helps at that scale. So weight tying is close to universal at small-to-mid scale and a genuine design choice, not a law, at the largest scale.
Tying also has a side effect that matters for interpretability: it is what makes the logit lens technique clean. Because the same matrix that produced the input vectors also reads out logits, an intermediate layer's residual stream (see the-residual-stream) can be projected through that shared matrix and produce a meaningful, if early, guess at the next token.
When it falls down
- One matrix serving two jobs is a constraint, not a free lunch. "What does this token mean as input" and "how likely is this token as an output given the context" are related but not identical questions; tying forces both to share geometry, which can be suboptimal in principle even though it tends to help in practice at moderate scale.
- Untied models lose the clean logit-lens interpretation. If
Uis notE^T, an early-layer residual stream projected through the final unembedding is reading a basis it was never trained to be compatible with at that depth, so the technique becomes noisier or invalid. - Frequency and semantics get entangled through the bias term. The unembedding step typically adds a learned bias per vocabulary entry separate from the tied weights, so raw token frequency is partly handled outside the tied geometry; tying does not mean frequency and meaning are cleanly separated.
- The saving is scale-dependent. At vocabulary sizes in the hundreds of thousands with very large
d_model, the embedding table itself can be enormous even before doubling it, so the relative benefit of tying shrinks compared to a small model where it can be a large share of total parameters.
Further reading
- Press and Wolf, 2016, Using the Output Embedding to Improve Language Models, arXiv:1608.05859 - the paper that popularised weight tying and measured its perplexity effect.
- Vaswani et al., 2017, Attention Is All You Need, arXiv:1706.03762 - the original transformer, which ties embedding and pre-softmax weights.