Unembedding Matrix

The unembedding matrix W_U is often tied to the transpose of the embedding matrix W_E through weight tying. This parameter sharing forces the input and output representations to live in the same space.

• Weight Tying: Reduces parameter count and enforces a shared semantic geometry between input tokens and output logits.
• Folding Norm: The unembedding vector for a token is its 'ideal' representation that the residual stream must align with to predict that token.
• Interpretation: The logit for a token is the dot product between the final residual stream state and that token's unembedding vector.

The Logit Lens technique applies the unembedding matrix to intermediate residual stream states, not just the final layer. This reveals the model's iterative prediction process.

• Early Layer Predictions: Applying W_U at early layers often shows the model predicting nonsensical or highly generic tokens.
• Convergence: As layers progress, the top predicted token converges toward the final output, showing how the model refines its belief.
• Interpretability Tool: Researchers use this to diagnose where specific factual recall or syntactic decisions occur in the forward pass.

The unembedding matrix is followed by a softmax function to produce a probability distribution. This creates a well-known representational bottleneck.

• Rank Limitation: The logit for a token is a linear function of the hidden state. The model cannot express arbitrary probability distributions; it is constrained by the dimensionality of the residual stream.
• Mixture of Softmaxes: Advanced architectures break this bottleneck by using multiple unembedding matrices and mixing their probabilities, allowing for more expressive output distributions.
• Impact: This limitation is a key reason why language models struggle with tasks requiring precise, multi-modal output distributions.

The rows of the unembedding matrix form a semantic basis for the vocabulary. The geometry of this space encodes linguistic relationships.

• Cosine Similarity: The cosine similarity between unembedding vectors for two tokens correlates with their semantic relatedness.
• Linear Analogies: Famous analogies like 'king - man + woman = queen' can be performed in the unembedding space, demonstrating that it captures structured relational knowledge.
• Concept Directions: Directions in the unembedding space can correspond to interpretable concepts, such as sentiment or tense, which can be used for steering.

The unembedding matrix is the final step for logit attribution methods that trace a prediction back to input tokens or model components.

• Direct Attribution: The contribution of a specific attention head or MLP neuron to the final logit is computed by multiplying its output by the unembedding matrix.
• Gradient-Based Saliency: The gradient of a target logit with respect to the input embeddings is computed through the unembedding matrix, revealing which input tokens were most influential.
• Causal Tracing: In causal tracing for factual recall, the unembedding matrix is used to project intermediate states to logits to measure the restoration of a corrupted fact.

By manipulating the logits before the softmax, the unembedding matrix enables direct control over the model's output distribution.

• Logit Bias: Adding a constant bias vector to the logits can suppress or promote specific tokens, a common technique in API-based inference.
• Contrastive Decoding: Methods like DoLa subtract the logits from an early layer's unembedding projection from the final layer's logits to amplify factual knowledge and reduce hallucinations.
• Steering Vectors: A direction computed in the residual stream can be mapped through the unembedding matrix to understand which tokens it promotes, validating the vector's semantic effect.