Weight Tying and Parameter Sharing

Parameters are the expensive resource. Using the same ones in more than one place is the most direct way to spend fewer of them, and whether it costs anything depends entirely on whether the two uses want the same function. Sometimes they provably do, sometimes they clearly do not, and the interesting cases are in between.

The embedding case, where it is almost free

A language model has an input embedding mapping token IDs to vectors and an output projection mapping the final hidden state to logits over the same vocabulary. Both are \(V \times d\) matrices over the same vocabulary, and tying them, using one matrix for both, is standard.

The justification is that the output projection computes a dot product between the hidden state and each token's row, so those rows already function as token representations. Requiring them to be the same representations the input uses is a compatible constraint, and empirically it improves perplexity on small models rather than merely saving memory, because the shared matrix receives gradient from both paths and is better estimated.

The saving is large where the vocabulary dominates. For a model with \(d = 768\) and \(V = 128{,}000\), each embedding matrix is 98M parameters, so tying saves 98M from a model that might total 350M. For a 70B model with the same vocabulary the same 98M is noise, which is why tying is near-universal in small models and often dropped in large ones, where the extra freedom is cheap and slightly better.

Cross-layer sharing, where it costs

ALBERT shares one transformer block's parameters across all layers (Lan et al., 2020, arXiv:1909.11942), turning the stack into a recurrent application of one function. Parameters fall by the layer count. Compute does not fall at all, because every layer still executes.

That combination is the crux. Cross-layer sharing reduces memory and leaves latency and FLOPs unchanged, so it helps only when parameter storage is the binding constraint, as on a memory-limited device. It costs capability, because layers in a trained transformer do measurably different things: early layers do syntactic and local work, middle layers do the bulk of factual and relational processing, late layers prepare the output distribution. Forcing one function to serve all three roles is a real restriction, and it shows up as a quality gap that grows with depth.

Partial schemes are the useful middle ground: share within groups of adjacent layers, which are more similar to each other, or share only the attention parameters while keeping per-layer MLPs, since MLPs carry most of the layer-specific knowledge.

When it breaks

Tied embeddings interact with normalisation and scale. The input embedding wants a scale suited to the residual stream at layer zero; the output projection wants a scale suited to producing well-calibrated logits. Tied models usually need an explicit scaling factor on one side, and getting it wrong produces a model that trains but has poorly calibrated output.

Tying constrains vocabulary changes. Extending a vocabulary post-training means growing both matrices consistently, and any adaptation that would benefit input and output differently, such as adding domain tokens that should be recognised but rarely generated, cannot be expressed.

Shared layers make interpretability and surgery harder. Layer pruning, early exit, and per-layer quantisation all assume layers are separable objects. With shared weights, changing one layer changes all of them, so most of the model-surgery toolkit does not apply.

Recurrent-depth models are a different idea wearing similar clothes. Applying one block a variable number of times, adaptively, is not primarily a parameter-saving technique: it trades compute for capability at inference, spending more passes on harder inputs. Evaluating it as parameter sharing misses the point, and evaluating parameter sharing as adaptive compute credits it with a benefit it does not have.