Weight Tying › The Two Embedding Matrices

Every language model that processes text through a transformer architecture faces the same fundamental challenge at its boundaries. Transformers operate on continuous vectors, but text consists of discrete tokens. Bridging this gap requires learned mappings in both directions: from tokens to vectors at the input, and from vectors back to tokens at the output. These two mappings are implemented as the input embedding matrix and the output projection matrix, and understanding them precisely is the first step toward appreciating why tying them together makes sense.

Think of the transformer as a black box that maps sequences of dense vectors into new sequences of dense vectors. Before the transformer can do its work, you need to convert your token indices into vectors it can process. After the transformer produces its output, you need to convert those output vectors back into probabilities over the vocabulary. The two embedding matrices handle these conversion tasks.

Both matrices have the same shape, both operate over the same vocabulary, and both are involved at every single forward pass through the model. Yet in naive implementations, they are treated as entirely separate learned objects. Weight tying questions whether this separation is necessary or even desirable.

The input embedding matrix $\mathbf{E} \in \mathbb{R}^{V \times d}$ is one of the most fundamental components in any neural language model. It stores a learned dense vector representation for each token in the vocabulary. When the model sees token index $t$, it retrieves the corresponding row from this matrix:

$$\mathbf{h}_{\text{in}} = \mathbf{E}[t]$$

where:

• $\mathbf{E}$: the input embedding matrix of shape $V \times d$, where every row contains the learned representation of one vocabulary token
• $t$: the input token index, an integer from 0 to $V - 1$
• $\mathbf{h}_{\text{in}}$: the resulting embedding vector of dimension $d$, which enters the transformer's processing pipeline
• $V$: the vocabulary size (50,257 for GPT-2, 32,000 for LLaMA-2, 128,000 for newer models)
• $d$: the embedding dimension, also called the model dimension or hidden dimension

This operation is purely a table lookup. Token 42 retrieves row 42 from the matrix. There is no computation beyond the indexing, which is one reason why embedding lookups are extremely fast in practice. The entire operation can be implemented as a gather operation on the weight matrix.

Over the course of training, the rows of $\mathbf{E}$ are shaped by gradient descent to capture the semantic and syntactic properties of each token. Tokens that appear in similar contexts will develop similar embeddings. Tokens that serve similar grammatical roles will develop representations that are geometrically related. The embedding matrix effectively becomes a learned dictionary that maps the discrete world of tokens to the continuous world of vector spaces.

The output projection matrix $\mathbf{W}_{\text{out}} \in \mathbb{R}^{V \times d}$, also called the language model head or LM head, performs the inverse operation. After the transformer has processed the input sequence and produced a final hidden state $\mathbf{h}_{\text{out}} \in \mathbb{R}^d$, the model needs to decide which token should come next. This decision is made by computing a score for every token in the vocabulary:

$$\mathbf{z} = \mathbf{W}_{\text{out}} \mathbf{h}_{\text{out}}$$

where:

• $\mathbf{z}$: the output logits vector of length $V$, one unnormalized score per vocabulary token
• $\mathbf{W}_{\text{out}}$: the output projection matrix of shape $V \times d$, where each row encodes the "target pattern" for one vocabulary token
• $\mathbf{h}_{\text{out}}$: the final hidden state from the last transformer layer, a vector of dimension $d$

Each row of $\mathbf{W}_{\text{out}}$ contains what you might call the "output embedding" for one vocabulary token. The full matrix multiplication computes the dot product between the hidden state and each token's output embedding in one vectorized operation, yielding a score for every token in the vocabulary simultaneously.

After computing these logits, we apply the softmax function to convert them into a proper probability distribution over the vocabulary for the next token prediction:

$$P(\text{token}_i \mid \text{context}) = \frac{\exp(z_i)}{\sum_{k=1}^{V} \exp(z_k)}$$

where:

• $P(\text{token}_i \mid \text{context})$: the probability of token $i$ being the correct next token given the context
• $z_i$: the logit for token $i$, computed as $\mathbf{w}_i \cdot \mathbf{h}_{\text{out}}$ where $\mathbf{w}_i$ is row $i$ of $\mathbf{W}_{\text{out}}$
• $\exp(z_i)$: the exponential applied to logit $i$, which ensures all values are positive
• $\sum_{k=1}^{V} \exp(z_k)$: the sum of all exponentials across the full vocabulary, serving as a normalizing constant

Why does this formula make sense? Notice that taking the exponential of each logit guarantees positive values regardless of the sign of the original logit, which is a prerequisite for valid probabilities. The denominator then ensures all values sum to exactly 1. Tokens with higher logits, those whose output embeddings more closely align with the hidden state, receive higher probabilities.

Notice the dimensional symmetry: both $\mathbf{E}$ and $\mathbf{W}_{\text{out}}$ have exactly the same shape, $V \times d$. Both have $V$ rows, one per vocabulary token. Both have $d$ columns, one per hidden dimension. Both map between the same two spaces: the discrete token space with $V$ elements and the continuous embedding space with $d$ dimensions.

This symmetry is not a coincidence. It reflects the deeper fact that both matrices perform the same abstract function: they define a correspondence between tokens and points in a $d$-dimensional vector space. The input embedding says "token $t$ lives at point $\mathbf{e}_t$." The output projection says "a hidden state should produce token $t$ with probability proportional to how much it resembles $\mathbf{w}_t$." These two statements describe the same geometric relationship from opposite directions.