Tie Input and Output Embeddings in Language Models

A language model needs token representations in two places. At the input, it converts token IDs into vectors. At the output, it converts a hidden vector into one score for every token in the vocabulary. A straightforward design gives these two operations separate parameter matrices, even though both matrices associate vocabulary items with vectors.

Weight tying removes that duplication by using the same matrix for both roles. The model looks up input embeddings from the matrix and later uses its transpose to produce output logits. This can remove a large block of parameters, but it also couples two parts of the model that would otherwise learn independently.

This article builds a practical mental model for tied embeddings, shows exactly what is shared, explains the parameter and memory implications, and covers the cases where tying is useful or structurally inappropriate.

Start with two vocabulary-sized matrices

Assume a language model has:

vocabulary size V = 50,000
hidden size     d = 4,096

An input embedding table has shape:

E_in: V x d

When token t enters the model, its initial vector is the corresponding row:

x = E_in[t]

After the transformer or other sequence model processes the input, suppose it produces a hidden vector h with d components. To predict the next token, an output projection can map h to V logits:

W_out: V x d
logits = W_out h

Ignoring an optional output bias, each row of W_out acts like a vector associated with one possible output token. The dot product between that row and h becomes the token’s logit.

With separate input and output matrices, these two vocabulary-sized parameter blocks contain:

2 x V x d

parameters.

For the example above, one matrix contains:

50,000 x 4,096 = 204,800,000 parameters

Two such matrices contain 409.6 million parameters. That makes the duplication worth examining, especially when the vocabulary and hidden size are large.

Reuse one matrix in both directions

With weight tying, the model keeps one matrix:

E: V x d

Input lookup still selects a row:

x = E[t]

Output prediction uses the same rows as output vectors:

logits = E h

Some notation writes the embedding matrix as d x V instead. In that convention, output prediction uses E^T h. The transpose is a matter of how the matrix is stored and written; the important property is that the input embedding parameters and output projection parameters are the same parameters.

A small example makes the relationship clearer. Suppose the vocabulary is:

0: cat
1: dog
2: runs

and each token has a two-dimensional shared vector:

cat  -> [ 0.8,  0.2]
dog  -> [ 0.7,  0.3]
runs -> [-0.1,  0.9]

The row for dog is used when dog appears as an input token. The same row is also used when scoring whether dog should be the next output token. If the final hidden state is:

h = [0.6, 0.4]

then the simplified output logits are dot products:

cat:  0.8*0.6 + 0.2*0.4 = 0.56
dog:  0.7*0.6 + 0.3*0.4 = 0.54
runs: -0.1*0.6 + 0.9*0.4 = 0.30

A softmax can then turn these logits into a probability distribution. The numbers are only a teaching example; real embedding dimensions are much larger, and logits may include other architectural details such as normalization, scaling, or bias.

Understand what tying changes during training

Sharing parameters does more than save space. It changes how those parameters receive gradients.

With separate matrices, an input embedding row is updated through the role that token representations play inside the network. The output matrix is updated through the prediction loss at the output. They can specialize independently.

With a tied matrix, both computational paths update the same parameter tensor:

input use  ----\
               -> shared E -> one parameter update
output use ----/

Conceptually, the gradient on E contains contributions from both roles. This does not mean the two gradient contributions are equal, nor that every token row receives both kinds of signal on every training example. It means optimization cannot change the input and output versions independently because there is only one version.

That coupling is the central modeling trade-off. It imposes the assumption that a shared token-vector space is suitable for both representing observed tokens and scoring predicted tokens.

Calculate the parameter saving correctly

If both untied matrices have shape V x d, tying removes one V x d parameter matrix. The parameter reduction is therefore:

saved parameters = V x d

For V = 50,000 and d = 4,096:

saved parameters = 204,800,000

If those parameters were stored as 16-bit values, the raw weight storage for that duplicate matrix would be about:

204,800,000 x 2 bytes = 409,600,000 bytes

That is about 409.6 MB in decimal units, before considering allocator overhead, serialization metadata, quantization metadata, or other implementation details.

Training can save more memory than the raw weight size alone when the removed parameter tensor would otherwise require gradients and optimizer state. The exact saving depends on the optimizer, precision scheme, sharding strategy, and framework implementation, so it should be measured rather than inferred from parameter count alone.

Tying also does not imply that inference becomes proportionally faster. The model still has to compute vocabulary logits unless another technique changes that operation. Weight tying primarily changes parameterization and storage; runtime effects depend on the serving implementation and hardware.

Make sure the shapes and vocabularies are compatible

The simplest form of weight tying requires the input embedding dimension to match the dimension consumed by the output projection. If the final hidden state has size d and the shared embedding table is V x d, the multiplication is directly compatible.

If those dimensions differ, a model needs an additional transformation or cannot use direct tying. For example:

embedding size = 1,024
final hidden size = 4,096

A single V x 1,024 matrix cannot directly multiply a 4,096-dimensional hidden state. An architecture could project the hidden state to 1,024 dimensions first, but that is an architectural choice, not automatic weight tying.

Vocabulary compatibility matters too. In a decoder-only model, the tokens consumed as input and predicted as output usually come from the same tokenizer vocabulary, which makes direct tying structurally natural.

Encoder-decoder models require more care. The source-side tokenizer, target-side tokenizer, encoder embeddings, decoder embeddings, and output projection are not necessarily identical. Sharing is possible only among components whose vocabulary definitions and dimensions are compatible. A model using different source and target vocabularies cannot simply treat their embedding rows as corresponding tokens.

Do not confuse parameter sharing with copying

A common implementation mistake is to initialize two matrices with equal values and call them tied.

These are different designs:

copy once:
E_in  = initial values
W_out = copy(initial values)

true tying:
E_in  and W_out refer to the same trainable parameters

After the first optimizer update, copied matrices can diverge. Truly tied weights remain identical because there is only one underlying parameter set.

This distinction also matters when loading checkpoints. A framework may represent tied tensors as aliases, may reconstruct the tie from model configuration, or may serialize values in an implementation-specific way. Do not assume that two similarly named tensors in a checkpoint prove that the runtime model is untied, or that equal initial values prove it is tied. Check the model’s documented architecture or inspect whether the parameters actually share storage in the framework you use.

Treat vocabulary changes as a shared operation

Adding or removing tokens changes the number of rows in an embedding table. In a tied model, that same row count defines the output vocabulary.

Suppose a model adds a token called <support_ticket>. The model needs a new input representation for that token and a corresponding output row if it is allowed to generate the token. With tied weights, those are the same new row.

This has two practical consequences.

First, resizing only one conceptual side is not meaningful for a directly tied matrix. The input vocabulary and output vocabulary remain coupled.

Second, initialization of new rows deserves attention. A newly added token does not become useful merely because a vector exists for it. Its representation must acquire appropriate behavior through training or another justified initialization strategy. Weight tying does not solve that learning problem.

Know when separate matrices are more appropriate

Weight tying is attractive when input and output tokens share a vocabulary and dimensional space, and reducing parameter count is valuable. It is not a universal requirement.

Separate matrices can be appropriate when the architecture intentionally gives input representation and output discrimination different parameter spaces. They are also necessary when shapes or vocabularies do not line up without extra transformations.

There is another practical reason not to change tying casually: compatibility. If you start from a pretrained model, whether its embeddings are tied is part of its architecture. Turning an untied checkpoint into a tied model by assigning one matrix to the other changes the model’s function unless the matrices already happen to be identical. Untying a pretrained tied model creates new degrees of freedom and changes subsequent optimization behavior. Either change may be a valid experiment, but it is not a neutral memory optimization.

For application developers who only call a hosted language-model API, weight tying is usually not a controllable setting. It is an internal model architecture choice. It becomes relevant when selecting, training, fine-tuning, converting, or serving models whose parameter structure you manage.

Avoid common mental-model mistakes

The most useful checks are simple:

• Tying does not merge tokens. Every vocabulary item still has its own row; the same table is merely reused for input and output roles.
• Tying does not eliminate the output softmax. The model still needs logits over the output vocabulary when full next-token probabilities are required.
• Tying is not the same as freezing. Shared weights remain trainable unless explicitly frozen.
• Tying is not the same as quantization. Tying reduces the number of distinct parameters; quantization changes how parameter values are represented.
• Parameter savings do not guarantee equal latency savings. Matrix multiplication, memory movement, batching, vocabulary size, and serving kernels still determine runtime behavior.

These distinctions help separate an architectural parameter-sharing technique from several optimizations that solve different problems.

Conclusion

Weight tying starts from a simple observation: language models often maintain one vocabulary-sized matrix to read tokens and another to score tokens. When the vocabulary and dimensions are compatible, those roles can share one matrix.

The immediate result is easy to calculate: direct tying removes V x d distinct parameters. The deeper effect is that input representation and output prediction now optimize the same token vectors. That coupling can be a useful architectural constraint, but it should be treated as part of the model design rather than as a free runtime optimization.

When evaluating a model, ask three questions: do the input and output vocabularies match, do their dimensions permit direct sharing, and was the model designed or trained with that sharing? If all three answers support tying, the technique is straightforward. If not, forcing the matrices together changes more than memory usage.