A language model needs two vocabulary-sized [V, d] matrices: an input embedding that turns a token id into a vector, and an output projection that turns the final hidden state into one logit per token. Weight tying (Press & Wolf 2017) makes them the same tensor.
untied: x = E[t] logits = h @ U.T // two matrices
tied: x = E[t] logits = h @ E.T // one
Two lines, and now E appears twice in the computation graph, so its gradient is a sum over both appearances and the += is load-bearing.
Hand-written forward and backward pass, every entry checked against central differences: https://dev48.infy.uk/dl/day67-weight-tying.html
Dropping a path is not a broken gradient
The output path is dense - dE[v] += dlogits[v] * h for every v, every step - and the input path is sparse, one row per token looked up. Drop the input path and the result is not garbage:
| gradient into E | wrong rows | rows still exactly right |
|---|---|---|
both paths, +=
|
0 | 64 of 64 |
| output path only | 8 | 56 of 64 (87.5%) |
| input path only | 64 | 0 |
It is exactly correct on every row belonging to a token absent from the batch, and wrong only on the eight that were present. A gradient check that samples random entries passes it, and a loss curve does not notice either: the broken model trains a smooth kink-free curve to 1.84 against 1.61.
A repeat bias nobody put there
Tying makes the logit for token t equal dot(h, E[t]), and a residual stream means h still carries a component of E[t], so that logit contains the largest self-dot-product in the row. The task has no fixed points - a token's true successor is never itself - so any mass on repetition is architecture, not data.
| model | mean rank of the just-read token, out of 64 |
|---|---|
| tied, res = 1 | 1.7 |
| untied | 31.4 |
| chance | 32.5 |
| tied, res = 0 (the control) | 35.0 |
A freshly initialised tied model's top prediction is "repeat the token you just read", and the res = 0 control is what pins that on the residual stream rather than on tying alone. Training suppresses it without removing it: after 600 steps the tied model still ranks the current token 20.3rd against the untied model's 31.9.
Free when the assumption holds
Tying asserts that a token's meaning as an input and as an output live in the same space. Made true or false on purpose, five seeds, scored as excess loss over each task's own entropy:
| task | tied | untied | ratio |
|---|---|---|---|
| ALIGNED | 0.2134 | 0.2143 | 1.00x at 32% fewer parameters |
| MISALIGNED | 0.1657 | 0.1100 | 1.51x |
Aligned, the per-seed ratios straddle one, 0.96 to 1.11. Misaligned, all five seeds agree, 1.37 to 1.64.
What the measurement killed
The first version of this page swept a coarse grid of (sigma, inScale), found nothing that put the residual-stream scale and the logit scale both near 1, and concluded that a tied matrix cannot serve both roles. Grid artifact: a finer sweep finds 19 settings that hit both. The surviving claim is much smaller - the coupling is non-linear, so doubling sigma multiplies the residual scale by 2.00 and the logit scale by 3.80. Tied models are fussier to initialise, not impossible.
The other quoted-not-measured line is that tying halves the parameters. It removes exactly V*d: 23.7% for GPT-2 small, 45.3% for a large-vocabulary toy, 0.5% for a character model, and never half. It matters least for the big models whose parameter counts worried you.
Part of a from-scratch series - one deep-learning idea a day, computed in-browser: https://dev48.infy.uk/deeplearningfromzero.php
Top comments (0)