🤖 AI Summary
This work addresses the challenge of fine-grained quantification of computational energy consumption across individual steps in chain-of-thought (CoT) reasoning within large language models. We propose the Step-Aware Reasoning Energy (SARE) framework, which leverages centered kernel alignment (CKA) on Gram matrices derived from token hidden states between adjacent Transformer layers to geometrically measure reasoning energy at the single-step level. This energy metric is embedded into semantic state transition trajectories for modeling. SARE enables, for the first time, precise characterization of per-step CoT energy consumption, revealing phase-transition-like dynamics during reasoning and demonstrating that internal energy dynamics encode predictive information beyond output confidence. Experiments across six reasoning benchmarks and three open-source large language models show that energy consumption varies significantly across step types, with erroneous trajectories exhibiting notably lower energy at critical nodes, and that SARE features consistently outperform output-confidence baselines.
📝 Abstract
Understanding how computational effort is allocated across individual chain-of-thought (CoT) reasoning steps remains an open challenge: existing interpretability methods rely on output-level signals or collapse processing depth into a single trajectory-level scalar, leaving step-wise effort opaque. We propose Step-Aware Reasoning Energy (SARE), a geometric framework that quantifies effort at the granularity of individual CoT steps via Centered Kernel Alignment (CKA) between Gram matrices of token hidden states across adjacent transformer layers, capturing inter-token relational structure without requiring eigenvector alignment or cluster correspondence. SARE further contextualizes this energy within reasoning's semantic progression by modeling CoT trajectories as transitions among latent semantic states. Across six reasoning benchmarks and three open-weight LLMs, we find that reasoning energy is highly non-uniform across step types, exhibiting phase-like transitions invisible to trajectory-level metrics; incorrect trajectories show systematically lower energy at critical reasoning junctions; and SARE-based features match or outperform output-based confidence baselines in most settings, indicating that internal geometric dynamics encode predictive information beyond surface-level signals.