ICML 2025

Measuring In-Context Computation Complexity via Hidden State Prediction
Vincent Herrmann 1 R´obert Csord´as 2 J¨urgen Schmidhuber 3 1
Abstract
Detecting when a neural sequence model does “in-
teresting” computation is an open problem. The
next token prediction loss is a poor indicator: Low
loss can stem from trivially predictable sequences
that are uninteresting, while high loss may reflect
unpredictable but also irrelevant information that
can be ignored by the model. We propose a better
metric: measuring the model’s ability to predict
its own future hidden states. We show empirically
that this metric—in contrast to the next token pre-
diction loss—correlates with the intuitive inter-
estingness of the task. To measure predictability,
we introduce the architecture-agnostic “predic-
tion of hidden states” (PHi) layer that serves as
an information bottleneck on the main pathway
of the network (e.g., the residual stream in Trans-
formers). We propose a novel learned predictive
prior that enables us to measure the novel infor-
mation gained in each computation step, which
serves as our metric. We show empirically that
our metric predicts the description length of for-
mal languages learned in-context, the complexity
of mathematical reasoning problems, and the cor-
rectness of self-generated reasoning chains.
1. Introduction
Neural sequence models—especially large language models
(LLMs)—have demonstrated remarkable capabilities, such
as solving complex reasoning tasks (Lewkowycz et al., 2022;
Wei et al., 2022; Zelikman et al., 2022) and performing in-
context learning (Brown et al., 2020; Zhao et al., 2021;
Liu et al., 2022; Kirsch et al., 2022). Recent progress in
mechanistic interpretability has offered glimpses into how
these models process information internally: for instance, by
identifying specific induction heads (Olsson et al., 2022) or
1IDSIA/USI/SUPSI, Lugano, Switzerland 2Stanford University,
Stanford, USA 3KAUST, Thuwal, Saudi-Arabia. Correspondence
to: Vinc...