Weight tying is one of the cheapest, most reliable wins in language-model design: instead of learning two separate V × d matrices — one to turn token IDs into vectors at the input and another to turn hidden states back into vocabulary logits at the output — you learn one matrix and use it in both places. The output projection is simply the transpose of the input embedding. That single decision deletes a full V × d parameter block, which for a small model is often the largest block in the whole network, and it usually improves perplexity rather than hurting it. This piece works through why the trick is sound (input and output token spaces are the same space), exactly how many parameters it saves with a numeric example, why small and CPU-bound models benefit most, the softmax-bottleneck caveat that tells you where its expressiveness limit lies, and how it generalizes into broader parameter-sharing schemes like ALBERT’s cross-layer sharing.

The two matrices most models keep separate

A transformer language model touches the vocabulary twice. At the input it holds an embedding matrix E of shape [V, d] — one row per vocabulary token, each row a d-dimensional vector. Embedding a token is just a row lookup: token ID t becomes E[t], a vector in ℝ^d. At the output, after the final transformer block produces a hidden state h ∈ ℝ^d, the model needs a score for every vocabulary token, so it multiplies by an output projection (the ‘unembedding’ or LM head) W_out of shape [d, V]: logits = h · W_out, giving one logit per token.

Left to themselves these are two independent learnable matrices, each with V × d entries. For a modern vocabulary (tens of thousands of tokens) and a modest hidden size, each is millions of parameters. The observation behind weight tying is that these two matrices are doing mirror-image jobs on the same object — the vocabulary — and there is no fundamental reason they should be learned separately. That symmetry is the whole opening for the optimization.

Advertisement

What tying actually does

Weight tying (also called tied embeddings or the shared input-output embedding) sets the output projection equal to the transpose of the input embedding: W_out = E^T. Concretely, the logit for token t becomes the dot product of the hidden state with that token’s input embedding row:

Input:   x_t   = E[t]            E: [V, d]   (row lookup)
Output:  logit_t = h · E[t]        logits = h · E^T   -> [V]
p = softmax(logits)                    probability over vocab

So there is exactly one matrix E in the model. It is read row-wise on the way in and used as a bank of classifier directions on the way out: the model predicts token t strongly when the final hidden state points in the same direction as t’s embedding. In code it is a one-liner — the LM head shares the storage of the embedding table — and in most frameworks a single flag or a shared parameter reference. The gradient from the output loss and the gradient from the input path now both flow into the same weights, which we will see is a feature, not a bug.

The parameter count, exactly

The arithmetic is the reason people bother. Without tying you store two matrices, so the vocabulary-facing parameter budget is:

untied:  |E| + |W_out| = V·d + d·V = 2·V·d
tied:    |E|            = V·d
saved:   2·V·d - V·d = V·d   (exactly one matrix)

Tying removes precisely V·d parameters — a full copy of the embedding table. Note that biases in the LM head, if present, are tiny (V numbers, and often omitted), so the V·d figure is the whole story. Whether that saving is negligible or enormous depends entirely on how large V·d is relative to the rest of the model. The bulk of a transformer’s non-embedding parameters lives in the attention and feed-forward blocks, roughly 12 · L · d^2 for L layers (the FFN’s 8d^2 plus attention’s 4d^2 per layer). Compare that quadratic-in-d, linear-in-L bulk against the V·d embedding, and you can see the balance tips toward the embedding when V is large and L and d are small — exactly the small-model regime.

A worked param-savings example

Take a concrete small model: vocabulary V = 32,000 (a typical SentencePiece size), hidden size d = 512, and L = 8 layers. Work the blocks out:

embedding table   V·d      = 32000 × 512      = 16.38M
output proj       V·d      = 32000 × 512      = 16.38M   (if untied)
transformer body  ~12·L·d^2 = 12 × 8 × 512^2   = 25.17M

untied total  = 16.38 + 16.38 + 25.17 = 57.9M
tied total    = 16.38 + 25.17          = 41.5M
saving        = 16.38M  (~28% of the whole model)

Tying shaves roughly 28% off this model’s parameter count — and it does so by removing weights that were, in the untied case, the single largest tensor tied for first place. Push d up to 4096 and stack 32 layers (a big model) and the body balloons to about 12·32·4096^2 ≈ 6.4B parameters while each embedding is only 32000·4096 ≈ 131M; now tying saves well under 2%. Same trick, wildly different payoff. The lesson is structural: the embedding is a fixed V·d cost that the rest of the model grows past, so the smaller the model, the more tying matters.

Why the trick is theoretically sound

The deeper justification is that the input and output token spaces coincide. Both matrices are asked to encode ‘what does token t mean, geometrically?’ The input embedding answers by placing t at a point in ℝ^d so that the model can compute with it; the output projection answers by placing a classifier direction in the same ℝ^d so that a hidden state pointing that way predicts t. These are two views of one latent geometry of the vocabulary. If token embeddings that are close in input space are also interchangeable in output space — and for natural language they largely are — then forcing the two matrices to be one is not a compromise but an inductive bias that reflects the true structure.

Press & Wolf (2017) and Inan et al. (2017) made this precise: the untied output projection, when trained well, ends up highly correlated with the input embedding anyway, so tying just bakes in a regularity the model would otherwise spend parameters and data rediscovering. Inan et al. framed it through the language-modeling loss itself, showing tying corresponds to a sensible similarity structure on the output distribution. The tie is a prior that says: a token’s ‘meaning’ is one thing, used symmetrically.

Regularization, not just compression

It would be easy to read tying as pure compression — fewer parameters, accept a small accuracy hit. The striking empirical fact is the opposite: on small and medium models, tying usually lowers perplexity. Fewer parameters, better results. The reason is that tying acts as a strong regularizer. An untied output projection has millions of free parameters, many of which correspond to tokens seen only a handful of times in training; those rows are poorly estimated and prone to overfitting. Tying forces every token’s output direction to be its (much better trained, because it is used on every occurrence at the input) embedding, sharing statistical strength across the two roles.

There is also a data-efficiency argument. Each embedding row now receives gradient signal from two sources — whenever the token appears in the input and whenever it is a prediction target — so rare tokens get roughly twice the learning signal per epoch. For low-resource settings, small corpora, and the compute-starved CPU-SLM regime this series cares about, that doubled signal on the long tail of the vocabulary is a real and free improvement in how well rare words are modeled.

The empirical effect on small models

Concretely, weight tying became standard because it kept winning. Press & Wolf reported perplexity improvements across word- and sub-word language models and neural machine translation with no downside, and the practice propagated: GPT-2 ties its embeddings, and countless small LMs since do the same by default. The effect is largest exactly where the embedding is a big fraction of the model. For a 125M-class model with a 50K vocabulary and d = 768, the two embedding tables would otherwise be roughly 2 × 50000 × 768 ≈ 77M parameters — more than half the network — so tying is not a tweak, it is a redesign of where the model’s capacity lives.

For the CPU-SLM practitioner the calculus is even sharper. Every parameter you do not store is memory you do not load and bandwidth you do not spend at inference, and on a CPU, memory bandwidth is usually the binding constraint. Cutting a 16M-parameter tensor out of a 42M model is a ~28% reduction in the weight footprint you stream per forward pass. Tying therefore improves quality, shrinks the file, and speeds up memory-bound inference simultaneously — the rare optimization with no obvious tradeoff for models of this size.

The softmax-bottleneck caveat

Tying is not free of theoretical limits, and the sharpest one is the softmax bottleneck (Yang et al., 2018). The final layer computes logits = h · E^T, and because h ∈ ℝ^d while there are V tokens, the matrix of log-probabilities the model can produce across all contexts has rank at most d. But the ‘true’ log-probability matrix of natural language appears to be high-rank — language is context-rich enough that a d-dimensional bottleneck genuinely limits expressiveness. When d is small (the SLM regime), this ceiling is real: no matter how you train, a single d-wide projection cannot represent every conditional distribution the data demands.

Tying interacts with this by coupling the two roles: the same vector must be a good input feature and a good output classifier direction, one more constraint on an already rank-limited head. The fix Yang et al. proposed is Mixture of Softmaxes — compute several context-dependent projections and mix them, breaking the rank ceiling. The practical takeaway: tying is nearly always worth it, but if you have squeezed d very small and hit a perplexity floor you cannot train past, the shared low-rank head may be the wall you are hitting.

Advertisement

When dimensions do not match: the projection layer

Plain tying assumes the input embedding width equals the model’s hidden size, so that E^T maps ℝ^d straight to V logits. Sometimes they differ on purpose. If token embeddings live in a smaller space ℝ^e with e < d — a factorized embedding — you keep a small [V, e] table and a shared [e, d] projection, so the vocabulary cost drops from V·d to V·e + e·d, a large win when V is huge. The same [V, e] table can still be tied to the output side.

ALBERT uses exactly this decomposition to decouple the vocabulary embedding size from the hidden size, arguing the two need not be equal: the embedding captures context-free token identity while the hidden state carries context, so a thin embedding feeding a wider model is a natural split. In neural machine translation, Press & Wolf go further with three-way weight tying — the source embedding, the target embedding, and the target output projection all share one matrix — which works when source and target share a subword vocabulary and cuts three tensors down to one.

Cross-layer sharing: ALBERT

Tied embeddings share weights across the two ends of the model; the other great parameter-sharing idea shares weights across depth. ALBERT (Lan et al., 2019) makes every transformer layer reuse the same parameters — one block of attention and FFN weights applied L times instead of L distinct blocks. Where a standard 12-layer encoder stores twelve independent ~12d^2 parameter sets, ALBERT stores one and iterates it, collapsing the body’s 12·L·d^2 cost to roughly 12·d^2 regardless of depth.

Combined with its factorized embedding, ALBERT reached BERT-large quality with an order of magnitude fewer parameters. The catch is that cross-layer sharing saves parameters and memory but not compute: you still run L forward passes through the shared block, so FLOPs and latency are unchanged — it shrinks the model on disk and in RAM, not the work per token. That is precisely the opposite tradeoff from a technique like distillation, and it makes cross-layer sharing attractive when your constraint is memory rather than throughput.

A brief tour of other sharing schemes

Parameter sharing is a design axis, not a single trick, and several schemes sit alongside embedding tying and ALBERT-style depth sharing. The Universal Transformer takes cross-layer sharing to its logical end — a single recurrent-in-depth block applied a dynamic number of times, so depth becomes iteration count. Grouped-query and multi-query attention share key and value projections across attention heads, shrinking not the parameter count so much as the KV cache that dominates decode memory — a sharing scheme aimed at inference memory rather than model size.

Going the other direction, Mixture-of-Experts deliberately un-shares the FFN into many experts and routes each token to a few, trading parameter count up for compute held roughly flat — the mirror image of ALBERT’s bargain. And adapter and LoRA methods share the entire frozen backbone across tasks, learning only tiny task-specific deltas. The through-line: every sharing decision picks which of parameters, compute, or memory to spend and which to conserve, and tied embeddings are the cleanest example because they spend almost nothing to save a lot.

Gradient flow and training dynamics

When E is used at both ends, its gradient is a sum of two contributions per step: the input-side gradient (from how the embedded token influenced everything downstream) and the output-side gradient (from the cross-entropy loss on the logits it produced). For a target token, the output-side term pushes the hidden state and that token’s row toward alignment; for a context token, the input-side term shapes the same row as a feature. These signals are generally compatible — both want a coherent representation of the token — which is why training stays stable and even improves.

One practical wrinkle is the scale of the shared matrix. Input embeddings and output projections sometimes prefer different magnitudes; the standard remedy is to scale the embedding lookup (for instance by √d, as in the original Transformer) so the same weights serve both roles at sensible scales. Another is the softmax temperature or a learned scalar on the logits. These are small knobs, but they matter: get the scaling wrong and a tied model can train sluggishly even though the parameter sharing itself is beneficial. When tying ‘does not help,’ a scale mismatch is a more common culprit than the tie being wrong.

Practical guidance for CPU-bound small models

For the small-model, CPU-inference setting, tied embeddings should be your default and near-automatic choice. The reasoning stacks up cleanly: the embedding is a large fraction of a small model, so tying delivers a big proportional parameter cut; that cut is directly a memory-footprint and memory-bandwidth win, which is what bounds CPU inference; and the change typically improves perplexity rather than costing it. Three benefits, no meaningful cost — it is close to a free lunch.

Combine it deliberately with the vocabulary size. Because the embedding cost is V·d, a leaner tokenizer (a smaller V) compounds with tying to shrink the dominant tensor further, though you must balance that against longer sequences from coarser tokens. If you are memory-constrained beyond what tying alone buys, reach next for a factorized embedding to break the V·d term into V·e + e·d, and consider cross-layer sharing if depth is inflating your body. Sequence these in order of payoff: tie first (free and beneficial), factorize second (cheap, mild quality cost), share layers last (biggest savings, real quality tradeoff).

Pitfalls and where tying stops paying

A few traps recur. First, tying saves parameters, not compute: the final matrix multiply h · E^T is still O(V·d) FLOPs whether or not the weights are shared, so a large vocabulary keeps the LM head expensive at inference no matter what — the win is memory, not speed of that layer. Second, the benefit shrinks toward irrelevance as models grow; on a multi-billion-parameter model with a small vocabulary fraction, tying is a rounding error and the coupling constraint may even cost a hair of quality worth measuring.

Third, watch the softmax bottleneck if you have pushed d very low — a shared, low-rank head is a real expressiveness ceiling, and no amount of training climbs past it. Fourth, verify your framework is truly sharing the tensor and not silently duplicating it; a ‘tied’ model that still stores two copies gets the regularization but not the memory savings. Handled with these caveats in mind, weight tying remains the highest-return-per-line change available to a small-model builder: one matrix, two jobs, less memory, and usually a better model.

Weight tying replaces the input embedding and the output projection — two V × d matrices doing mirror-image jobs on the same vocabulary — with a single shared matrix, deleting exactly one V × d parameter block. It is sound because the input and output token spaces coincide: a token’s meaning is one geometry used symmetrically. On small models the embedding is a large slice of the network, so tying can cut a quarter or more of the parameters while improving perplexity, acting as a regularizer that doubles the learning signal on rare tokens — a genuine win for memory-bandwidth-bound CPU inference. Its limit is the softmax bottleneck, where a shared low-rank head caps expressiveness at very small d. The idea generalizes: factorized embeddings split the V × d cost, and cross-layer schemes like ALBERT share weights across depth to save memory (though not compute). Tie first — it is the cleanest, highest-return sharing decision a small-model builder can make.