Why Use Weight Tying?

1. What is an embedding matrix?

Suppose a vocabulary that contains just three words.

This entire table is the embedding matrix, E.

E = ⎡ 0.2  0.8 ⎤
    ⎢ 0.5  0.3 ⎥
    ⎣ 0.9  0.1 ⎦

In general, E has shape V × d, which is vocabulary size x embedding dimension.

Given token ID 0, the model retrieves the first row, [0.2, 0.8]. This operation is a "lookup".

The embedding matrix is kinda like a dictionary, but not exactly because the values don't stay fixed forever.

• The mapping between token IDs and words is generally fixed.
• The numbers in each embedding vector are updated during training.
• During ordinary inference after training, those numbers stay fixed.

The embedding matrix is therefore part of the model’s learned weights.

2. What does the LM head actually compute?

The LM head is the model’s final layer, that turns the transformer’s output into a score for every token in the vocabulary, so the model can pick the next token.

To illustrate the processing steps, imagine giving a model the text “I like”.

The transformer produces a hidden vector h that reflects the preceding context. The LM head uses it to score the possible next tokens. Using the three example tokens above, it assigns a score to cat, dog, and apple.

In a model that uses weight tying, this means computing the dot product of the hidden vector with each row of the embedding matrix.

logits = h × ET

h: 1 × d, E^T: d × V, result: 1 × V

Here are h and E again:

h = [ 0.3  0.7 ]

E = ⎡ 0.2  0.8 ⎤  cat
    ⎢ 0.5  0.3 ⎥  dog
    ⎣ 0.9  0.1 ⎦  apple

For this example, the scores are:

cat has the highest score. Applying softmax turns these scores into a probability distribution.

The LM head computes a score for every token in the vocabulary by taking the dot product of the hidden vector with that token’s weight vector. In practice, this is done efficiently through matrix multiplication, rather than by iterating over tokens one at a time.

3. Why tie the weights?

A model can predict the next token without weight tying. Weight tying means reusing the embedding matrix E as the LM head’s weights. Without it, the input embedding and the LM head are two separate layers: E looks up a token’s input vector, while the LM head uses its own matrix, W, to score every token from h.

The two roles answer different questions:

• Input embedding: “Which vector represents cat?”
• LM head: “What score should cat receive as the next token in this context?”

Both matrices have one row per token in the same vocabulary, so a model can share them. However, because E and W do different jobs, the best values for the two are not necessarily the same. E’s cat row is tuned to be a useful input for the transformer, while W’s cat row is tuned to be the direction h should point when cat comes next. That is why some models do not use weight tying.

So why tie them? There are two main reasons:

• Fewer parameters. With a large vocabulary, V × d is big. GPT-2 small has 124M parameters with tying; a separate W would add about 38.6M more (50,257 × 768).
• Often no loss in quality. In practice, sharing usually works about as well, and sometimes better. Every row keeps learning through the LM head, even for rare tokens that seldom appear in the input.

The savings matter most in small models, where the embedding is a large share of all parameters. Very large models often skip tying, since a second V × d matrix is a small fraction of their size.

When are the weights updated?

During training, each step computes the loss and updates the weights through backpropagation. Which rows get updated depends on the setup:

During inference, none of these weights change.

Important: Weight tying does not happen automatically just because two matrices have the same shape. The model’s designer explicitly makes the two modules share the same weight parameters.

4. Does the LM head run for every token?

It helps to distinguish training from inference.

Training

In standard next-token prediction training, the model computes next-token scores at every position in a sequence, so each position gets a score for every token in the vocabulary. For a sequence with three tokens, the hidden states and logits have the following shapes:

H ∈ ℝ3 × d
Logits = H × ET ∈ ℝ3 × V

Each of the three rows of H is one position’s hidden vector. Multiplying by E^T turns each row into V scores, giving a 3 × V table. So yes: during training, the LM head runs at every position, because every position’s prediction is compared with the actual next token and adds to the loss.

Inference

During autoregressive generation, the model typically uses only the hidden vector at the final position. It scores the vocabulary from that vector, picks one next token, appends it to the input, and repeats.

h ∈ ℝ1 × d
Logits = h × ET ∈ ℝ1 × V

Here h is just the final position’s hidden vector, so the result is a single row of V scores. So no: during generation, the LM head runs only once per new token, at the final position. Earlier positions still pass through the transformer, but their next-token scores are not needed.

Key takeaways

• Input embedding: Looks up the vector corresponding to a token ID in the embedding matrix.
• Transformer: Transforms the input vectors into hidden vectors that reflect the context.
• LM head: Converts a hidden vector into token scores across the entire vocabulary.
• Weight tying: Lets the input embedding and the LM head share one weight matrix while performing different operations.

Even with weight tying, both the input lookup and the LM head’s score computation are necessary. Weight tying removes duplicated parameters; the LM head’s computation still takes place.