The Logit Lens
What if we don't wait for the final layer to see the prediction? By attaching the final vocabulary decoder to the middle layers, we can watch the model 'change its mind' in real-time.
Why Does This Exist?
In a standard Large Language Model, the process of predicting a word takes many steps. If the model has 32 layers, the input vector travels through Layer 1, then Layer 2, all the way to Layer 32. Finally, the Layer 32 vector is passed through an Unembedding Matrix (the final classification head), which converts the mathematical vector into human-readable words from the vocabulary.
If you look at the vectors in Layer 15, they are completely unreadable to humans. Researchers traditionally used probing-classifiers to understand them.
But in 2020, Nostalgebraist (an AI researcher) asked a brilliant, simple question: "What if we just take the Unembedding Matrix from the very end of the model, and prematurely attach it to Layer 15?"
The result was the Logit Lens. It allows us to instantly translate the hidden, unreadable thoughts of the middle layers into human words, giving us a step-by-step transcript of the model's internal reasoning process.
Think of It Like This
Think of It Like This
Think of a Transformer like a massive, 32-story assembly line building a car.
Normally, you are only allowed to see the final, fully assembled car as it rolls out of the factory on the 32nd floor. You have no idea what it looked like on the 10th floor.
The Logit Lens is like ripping the final "Paint and Polish" machine off the 32nd floor, carrying it down to the 10th floor, and forcing it to paint whatever unfinished chassis is currently sitting there. By looking at the half-painted, half-finished car, you can understand exactly how far along the assembly process the factory is at that specific moment.
How It Actually Works
The Logit Lens relies on the architectural quirk that a Transformer's hidden state maintains the exact same dimensions across all layers (e.g., a vector of size 4,096). Because the dimensions never change, the final Unembedding Matrix can be mathematically applied to any layer.
1. The Setup
You give the model a prompt: "The capital of France is".
2. Tracing the Layers
As the data flows through the network, you pause at every single layer, extract the hidden state, and pass it through the Logit Lens (the Unembedding Matrix).
You record the top prediction at every step:
- Layer 1: "Apple" (The model is just processing raw syntax; its prediction is random noise).
- Layer 5: "City" (The model has realized the context is about geography).
- Layer 10: "London" (The model knows it needs a European capital, but hasn't fully processed "France" yet).
- Layer 15: "Paris" (The model has solved the puzzle halfway through the network).
- Layer 32: "Paris" (The model spends the remaining 17 layers simply reinforcing its confidence).
3. What We Learned
The Logit Lens was a massive breakthrough in Mechanistic Interpretability. It proved that LLMs do not wait until the final layer to make decisions. They often solve the prompt in the very early layers, and spend the deeper layers refining the grammar, tone, and confidence of the output.
Show Me the Code
Implementing the Logit Lens in PyTorch using the Hugging Face transformers library is surprisingly easy, because you just need to grab the lm_head (the Unembedding Matrix) and apply it to the intermediate hidden states.
import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer
# 1. Load model and tokenizermodel = AutoModelForCausalLM.from_pretrained("gpt2")tokenizer = AutoTokenizer.from_pretrained("gpt2")
# 2. Prepare the inputtext = "The capital of France is"inputs = tokenizer(text, return_tensors="pt")
# 3. Run the model and ask it to return ALL hidden stateswith torch.no_grad(): outputs = model(**inputs, output_hidden_states=True)
# 4. Extract the final Unembedding Matrix (the "Lens")unembedding_matrix = model.get_output_embeddings()
# 5. Apply the Lens to every layerprint(f"Tracking the prediction for the final token across layers:")for layer_idx, hidden_state in enumerate(outputs.hidden_states): # Grab the vector for the very last word in the prompt final_token_vector = hidden_state[0, -1, :] # Apply the Logit Lens! logits = unembedding_matrix(final_token_vector) # Find the highest probability word predicted_token_id = torch.argmax(logits).item() predicted_word = tokenizer.decode(predicted_token_id) print(f"Layer {layer_idx:02d} thinks the next word is: '{predicted_word}'")Watch Out For
The Residual Stream Illusion
Transformers use "Residual Connections," meaning Layer 15 is actually just Layer 14 plus a small update. Because the Logit Lens reads this accumulated residual stream, it is very accurate. However, the Logit Lens struggles to decode the updates themselves (the outputs of the FFNs or Attention heads before they are added to the stream). Reading those intermediate updates requires more advanced techniques.
The Quick Version
- The hidden states of an LLM's middle layers are unreadable vectors of numbers.
- The Logit Lens is a technique that attaches the model's final vocabulary decoder to these early layers.
- This forces the early layers to prematurely output human-readable words.
- By tracking these words layer by layer, researchers can watch the model's "train of thought" evolve from random noise, to conceptual understanding, to the final correct answer.
- It proved that LLMs often solve factual queries very early in the network, using the deep layers merely for refinement.
What to Read Next
activation-steering— Now that we can read the model's mind using the Logit Lens, can we inject our own vectors into the middle layers to brainwash it?sparse-autoencoders— The Logit Lens can only translate vectors into the vocabulary. What if the vector represents a concept that doesn't have a specific word?