Large language models improve decoding by selecting informative hidden layers
Surfing the Activation Wave: Analyzing the Information and Geometry of LLM Hidden States for Speculative Decoding
Computation and LanguageArtificial Intelligence
Summary
When large language models generate text, different layers inside the model hold different useful information. The researchers found a way to pick which layers to use for faster and better text predictions by measuring how much each layer’s information matters and how different they are from each other. They also developed a method to ignore parts that give less useful clues, improving the process even more. Their approach beats common methods and helps explain why some standard choices still work well.
What this means in practice
- •For llm infrastructure teams: Improve speculative decoding speed and accuracy by selecting the best hidden layer sets using entropy and alignment analysis.
- •For nlp platform engineers: Optimize text generation efficiency with a principled method for hidden state filtering and layer selection in LLMs.
Authors
Raghavv Goel, Sudhanshu Agrawal, Mukul Gagrani, Risheek Garrepalli, Mingu Lee
Abstract
Decoder hidden states from different layers of large language models (LLMs) often encode complementary information. However, existing speculative decoding methods typically rely on a small set of target-model layers chosen using fixed, position-based heuristics. Selecting an optimal subset of layers is inherently a combinatorial problem, making exhaustive evaluation infeasible even for moderately deep models. In this work, we present a principled framework for multi-layer hidden-state selection for speculative decoding based on two complementary criteria: informativeness, measured through matrix-entropy analysis, and diversity, quantified using centered kernel alignment (CKA). We investigate layer-wise Shannon entropy, aggregate entropy under hidden-state concatenation, and pairwise CKA dynamics across depth, revealing distinct geometric and informational compartments within decoder representations. Motivated by these findings, we propose entropy and CKA-based layer selection algorithms that favor hidden-state sets that are both highly informative and minimally redundant. As part of this framework, we introduce token-entropy filtering, which removes hidden states associated with low predictive-entropy tokens, thereby improving the quality of layer-selection statistics. In EAGLE3-style speculative decoding, our selected sets outperform fixed baseline choices, including the default $[2, N//2, N-3]$ rule, while providing a clearer explanation of why some defaults can remain competitive. Our results provide a practical recipe for multi-layer hidden-state selection for speculative decoding, insights into the geometry and information in LLM decoder states, and a foundation for extending similar analysis to related applications.