Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Logit Lens › Layer-by-Layer Entropy Analysis

Michael BrenndoerferFebruary 12, 20264 min read

Language AI HandbookAll chapters

Part 59 of 67Interpretability

Chapter 452Logit Lens › Layer-by-Layer Entropy Analysis

Reading section 8 of 12 · Layer-by-Layer Entropy Analysis · 4 minSection 8 of 12 · 4 min

Section 8 of 12 · 4 min read

Logit Lens › Layer-by-Layer Entropy AnalysisOne section · Sections in this chapter8 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Beyond tracking specific tokens, we can analyze the entropy of the logit lens distribution at each layer. Entropy is a principled measure of uncertainty: high entropy means the probability mass is spread across many tokens, while low entropy means it is concentrated on a few.

For a probability distribution $p$ over a vocabulary of size $V$ , the entropy in nats is:

H(p) = -\sum_{v=1}^{V} p_v \log p_v

where:

$p_v$ is the probability of token $v$ under the logit lens distribution at a given layer
The maximum possible entropy is $\log V$ (uniform distribution over all tokens)
The minimum entropy is 0 (all probability on a single token)

For GPT-2's vocabulary of 50,257 tokens, the maximum entropy is $\log(50257) \approx 10.8$ nats.

In[16]:

Code

def compute_layer_entropy(prompt, position=-1):
    """
    Compute the entropy of the logit lens probability distribution
    at each layer for a given token position.
    """
    hidden_states, tokens, input_ids = get_hidden_states(prompt)

    entropies = []
    for layer_idx, hs in enumerate(hidden_states):
        h = hs[0, position, :]  # (d,)
        _, probs = logit_lens_decode(
            h.unsqueeze(0),
            W_U,
            ln_weight,
            ln_bias,
            apply_final_norm=layer_idx < len(hidden_states) - 1,
        )
        probs_np = probs[0].numpy()
        # Entropy in nats: H = -sum(p * log(p))
        eps = 1e-10  # numerical stability
        entropy = -np.sum(probs_np * np.log(probs_np + eps))
        entropies.append(float(entropy))

    return entropies, tokens


entropies, tokens_ent = compute_layer_entropy(prompt, position=-1)

Out[17]:

Visualization

Line chart of non-monotonic prediction entropy across 13 GPT-2 states with a dashed theoretical maximum entropy reference line. — Entropy of the logit lens probability distribution across GPT-2 layers for the prompt ''The Eiffel Tower is located in the city of''. The trajectory is non-monotonic: layers alternately sharpen and broaden the distribution as the model moves from generic continuations toward competing city names. The dashed line marks the vocabulary-wide maximum entropy.

The entropy curve is non-monotonic for this prompt. Drops identify layers that concentrate probability, while rises show layers that reopen competition among plausible continuations. The final distribution can therefore be semantically better even when its entropy is higher than an earlier, confidently generic prediction.

The entropy at the embedding layer (index 0) is also informative as a baseline. It tells you how much information the model can extract from just the token identity and position, before any contextual processing. For common words in predictable positions, this baseline entropy can be surprisingly low: the word "the" is almost always followed by a noun phrase, so even the embedding-level entropy for predictions after "the" is not at the theoretical maximum.

Using Entropy to Compare PromptsLink Copied

Entropy analysis across layers is useful for comparing how "hard" different prompts are for the model. A prompt where entropy drops sharply in layer 4 is "easy": the model commits early. A prompt where entropy stays high until layer 11 is "hard": the model keeps processing until the last moment. This comparison can reveal patterns in where the model's knowledge and uncertainty are concentrated, and can help identify which types of inputs are most taxing for the model's processing pipeline.

Key ParametersLink Copied

The logit lens implementation involves several important design choices:

output_hidden_states=True: Must be set in the HuggingFace model call to extract intermediate representations.
LayerNorm application: Using the final layer's LayerNorm parameters is standard for the logit lens. Some implementations experiment with using no normalization or the layer's own normalization; results vary.
Weight tying: When the embedding and unembedding matrices are tied (as in GPT-2), the logit lens has a clean geometric interpretation as dot product with token embeddings. In models without weight tying, the unembedding matrix must be accessed separately.
Top-k vs. argmax: Looking at top-5 predictions captures more of the probability mass in early layers, where the distribution is flat. Argmax alone can be misleading when the model is uncertain.