Introduction

Weight tying in a language model means reusing the token embedding matrix as the output projection matrix before softmax. This reduces parameters and often improves generalization, but it only works cleanly when the hidden-state width matches the embedding width or when you insert a projection layer to make them match.

Why Weight Tying Makes Sense

In a basic language model, you often have two matrices:

• an embedding table of shape vocab_size x embed_dim
• an output projection of shape embed_dim x vocab_size

Those shapes are transposes of each other. Instead of learning both separately, you can learn one embedding matrix and reuse its transpose to compute logits over the vocabulary.

That gives you:

• fewer trainable parameters
• tighter coupling between input and output token representations
• an architecture used in many practical language-model designs

The core constraint is dimensional consistency. The model state you feed into the output projection must have width embed_dim.

A Custom Tied Output Layer

Keras does not provide a single built-in flag for this exact setup, so the usual solution is a custom layer that receives an Embedding layer reference and reuses its weights.

1import keras
2from keras import ops
3
4
5class TiedOutputProjection(keras.layers.Layer):
6    def __init__(self, embedding_layer, **kwargs):
7        super().__init__(**kwargs)
8        self.embedding_layer = embedding_layer
9
10    def build(self, input_shape):
11        vocab_size = self.embedding_layer.input_dim
12        self.bias = self.add_weight(
13            name="bias",
14            shape=(vocab_size,),
15            initializer="zeros",
16            trainable=True,
17        )
18
19    def call(self, inputs):
20        embedding_matrix = self.embedding_layer.embeddings
21        logits = ops.matmul(inputs, ops.transpose(embedding_matrix))
22        return logits + self.bias

This layer creates only the bias term. The projection weights are taken directly from the embedding matrix.

Using It in a Simple Next-Token Model

Here is a small runnable example for sequence-level prediction:

1import keras
2
3vocab_size = 100
4embed_dim = 32
5sequence_length = 10
6
7tokens = keras.Input(shape=(sequence_length,), dtype="int32")
8
9embedding = keras.layers.Embedding(vocab_size, embed_dim, name="token_embedding")
10x = embedding(tokens)
11x = keras.layers.LSTM(embed_dim)(x)
12logits = TiedOutputProjection(embedding)(x)
13outputs = keras.layers.Activation("softmax")(logits)
14
15model = keras.Model(tokens, outputs)
16model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
17model.summary()

The important part is LSTM(embed_dim). Its output width matches the embedding width, so multiplying by the transposed embedding table is valid.

If your hidden state has a different width, add a learned projection first:

x = keras.layers.Dense(embed_dim)(x)

That preserves weight tying while letting the rest of the model use a different hidden size internally.

Time-Step Output Versus Single Output

Some models predict one token per sequence, while others predict a token at every time step. Weight tying works in both cases, but the tensor shapes differ.

For per-time-step prediction, keep the sequence dimension and apply the tied projection over the last axis. A simple way is to let the recurrent layer return sequences:

x = keras.layers.LSTM(embed_dim, return_sequences=True)(x)
logits = TiedOutputProjection(embedding)(x)

The matmul still works because it multiplies the last dimension of x by the transposed embedding matrix, producing logits of shape batch by time by vocabulary.

So the underlying idea does not change. Only the tensor rank changes.

Why a Normal Dense Layer Is Not Enough

A standard Dense(vocab_size) layer creates its own weight matrix. Even if its shape matches the transpose of the embedding table, the weights are not tied unless the layer explicitly reuses the embedding variable.

This is an important distinction. Similar shape is not shared weight.

If you want true weight tying, the output layer must read from the embedding layer's variable directly, as the custom layer above does.

Training Considerations

When weights are tied, gradient updates from the output projection also update the embedding table. That is the whole point: the same parameters serve both input and output roles.

This can improve parameter efficiency, but it also means the embedding matrix is carrying two modeling responsibilities. In practice that is often beneficial, but it is a real architectural choice rather than a free trick.

Also remember that the output bias is still typically learned separately. Weight tying usually replaces the output projection matrix, not the bias vector.

Common Pitfalls

The most common mistake is trying to tie weights when the hidden dimension and embedding dimension do not match. If the shapes are incompatible, add a projection layer first.

Another issue is assuming a Dense(vocab_size) layer automatically ties because it has the right shape. It does not; it owns separate weights unless you explicitly reuse the embedding variable.

People also often forget the output bias term. Tied projection weights and output bias are separate design elements.

Finally, be careful with serialization and custom objects. If you save and reload the model, make sure the custom layer is registered or provided when loading.

Summary

• Weight tying reuses the embedding matrix as the output projection before softmax.
• In Keras, the usual implementation is a custom layer that multiplies by the transpose of the embedding weights.
• The hidden width must match the embedding width, or you need a projection layer.
• A normal dense output layer does not create shared weights automatically.
• Tied weights reduce parameters and are common in language-model architectures.