The pragmatic tradeoff of tied embeddings

Tied embeddings are one such tradeoff.

It comes from a simple observation: "we have a 617 million parameter embedding matrix^* on both sides of our nn. why not just make them the same matrix?" ^*(in GPT 3)

In other words: since the embedding matrix encodes semantic meaning in words, it serves roughly the same purpose for both input and output predictions.

Transformers have two matrices that deal with tokens, an

• embedding matrix $W_E$: converts tokens to vectors (input), and an
• unembedding matrix $W_U$: converts residual stream to logits (output).

With tied embeddings, we use the same weights: $W_U = W_E^T$. It the ML brain because it appears to be an elegant way to reduce parameters and add symmetry.

Let's make this concrete with a toy example vocabulary:

Tokens are represented as one-hot vectors, where a single 1 indicates which token we're referring to:

The embedding matrix $W_E$ has shape (vocab_size, embedding_dim). Each row is the learned embedding for one token (these are made up):

To get a token's embedding, we multiply its one-hot vector by $W_E$. This just selects the corresponding row:

At the output, the unembedding matrix converts a residual stream vector to logits. With tied embeddings, this is $W_E^T$:

Each logit is the dot product of the residual with a token's embedding. Higher dot product means the model thinks that token is more likely:

In a transformer, information flows through attention and MLP layers. But residual connections mean that there's always a direct linear path from input to output:

Even in a deep model, this path exists. If we ignore all the attention/MLP layers, the model is simply computing:

The matrix $W_E \cdot W_U$ is a vocab_size × vocab_size matrix. Entry [i, j] tells us: given input token i, what's the logit for output token j?

"Given the current token, what's the most likely next token?"

This is called bigram statistics. It's the conditional probabilities $P(\text{next\_token} \mid \text{current\_token})$. A bigram is any two consecutive tokens in text.

For our vocabulary, realistic bigram probabilities look like this:

Notice the critical observation: bigrams are asymmetric. The order matters:

Visualizing this asymmetry. Notice how the matrix is NOT symmetric across the diagonal:

Here's the critical issue. With tied embeddings, $W_U = W_E^T$, so the direct path matrix becomes:

And $W_E \cdot W_E^T$ is always symmetric. Let's see why.

This follows from basic linear algebra. The (i, j) entry of $W_E \cdot W_E^T$ is:

And dot products are commutative: $a \cdot b = b \cdot a$. Therefore:

No matter what values $W_E$ has, $W_E \cdot W_E^T$ is always symmetric.

This is bad. We lose expressivity because these serve fundamentally different purposes. Tying them forces a single representation to do both jobs.

With untied embeddings, $W_U$ is a separate learnable matrix. Now the direct path is

where $W_E$ and $W_U$ are independent. This product can be any matrix, including asymmetric ones.

With untied embeddings, we can solve for a $W_U$ that approximates the bigram probabilities:

The untied result can now represent asymmetric relationships:

Smaller models still use them because the lack of expressivity means that the nonlinearity added through attention and MLP layers is enough to not bump against the asymmetry constraints.

Some smaller models also do this by directly adding a MLP layer before the residual stream which breaks the symmetry as well.

MLP₀ learns a transformation M, making the effective path $W_E \cdot M \cdot W_E^T$, which CAN be asymmetric.

Let's see how adding an MLP transformation breaks the symmetry:

Explore the parameter savings from tied embeddings:

Tied embeddings are a practical tradeoff, not a principled design choice. They save parameters and memory, which matters for smaller models where early layers can compensate for the symmetry constraint. But for large models chasing maximum performance, the math is clear: language is asymmetric, and untied embeddings can represent that.

• Tied: GPT-2, Gemma, original Transformer
• Untied: LLaMA 1/2/3, Mistral
• Unknown: GPT-3/4, Claude (not publicly disclosed)