Untied embeddings, contrasted against weight tying

The problem, from first principles

A language model needs TWO conceptually different matrices shaped (vocab_size, d_model): an INPUT embedding table ([04-seq-modeling/02-embeddings/01-token-embedding-lookup], looking a token id UP to get a vector) and an OUTPUT projection (turning the model's final hidden state into a score, a "logit," for every word in the vocabulary, so a softmax over those logits gives next-token probabilities). Press & Wolf (2016) observed something genuinely useful: these two matrices, despite serving different DIRECTIONS of the same lookup (id-to-vector vs. vector-to-id-scores), can literally be the SAME matrix, "tied" together, rather than two entirely separate, independently-learned parameter sets, without hurting model quality, and while meaningfully REDUCING the total parameter count (a vocab_size * d_model-sized matrix, often a substantial fraction of a smaller model's total parameters, gets counted only once instead of twice).

The alternative, UNTIED embeddings, uses two entirely separate matrices: the model learns its own best input representation independently from its own best output-scoring function, at the cost of the extra parameters. Whether tying helps or hurts is itself an empirical, model-size-dependent question, smaller models often benefit measurably from tying (proportionally, the savings matter more, and sharing information between the "read" and "write" directions of the vocabulary can act as a useful regularizer), while very large models are more often trained untied, having plenty of capacity to spare and, in some cases, benefiting from letting the two roles specialize independently.

From theory to code

Implement output_projection_tied(hidden_states, embedding_table), projecting hidden_states to vocabulary logits by reusing an EXISTING embedding table (transposed), and output_projection_untied(hidden_states, output_weight), the same projection using a SEPARATE, independently-learned matrix.

Constraints

• Both functions compute hidden_states @ matrix.T, matrix being embedding_table in the tied case and output_weight in the untied case; both matrices are shaped (vocab_size, d_model).
• The TIED version must genuinely reuse the SAME matrix object passed in as the embedding table, not a copy: if that table changes (as it would during training, since the SAME parameter also receives gradient updates from the input-embedding side), the output projection changes too, automatically.
• The UNTIED version's output_weight is entirely independent: changes to some other, unrelated embedding table must have no effect on it whatsoever.

Hints