Entropy-Lens: Uncovering Decision Strategies in LLMs
Riccardo Ali, Francesco Caso, Christopher Irwin, Pietro Liò
cs.LG, cs.AI, cs.CV
2025-02-23
Cambridge's Entropy-Lens turns logit-lens entropy into a depth profile: family-specific expansion/pruning, task- and format-separable, with expansion layers more critical on MMLU.
Most interpretability work lives in embedding space. Token-space dynamics are harder: next-token distributions are high-dimensional and unordered, so ordinary moments misbehave. The Cambridge group compresses each layer to one number. Project the residual stream through the model's own unembedding (logit lens), then record Shannon entropy. The resulting per-layer sequence is an entropy profile.
Entropy going up is read as the candidate set expanding; entropy going down as pruning. Candidates are tokens inside the top-p nucleus with p = 0.6. Spearman correlations back the reading: 0.8250 / 0.7365 on Llama-3.2-1B / 3B and 0.8804 / 0.8540 on Gemma-2-2B / 9B. Adjacent layers also share a large fraction of their top-p sets, so the rise and fall is mostly the same shortlist growing and shrinking.
For every generated token and every layer, compute Shannon entropy. Concatenate profiles across tokens in a generation, then hand the vector to a kNN used only as a diagnostic: high AUC means the profiles separate. Frozen off-the-shelf Instruct models, no gradients, no trained probes. Shannon entropy is the default; Rényi entropy at α = 0.5 / 1 / 5 leaves the conclusions intact.
Across 12 decoder-only models from about 100M to 9B (GPT, Gemma, Llama, Qwen), blank-prompt profiles cluster by family, not by size. After depth is normalized, same-family models of different sizes line up: smaller models look like a coarser sampling of the same expansion–pruning schedule. GPT-family profiles often start high-entropy and sharpen; Llama-family profiles start low, plateau, then refine.
Task type is visible. On TinyStories, continue-the-story, count-the-words, and extract-the-gist, 800 prompts each, three-class kNN AUCs are 97.66 / 98.38 for Gemma-2 2.1B / 8.9B, 94.94 / 94.77 for Llama-3.2 1B / 3B, 96.10 for Llama-3-8B, and 97.07 for Phi-3 3.6B. Word-scrambled stories are a control, so the signal is not just lexical. Output format separates too: poem, scientific blurb, chat log. Gemma-2-2B-it reaches 98.7 ± 1.1 AUC with Shannon entropy; Llama-3.2-1B-it reaches 97.8 ± 2.4.
Profiles also carry correctness signal on MMLU. Llama accuracy is 50.89–60.62 across three prompt styles, with kNN AUC 67.23–73.61; Gemma accuracy is about 55–56, with kNN AUC 68.36–72.78. Skipping layers of maximal entropy increase (expansion) versus decrease (pruning), against random skips, shows unequal cost. On Llama variants, skipping the single largest-expansion layer collapses MMLU to chance. Gemma2-2B-it is the exception: pruning layers matter more.
One forward pass, no trained probe, and you can see which layers widen the shortlist and which shrink it. Family fingerprints help when picking a base model or scaling within a line: size looks like resolution, not a different algorithm. If you skip or early-exit layers, protect expansion phases, at least on Llama.
The authors do not know which architectural or training choices stamp a family profile. Whether SFT or RLHF can flip the expansion-versus-pruning ranking is open; Gemma2-2B-it is the one reversal. The main study is decoder-only Transformers; ViTs and MoE are only touched in the appendix. High kNN AUC shows separability, not a causal proof that entropy changes "are" expansion and pruning. Layer skipping zeros a block's residual write, which is not the same as trained early exit.