Tied Input/Output Embeddings

A token embedding table gives each vocabulary item a vector. A language model can reuse that same table to turn its final hidden states into a score for each vocabulary item. This sharing is called weight tying.

The score for a token is the dot product of its embedding vector and the hidden state. A larger score means a larger probability if softmax is applied later. This exercise returns the scores, called logits, without softmax.

Your task

Implement tied_lm_head(hidden, embedding_matrix):

hidden is a floating-point tensor with shape (..., d_model). It can be a single vector, a batch of vectors, or a batch of token sequences.
embedding_matrix is the shared table with shape (vocab_size, d_model).
Return hidden @ embedding_matrix.T, with shape (..., vocab_size).

Both inputs have the same floating-point dtype. Preserve all leading dimensions of hidden; only its last dimension changes. The table is shared across every batch and sequence position.

logits[..., v] = sum(hidden[..., d] * embedding_matrix[v, d] over d)

Do not add a bias, normalize the vectors, scale the dot products or apply softmax. The supplied values are ready to use. Logits may be negative and do not have to sum to one.

For example, with hidden=[2, -1] and table rows [1, 0], [0, 1] and [1, 1], return [2, -1, 1]. These are three scores for one hidden vector, so the output shape is (3,), not (1, 3).

What is shared?

An untied model has separate parameters for its input table and output weights. Tying them removes one independently stored matrix of vocab_size * d_model parameters. During training, both uses contribute gradients to the same table. This function implements the output projection only; model construction is responsible for passing the shared parameter rather than a separate copy. Weight tying is an architectural choice, not a requirement for language models.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Look up a learned vector for each token

An embedding table E has shape (vocabulary_size, embedding_dim). Integer token ID i selects row E_i. IDs are labels for rows, not measurements: token 100 does not have twice the meaning of token 50.

A lookup is equivalent to multiplying a one-hot token vector by E, but avoids constructing the mostly zero vector. Looking up IDs [2, 2, 5] returns rows [E_2, E_2, E_5]. Both uses of row 2 contribute gradients to that row during training, subject to any scaling or special padding behavior in the task.

Averaging token embeddings gives one fixed-width vector for a sequence. It loses order information: the average for tokens [2, 5] is the same as for [5, 2].

Reuse the embedding weights to predict tokens

Weight tying uses the same table for input embeddings and output scores. For a hidden vector h with width embedding_dim, the product hE^T gives one score, or logit, for every vocabulary entry.

An input lookup selects one row. The output projection compares the hidden vector with every row. Sharing the table saves a separate output weight matrix, but means both uses must contribute gradients to the same parameter. Two copies that are updated independently are no longer tied.

Check that token IDs are in range and have the required integer dtype. When pooling padded sequences, exclude padding from both the sum and the count. Even a zero padding vector lowers an ordinary average if it is included in the denominator. Follow the specified behavior for a sequence with no valid tokens.

Detecting runtime…

Tied Input/Output Embeddings

Your task

What is shared?

Look up a learned vector for each token

Reuse the embedding weights to predict tokens

Your code, saved here

Stuck?

JAX reference solution