Measuring In-Context Computation Complexity via Hidden State Prediction
Authors: Vincent Herrmann
Róbert Csordás
Jürgen Schmidhuber
Venue: ICML 2025 poster (2025)
Keywords: in-context learning
interpretability
transformers
Primary area: deep_learning->sequential_models_time_series
TL;DR: Hidden state unpredictability in sequence models is a meaningful measure for in-context reasoning complexity

Abstract: Detecting when a neural sequence model does "interesting" 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 prediction loss–correlates with the intuitive interestingness of the task. To measure predictability, we introduce the architecture-agnostic "prediction of hidden states" (PHi) layer that serves as an information bottleneck on the main pathway of the network (e.g., the residual stream in Transformers). We propose a novel learned predictive prior that enables us to measure the novel information gained in each computation step, which serves as our metric. We show empirically that our metric predicts the description length of formal languages learned in-context, the complexity of mathematical reasoning problems, and the correctness of self-generated reasoning chains.

Measuring In-Context Computation Complexity via Hidden State Prediction
Vincent Herrmann Róbert Csordás?
Jürgen Schmidhuber 31
Abstract
Co-tentttt Learning
Literature
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 metrie: 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 correlate with the intuitive inter- estingness of the task. To measure predictability, we introduce the architecture- agnossic "predic- tion of hidden states" (PHi) layer that serves information bottleneck 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 metric predicts the description length of for- mal languages learned in-context, the complexity of mathematical reasoning problems, and the cor- rectness of celf-generated reasoning chains.
10
600
Interesting
Tasks
PHI Loss
Boring Tasks
Memorized Programs
Random
500
Memorized Sequences
Trivial
400
Memorized
Special ized Model
Random
300
Pre-trained LLM
10" Next Token
Figure Interesting tasks, like in in-context Iearning, or modeling of code and literature, exhibit high hidden state prediction (PHi) loss, while boring or trivial tasks, such as retrieving memorized sequences or modeling random structureless data, show low PHi loss. Next token loss provides no meaningful insight into task complexity. Results for specialized transformer model (blue) and pre-trained LLM (green), with PHi loss scales differing due to hidden state size. See Sections 3.1.1 and 3.2.1 for details.
n-gram mechanisms (Akyürek al., 2024) that capture non- trivial aspects of the underlying computation. Nevertheless, quantifying in general way when and how much truly interesting or complex computation is happening within model remains challenging.
1. Introduction
A natural first instinct to use next-token prediction loss to gauge complexity: Computations that yield high loss might appear "hard," whereas those with low loss might be "easy". However, this notion of hardness does not necessar- ily reflect the complexity or "interestingness" of underlying computation, see e.g., Schmidhuber (1991a;b), For example, a random sequence has a high next token prediction loss, but at the same time, the model does not have to compute anything because there nothing that can do to help the predictions. Hence, next-token prediction alone fails cap- ture the complexity of in-context computation performed by the model.
Neural sequence models-especiiall large language models (LLMs)-ae 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
'IDSIA/USI/SUPSI Lugano, Switzerland Stanford University, Stanford, USA "KAUST, Thuwal, Saudi-Arabia, Correspondence to: Vincent Hermman <vincent. idsia.ch>.
We therefore propose a new metrie: hidden-state predictabil- ity. Concretely, we measure how well the model can predict its own future hidden states, rather than future tokens. To achieve this, we introduce PHi ("Prediction of Hidden
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
States") layer, which imposes an information bottleneck on the hidden representations, and predicts them autore- gressively. The layer is placed in the information pathway of the autoregressive sequence model. Under this scheme, the activations flowing through the PHi layer have to be sumpressed-informaae
X
Xg
Upper Sugaris Modill Lagers
0g
is included in the hidden state only expected to be relevant for generating future tokens. The model is encouraged to encode the "in-context program" needed for sequence prediction in the hidden states. The hidden state prediction loss (PHi loss) allows us to quantify the complexity, or description length, of this program.
42
Cansal Self Attezzion
PHi
Layer
4
22
23
4v
4V
Our approach offers several advantages: (i) It is architecture- agnostic and can be inserted in a wide range of autoregres- sive models, including transformers (Vaswani et al., 2017: Schmidhuber, 1992b; Schlag et al., 2021) and RNNs (El- man, 1990; Jordan, 1986; Hochreiter & Schmidhuber, 1997). (ii) The model and the PHi layer can be jointly trained from scratch, or the PHi layer can be added post-hoc to pretrained LLMs. (iii) provides a way to measure when the model is doing complex computation by examining the unpredictabil- ity-and hence information content-of the hidden states.
hg
Lowre Sequence Model Lagens R,
X2
The structure our layer. It can inserted in the middle of any next token prediction architecture, and it reconstructs its hidden states through an information bottleneck. It consists of posterior encoder (xe | ht) that predicts the latent code &t, decoder ag used to reconstruct the hidden state 1ú = ag(zz). and learned autoregressive prior Px 21, 2t-1) predicting 2, from the past latent codes. We propose to use the KI. divergence between the posterior and the prior quantify the complexity of the "in- inn-conextt computation" performed by the model.
Figure 2.
We provide empirical evidence, using both smaller-seut transformers/RNNN and large pre-trained language models, that our hidden-state state unpredictability metric correlates with task complexity in various domains. Concretely, we show that:
Tasks that require no interesting data-processing such as reciting memorized sequences, executing memo- rized programs in-context, or modeling random noise, are associated with low PHi loss. Interesting tasks such as in-context learning or modeling complex natural lan- guage and code yield high PHi loss (see Figure 1).
effectiveness of our approach on variety of tasks, showcas- ing how hidden-state predictability correlates with task com- plexity. Section 4 reviews related work before we discuss limitations and avenues for further research in Section 5. We conclude with Section 6. Our code is publicly avail- able (https:/ predictiin--hidden-stta
For in-context learning formal stochastic languages, PHi loss is predictive (beyond next token loss) of the complexity of the language.
2. Method
In the following, Section 2.1 introduces our PHi layer, which serves as basis of our novel metric. Then, we build on this layer to introduce our proposed quantification of "in- terestingness" in Section 2.2. In Section 2.3, we discuss the training of our layer, and finish with Section 2.4 by discussing the connection between our proposed PHi loss and the next token prediction loss.
For step-by-step solutions of math problems, PHi loss is predictive (again, beyond next token loss) of the hardness of the problem.
When an LLM itself generates step-by-step solutions to math problems, solutions with high PHi loss have significantly increased chance of being correct.
2.1. Prediction of Hidden States (PHi)
Taken together, these insights shed light on the internal computations of powerful sequence models and offer a prin- cipled tool for detecting "interesting" in-context behaviors. The remainder of this paper is organized as follows: In Sec- tion 2, we introduce the PHi layer and describe how it is used to measure predictability, and how it is integrated into an autoregressive model. In Section 3, we demonstrate the
We consider standard autoregressive sequence model mo. with parameters 0, so that
po(xI.))
II To( xer).
(1)
2

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
In practice, Te(- Xite) can be realized by any neural se-
ingness of the computation performed in specific timestep.
quence model such as decoder-only transformer RNN. As long as we has at least two sequence layers, we can split this model into a bottom set of layers and a top set of layers.
When quantifying the interestingness of an entire sequence, we calculate the mean of the above metric over all tokens:
LPm
Σ (pun()) T
1
(3)
Bottom Layers: The function Bg, with parameters 3, pro- cesses the sequence (I'....., 2x) into a hidden representa- One can think of Bg as the first few layers of the network.
tion h4 (3
We report this metric in our experiments in Section 3.
2.3. Training Objective
Top Layers:
The top (remaining) layers, denoted T7, with
parameters T, use hidden representations {fi,..., h } to parameterize the next-token distribution: no xet) =
We train the entire system (bottom layers B3. top layers Tr, and PHi layer) by combining two losses. First, we need to ensure that ensures that the transformed hidden states (M,...., still suffice to predict the token Iy. We do this by using the standard negative log- ligliikd next token prediction loss used in language modeling:
PHi Layer: We introduce a new layer that has distinct func- tions: It (a) introduces information bottleneck the current hidden state h, and (b) predicts future hidden states. We call this a PHi layer (see Figure 2). It is sandwiched between B3 and Tr, and transforms { h into a new sequence of representations {h/...., K3.. The layer is inspired by variational autoencoders (Kingma& & Welling, 2014), but instead of a fixed prior, we use an autoregres- sive, learned prior, with access the past states. First, we introduce latent variables {z1,...,22 ] at each time step. The learned encoder qu qu(22 h2) maps each hidden state hy to a posterior over 2,. The autoregressive, learned prior Px (2t .,,...,, 2)-) generates a distribution over 2, based on the previous history of latent variables. Thus, the prior serves as a predictor of future hidden states. In practice, the prior can be implemented by any causal sequence processing module, such as self-attention combined with a step-wise multi-layer percepetron (MLP). A decoder ag, parametrized by a small neural network, is used to reconstruct the hidden states M/ = ag(22 from the latents 2, ~ q+(+ ht).
(4) In parallel, we also minimize the PHi loss (Eq. 2). This encourages each 2, to be predictable from its predecessors Z<e by minimizing the KL divergence between the posterior and the autoregressive prior. Thereby, the "unnecessary" information that is not expected to be relevant for future predictions is discouraged in 2t.
(t) - log logoTT (ara 1hi..... 11-11
The total loss for time step t is:
L(I) = (t) + LPHi(t)
(5)
In effect, the model is incentivized to maintain enough in- formation in the latent (and hence h2) to predict the next token, but not so much as to deviate from the learned autore- gressive prior over 2. This enforces an information bottle- neck on the hidden states while preserving good predictive performance on the original sequence task.
The PHi layer can easily be inserted into most autoregressive sequence modeling architectures. The sequence model me can be trained jointly with the PHi layer, alternatively can be pre-trained, optimizing only the new PHi layer while holding the model's weights fixed.
We treat the latent variables {=e} in fashion similar to variational autoencoder (VAE) by sampling them from the approximate posterior go (= he). To enable gradient backpropagation through these samples, we use the stan- dard reparameterization trick (Kingma & Welling, 2014; Rezende et al., 2014). Specifically, the posterior is modeled as a Gaussian with mean (ha) and diagonal covariance ow (he )2, we draw an auxiliary noise variable (O0, and set 2, p (h.) + (he) e. This approach provides low- varvanriance gradient estimator for updates to the parame- ters v. Consequently, the entire network (including bottom layers B3, the prior Px, and the output transform ag) can be trained end to end via standard stochastic gradient descent.
2.2. Quantifying "Interestingness"
We propose to quantify interestingness by measuring the amount of information gained from the posterior q, hy) over the prior (x 21...., 11.) This is precisely the KL divergence DKL (qv ll We define the PHi loss as the the information gained in given timestep t:
LpHi (t) = DKa( (10 (1) || px(x 21, 21,..., 21-1) (2)
2.4. The Connection Between LNL. and Lem
Intuitively, this measures the nats of novel information learned in timestep that is not predictable from the past. We propose to use this quantity as the measure of the interest-
Because the token sequence G1r is modeled autoregres- sively by (1), we can use a standard entropy-coding tech- nique (e.g. arithmetic coding, compare Deletang et al.
3

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
(2024)) to encode T1.T in
associated token, move the edge's target state, and repeat for a random length of 10-50 tokens. In in-context language learning, for each input, new random PFA is generated, and the model is presented with multiple examples gener- ated from this PFA. To achieve good next-token predictions on subsequent examples, the model must infer the structure of the PFA in-context.
T Σ (t)
(-log log no(as x<e))
nats.
The PHi layer allows to proceed similarly with the hid- den state sequence, more precisely with the sequence of posteriors {qe (-Ihe), E T....T) over the latent variables 21.T- The number of nats required to encode the entire latent sequence is
We quantify the theoretical complexity of the "program" that must be learned in-context by calculating the number of bits necessary to encode the PFA. Let a PFA A have states and m edges, each labeled with token from a chosen subset of size v out of a total vocabulary of size V. We encode A by (1) specifying log, (*) bits to choose which u tokens are used; (2) listing m edges, each requiring log, n bits for the origin state, log, n bits for the target state, and log2 u bits for the token. Hence, the complexity is
Σ DRu( (gl(- he) Px( 221)) = Σ LPHi (t).
This description length of the hidden states allows us to quantify the amount of information that the model is extract- ing from the imput sequence and is used to predict the remain- der of the sequence. Note that the token zg at any particular LPH nats by first encoding the latent sequence 2 - 1 and then encoding xi given T, (- , without explicitly en-
C(A) = log (M) + m(2 logn + log v). (6)
position i alone can be encoded in LNLL (f) + f-1
Since A changes with the problem instance, the model must synthesize in its hidden states an in-context representation of complexity C(A) to predict future tokens effectively.
coding the past sequence of tokens.
1. BORING VS. INTERESTING TASKS
3. Experiments
We first investigate how the PHi loss differs among tasks of varying complexity. Concretely, train our Transformer and LSTM models on a mixture of four tasks: (1) Retrieval of Memorized Subsequences: A small, fixed set of ran- dom subsequences is repeatedly used during training and testing. Each overall sequence is a random concatenation of these subsequences. The model can simply memorize this small set in its static parameters and requires almost (2) Retrieval of Memorized Programs: A small set A of automata (PFAs) is fixed dur- ing training and testing. Each sequence is generated by an automaton in A. Again, the model needs only to identify which automaton from A is being used and then rely on its memorized representation (minimal in-context inference). (3) In-Context Language Learning: Each sequence is gen- erated from new, unfamiliar PFA A. The model must learn A in-context, effectively synthesizing a program of complexity C(A). This is a non-trivial in in-context learning task. (4) Random Examples: Each sequence consists of random tokens with no structure to learn. While next-token prediction is difficult (leading to high loss), no meaningful in in-context computation required.
We split the experiment section into two parts: Section 3.1 covers smaller scale sequence models that are trained jointly with the PHi layer from scratch. Section 3.2 investigates inserting a PHi layer into a pre-trained LLM.
3.1. Experiments with Fully Trained Sequence Models
We begin our evaluation by training autoregressive sequence models from scratch with the proposed hidden-state predic- tion framework. Specifically, we consider two architectures: a decoder-only transformer and an LSTM. The transformer has 12 layers and a hidden state size of 768. The LSTM has 2 layers and the same hidden state size. In each network, we insert single PHi layer in the middle of the model. Exact hyperparameters and further details can be found in Appendix A. Our primary goal is to explore how the PHi approach responds to different tasks, especially those that require non-trivial in-context computation.
in-context computation.
Background: In-Context Language Learning with Prob- abilistic Finite Automata We adopt the in-context lan- guage learning setup proposed by (Akyürek et al., 2024), where each problem instance is defined by probabilistic fi- nite automaton (PFA). A PFA has a set of states and directed edges, each labeled with a token from some vocabulary. To generate an example from a PFA, we start at a random initial state and randomly sample an outgoing edge (with uniform probability among that state's s edges). We emit the
Only Task 3 demands truly interesting in-context computa- tion: The model must discover the PFA's structure to predict future tokens effectively. By contrast, Tasks 1-2 can be solved by looking up memorized routine, and Task 4 is just uniform noise.
Figure 3 shows the next-token prediction loss for both archi-
4

---

Mcasuring In-Contex, Computation Complexity via Hidden State Prediction
3.0
Transtomer
LSTH
0.030
Next Toxen Loss
2.0
0.025
Complexity
1.5
PHI Loss
0.020
1.0
0.015
0.5
0.010
0.0
Seg.
Prog.
ICCL
Rondom
0.005
2
Figure 3. Next-token prediction loss on each the four tasks, for both the Transformer and the LSTM. Memorized tasks yield the lowest loss, random is the highest, and the in-context language learning task is intermediate. Bootstrapped mean with 95% confi- dence intervals across 10 runs.
0.0
0.5
1.0
1.5
2.0
2.5
3.0
Next Token Loss
Figure 5. Token-wise PHi loss (y-axis) versus binned next-token losses (X-axis), stratified by PFA complexity (color, grouped into 10 levels, from (simple) to 10 (complex)). Across all next-token prediction losses, more complex PFAS result a higher PHi loss.
Transforme
2.0
LSTT
strated that "boring" tasks (memorized or random) lead to more compressible hidden states (i.c., lower PHi loss), whereas "interesting" tasks (requiring non-trivial in-context learning) exhibit higher PHi loss. Now ask: Within the realm of in-context learning itself, does the degree of complexity correlate with the magnitude of the PHi loss?
Rolatve Pho LOSS
1.5
1.0
0.5
0.0
Mam. Seq.
Mam. Prog.
ICCL
Random
We again train our Transformer model on PFAS of varying complexity. Then we test the trained model on 1000 prob- lem instances, each based on a different unfamiliar PFA. For each sequence, we track token- wise PHi loss as well next- token prediction loss. Because more complex PFAs also tend have higher irreducible uncertainty (and thus higher next-token loss), we want to show that PHi loss explains task complexity beyond next-token loss. To accomplish this, we bin tokens by their next-token prediction loss and compare the average PHi loss across sequences different complexity levels. Figure displays the token-witg wise PHi loss versus binned next-token losses, stratified by PFA complex- ity (grouped into 10 levels, from 1 (simple) to 10 (complex)). We see clear trend: Even after controlling for next-token loss, more complex PFAS yield systematically higher PHi loss. Similarly, a partial correlation analysis (controlling for mean next-token loss) reveals a significant relationship between sequence-win mean PHi loss and PFA complexity (r 0.37, with 95% confidence interval of [0.32, 0.43)). These findings confirm that PHi loss tracks the amount of in-context computation needed, beyond what is implied by simple token-level uncertainty.
Figure 4. PHi loss, relative to the performance on the memorized sequences, for the four tasks. In-context language learning shows a significantly higher PHi loss than the other tasks, indi- cating that the model performing non-trivial computation in its hidden states to infer the unknown automaton. Bootstrapped mean with 95% confidence intervals across 10 runs.
tectures. As expected, Tasks 1 and 2 yield low loss, while Task 4 is the highest due to random sequences. Task 3 has an intermediate loss. Note, however, that intermediate loss alone does not clarify whether the model has learned an intricate structure ponents. In Figure 4 we can observe that we have high PHi loss only for Task 3. This is in accordance with our claim that hidden states difficult to predict when complex program is generated in-context (Task 3), but not when easy tasks are being performed (Tasks 1, 2 & 4). These results are consistent for both the Transformer and the LSTM model.
simply mixing easy and random com-
3.1.2. SIMPLE VS. COMPLEX TASKS
Here, we focus exclusively on in-context language learning (Task from Section 1.1.1) to investigate whether more complex tasks lead to higher PHi loss. Recall that the complexity of stochastic language can be quantified by the code-length C(A) of the corresponding PFA (see Eq. 6). Intuitively, larger or more intricate automata require more complex "in-context programs" to achieve low next-token prediction errors. The previous experiment 13. demon-
3.2. Experiments with Large Language Models
Next, investigate whether our findings generalize to LLMs that have already been pre-trained on extensive cor- pora. The model weights are held fixed, and we insert and train a single PHi layer within the existing architec- ture. Specifically, we use the instruction-tuned "Llama 3.2" model with 3B parameters (Dubey et al., 2024), and train a
5

---

Mcasuring In-Contex, Computation Complexity via Hidden State Prediction
Trivial
800
Trivial
Licenses
Licenses
Next Token Loss
Random
600
Random
Literature
Literature
PHi Loss
Code
Code
400
200
10
15
20
25
5
10 15
20
25
PHi Layer Position in Llama 3B Model
PHi Layer Position in Llama 38 Model
Figure 6. Next token prediction loss in a pretrained Llama 3B model with the PHi layer inserted at different layer positions. Dot- ted lines show the performance of the base model without PHi layer. Note that inserting in the lower layers causes sigmificant degradation in the model quality, most likely due to posterior col- lapse. We see no correlation between next token prediction loss and the fact of whether task is simple (Trivial, Licenses, Random) or complex (Literature, Code).
Figure 7. PHi loss for different tasks at different layers of a pre- trained Llama 3B model. For layers 18-24 (the highlighted region). we can see the most differentiation between PHi losses for differ- ent tasks, these layers, tasks that are intuitively "uninteresting", like predicting random sequences, memorized licenses, or trivial, memorized tasks, have a lower PHi loss compared to interesting tasks like code or literature prediction.
syntax, suggesting a non-trivial in-context program.
single PHi layer for 10,000 steps on a mixture of reasoning and natural-language data (see Appendix B for details).
We group tasks (a)-() as "boring" (no new in-context pro- gram) and tasks (e) as "interesting" (likely to require more in-context synthesis).
Unlike training from scratch, inserting PHi layer into pre-trained model raises design questions: Where in the model (which layer) do we impose the information bottle- neck? The pre- rre-trained model did not "expect" an internal bottleneck, some layers might not be able to accommodate it. Furthermore, which tasks are actually easy or difficult, simple or complex to an LLM? What kinds of texts and programs does it have memorized?
We insert the PHi layer at various depths in the Llama model, then measure PHi loss and its effect on the next-token prediction loss on each task. Figure 6 shows that placing the PHi layer in early layers can lead to posterior collapse: Next-token loss rises significantly (the bottleneck disrupts the flow of information), yet very low PHi loss. In contrast, placing the PHi layer in the later layers yields minimal impact on next-token accuracy and avoids collapse. In our analysis and the remaining experiments, we focus on the instances where the PHi layer is placed after layers 18-24. There, the PHi losses are most differentiated between tasks and we observe no posterior collapse. Figure/ 7 shows that the PHi loss indeed higher for unknown literature and code than for trivial memorized, or random tasks. Mean- while, the next-token loss (6) shows consistent pattern that discriminates "interesting" tasks from "boring" ones.
3.2.1. DIFFERENT LANGUAGE MODELING TASKS
The first experiment involving LLMs seeks to address these questions. We design test set comprising five types of se- quences, with varying expectations about memorization and in-context reasoning complexity: (a) Trivial Tasks: Exam- ples such as multiplication tables, counting to 50 in Spanish, or listing weekdays in January 2025. We expect the model to have these "programs" firmly memorized and easily trievable. (b) Memorized Licenses: Popular licenses GitHub that appear frequently in pre-training data, which the model can reproduce verbatim without additional in- context reasoning. (c) Random Words: Shuffled tokens lacking syntactic or semantic structure. Although this results in high next-token loss, there no interesting structure to learn in-context. (d) Unknown Literature: Recent Project Gutenberg chapters assumed unseen by the model. Natural language, especially literary text, often contains rich struc- ture that benefits from deeper in-context reasoning, (e) Un- known Code: A private python codebase, never encoun- tered during pre-training. Programming languages typically require the model to capture longer-range dependencies and
3.2.2. SIMPLE VS. COMPLEX TASKS IN LLMS
We next explore whether we can again distinguish simple versus complex tasks using PHi loss, this time a pre- trained LLM. We use the same Llama-based setup (Sec- tion 3.2.1). inserting PHi layer into the upper portion of the network while keeping the rest of the weights fixed.
PFA -ased Tasks. Following a methodology similar to the Experiment 3.1.2, we instruct the LLM to generate data from a probabilistic finite automaton (PFA), then present with an examples generated by particular problem instance.
6

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
63%
10
0.15
10"
60%
Correlation
x
PFA Somneeee
0.30
PHi Loss
2 102
Chance Choesing
55%
0.05
10P
-4
1,6 10?
50%
All Questions
1.4
Counterintultive Questions
10
10
10
45%
18
20
24
Token
PHi Layer Position in Laaa 38 Model
Figure 9. Chance choosing correct answer to mathematical reasoning questions from the GSM-8 8k dataset when selecting the one with the higher PHi loss between two option (one correct one false). The step-by-step rationales generated by the model itself. "Counterintuitive" are only those questions for which the answer with the lower next token prediction loss is wrong. In purple, the partial correlation between PHi loss and answer correctness- controlled for token prediction loss-is shown. Answers with high PHi loss are clearly more likely to be correct.
Figure 8. Token-wise PHi loss (y-axis) versus binned next-token losses (x-axis), stratified by PEA complexity (color). Results Llama 3B with the PHi layer placed after layer 20, For most next- token loss bins, more complex PFAs result higher PHi loss.
For this, we observe token-wise PHi loss and next-token prediction loss. We again bin tokens by their next-token loss and group PFA instances by complexity, C(A). Figure shows that, much like the smaller but specialized model shown in Figure 5, complex PFAs yield higher PHi losses overall, even after controlling for next-token loss.
the correct answer by random guessing in this setting is of course 50%). should be noted that when, from all solu- tion pairs, we always pick the ones with lower mean next token prediction loss, our chance of being correct is 71 %. This is surprising, since in general, PHi loss and next token prediction loss positively correlated. It suggests that the two losses capture different important aspects of the gener- ated rationale: A good solution to a mathematical question should be coherent (low next token loss) and interesting (high PHi loss). Also here, we perform a partial correla- tion analysis and find that PHi loss is highly predictive for correctness when controlling for next token prediction loss, across all tested PHi layer positions. Additionally, con- struct a subset of "counterintuitive" questions. These are questions for which the more likely answer (the one with lower next token prediction loss) wrong. Even for these question, we observe the increased chance of choosing the correct answer when selecting the one with higher PHi loss. These findings suggest that, for unfamiliar or challenging math problems, the model must synthesize in-context program (and hence incur higher PHi loss) to achieve correct solutions.
Mathematical Reasoning. We further test whether PHi loss tracks the complexity of mathematical reasoning. Specifically, use the MATH dataset (Hendrycks et al., 2021), which consists of math problems labeled from Level (easy) to Level 5 (hard), along with detailed reason- ing solutions. We measure the partial correlation between PHi loss on the step-by-step solution and the difficulty of the problem and find across all tested layers highly significant positive relationship (r 0.079, with 95% confidence interval of [0.07, 0.09).. In line with our overall hypothe- sis, more difficult math problems appear to demand a more complex "in-context program" in the hidden states, leading to higher PHi loss.
3.2.3. CORRECT VS. ERRONEOUS RATIONALES
Finally, we examine how PHi loss relates chains-of- thought generated by the model itself. Specifically, we look at the correctness of generated answers to questions from GSM-8k (Cobbe et al., 2021), a dataset of grade-school math problems. For each test question, we sample sev- eral solution candidates from the model, giving example problems with solutions in-context and employing chain-of- thought prompting. From these candidates, we assemble random pairs of rationales for each problem, with one of the generated rationales leading to correct final answer and the other one to an incorrect answer.
complex
We also perform the same experiments on the MATH dataset, previously used in Section 3.2.2. There, we obtain very similar results (see Appendix C.3). The classification of questions by difficulty allows for a more granular analysis. This breakdown reveals intriguing pattern in the coun- terintuitive subset: For easy questions, high PHi loss does not correlate with correct rationales, whereas for difficult questions, there a strong positive correlation between high PHi loss and correctness. We propose the following tenta- tive explanation: difficult counterintuitive questions, the
Figure shows that choosing the answer for which the ra- tionale has higher PHi loss leads to a significantly increased chance of picking the correct one (the probability of picking

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
model tends generate misleading rationales that are overly simplistic. Arriving at the correct answer requires "effort," " which is reflected in the PHi loss. In contrast, for easy coun- terintuitive questions, failure is not due to oversimplified rationales but might stem for example from genuine mis- conceptions within the model. Put differently, for difficult questions, we expect the correct answer to be "interesting" demanding complex reasoning-wheress
time separation. To uncover the mechanical underpinning of such capabilities, Akyürek et al. (2024) investigate the in- context learning of formal languages. They generate finite state automata on the fly and present transitions sampled from them to the model, and the model has to infer the structure of the automata in context. The authors call this setup "in-context language learning". We adapt their setup for our experiments in Sec. 3.1.
whereas for easy questions, this is not necessarily the case.
Prequential coding (Blier & Ollivier, 2018) is proposed to measure the complexity of neural models. It measures the amount of information needed to encode a network and its training data together, starting from an untrained model, training it online, and using the improving estimates from the model to compress the training data. The authors show that this results in good compression ratios, and more compressible models tend to generalize better. Building on these findings, Elmoznino et al. (2024) shows that the next token prediction loss naturally minimizes both the training error and the prequential code length of the implicit model learned by in-context learning, which can serve as a form of Occam's 's razor that encourages generalization.
4. Related Work
In early work on hidden-state prediction, Schmidhuber (1992a) proposed using one RNN to predict not only its next input but also the hidden state of another "chunker" RNN, which operates at a higher temporal level of abstrac- tion. However, the chunker's ' hidden states had no explicit incentive to be predictable by the lower-level RNN. This approach was explored in the context of a -on-geeerative autoencoder, encouraged to produce an internal representa- tion that is informative about its current input while at the same time being predictable from previous inputs (Schmid- huber & Prelinger, 1993). More recently, self-predictive representations have been proposed in the context of rein- forcement learning to encourage state representations that capture the important structure of the environment dynam- ics (Schwarzer et al., 2020; Guo et al., 2020; Ni et al., 2024).
Our work can be seen as explicitly quantifying and minimiz- ing an upper bound on the complexity of this implicit model (i.e., the "program") ") that is generated in-context to predict the next tokens. This perspective, of jointly minimizing a two-part description the complexity of a model plus the complexity of the data given that model (Grünwald, 2007)- relates closely to a concept from algorithmic information theory called sophistication (Koppel, 1987; Kolmogorov, 1974). This uncomputable quantity is the complexity of the simplest model which minimizes the two part descrip- tion length of some data; it captures the amount of struc- tured, meaningful or "interesting" information contained the data (Vitányi, 2006). Sophistication is closely related to logical depth (Bennett, 1988; Antunes et al., 2017), an- other concept that seeks to quantify interestingness. Other views of interestingness relate it to discoverable compres- sion progress (Schmidhuber, 1999; 2009), or indeed as we do in this work- to information gain (Storck et al., 1995; Itti & Baldi, 2005; Herrmann et al., 2023).
Variational autoencoders (VAEs, Kingma & Welling (2014)) are generative models that combine variational Bayesian methods and autoencoder architectures. Their basis is the parametrization trick (Kingma & Welling, 2014; Rezende et al., 2014), which allows backpropagation through stochastic sampling operations. VAEs typically use fixed, non-informative priors, although structured priors that more closely resemble our autorregressive prior have been proposed (Johnson et al., 2016; Li & Mandt, 2018).
Various forms of information bottlenecks (Tishby & Za- slavsky, 2015) and autoencoders are widely used in mecha- nistic interpretability literature. Jiang et al. (2020) propose using an information bottleneck for token-level attribution of BERT models (Devlin et al., 2019). Recently, sparse autoencoders (Bricken et al., 2023) are proposed to decom- pose features in superposition (Elhage et al., 2022), that are common in large language models. (Henderson & Fehr, 2023) uses a VAE-based method to regularize the atten- tion (Bahdanau et al., 2015; Schmidhuber, 1992b; 1993) in transformers.
5. Limitations & Future Work
Although our proposed PHi loss shows a clear correlation to the intuitive notion of "interestingnes"" its precise connec- tion to formal concepts such as sophistication remains an open question. One particular issue is that of copying, or re- dundancy in general: When random sequences are repeated, the model must pass the information through the PHi bottle- neck in order to predict the second instance of sequence, leading increased PHi loss. An empirical demonstration of this effect can be found in Appendix C. However, under most definitions, such redundancy does not constitute true
In-context learning (Brown et al., 2020; Zhao et al., 2021; Liu et al., 2022) is a well-known ability of large language models that enables learning from examples presented as inputs during inference time without explicit training. This is in contrast with the standard paradigm of train and test
8

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
"interestingness.. Addressing this conceptual challenge- both theoretically and caapirrcally--mmiiin an important direction for future work.
ingful notions of task complexity, such as the description length of probabilistic finite automata, and aligns well with intuitive conceptions of complexity across various tasks, We anticipate that the PHi loss could serve as a powerful objective for applications that lack direct external feedback.
At present, the hidden states need many more bits to encode than the discrete tokens, most likely due to their high dimensionalityy. Exploring alternative forms of information bottlenecks (e.g., quantized autoencoders (van den Oord al., 2017)) could help bring the two losses onto more directly comparable scales. Moreover, the correlation between the next token prediction loss and the PHi loss should be explored further, possibly compensating for it when computing an aggregate score.
A Acknowledgements
We thank Eric Alcaide and Anand Gopalakrishnan for fruit- ful discussions. This research was funded by Swiss Na- tional Science Foundation grant: 200021 192356, project NEUSYM and the ERC Advanced grant no: 742870, Al- goRNN. We also thank NVIDIA Corporation for donat- ing DGX machines as part of the Pioneers of Al Research Award.
While we observe promising results by inserting a PHi layer into a pre-trained model, we rely on heuristics for deciding where in the network to place the bottleneck. Further empir- ical and theoretical studies are necessary to identify optimal bottleneck positions, and it would be valuable to train large models from scratch (or at least fine-tune them) so that they can fully accommodate the new information bottleneck. For simplicity, we aggregate the PHi loss over individual to- kens by taking the mean. This works well for fixed-size sequences, but for variable-length sequences, alternative aggregation methods (c.g., summation, top-k averaging, or thresholding by next-token loss) could also be considered. Finally, both "interestingness" and "task complexity" are inherently difficult to define. Consequently, evaluating our approach poses unique challenges. Developing new datasets with explicit complexity gradations- aligned to our intu- itive sense of interestingness remains an important avenue for future work.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
References
Akyürek, E., Wang, B., Kim, Y., and Andreas, J. In-context language learning: Architectures and algorithms. In Proc. Inf. Conf. Machine Learning (ICML), Vienna, Austria, July 2024.
Antunes, L.,, Bauwens, B., Souto, A., and Teixeira, A. So- phistication vs logical depth. Theory of Computing Sys- tems, 60:280-298. 2017.
A robust measure of interestingness could prove invaluable for agents that explore the world in an open-ended fash- ion (Lehman& & Stanley, 2008; Hughes et al., 2024). Many ultimately important tasks provide neither external rewards nor verifiable solutions. Future work should investigate us- ing the PHi loss of world model as an intrinsic reward signal for exploration in reinforcement learning, or as self-supervised objective for reasoners that must learn in the absence of external feedback.
Bahdanau, D., Cho, K., and Bengio, Y. Neural machine translation by jointly learning to align and translate. In Inf. Conf. on Learning Representations (ICLR), San Diego, CA, USA, May 2015.
Bennett, C. H. Logical depth and physical complexity. In The Universal Turing Machine: A Half Century Survey, pp. 227-258. Oxford University Press, 1988.
6. Conclusion
Blier, L. and Ollivier, Y. The description length of deep learning models. In Proc. Advances in Neural Infor- mation Processing Systems (NeurIPS), pp. 2220-2230. Montréal, Canada, December 2018.
We introduce the PHi loss, a novel information-theorett metric for quantifying the complexity or ""nterestingness" of in-context computation within neural sequence models. By augmenting standard architectures (e.g., Transformers or RNNs) with our self-predictive information bottleneck (the PHi layer), we can measure when models encode non-trivial structure in their hidden states. The PHi layer can be trained both from scratch and post-hoc (with pretrained models like Llama), allowing flexibility in how our method is used. Our experiments demonstrate that PHi loss correlates with mean-
Bricken, T., Templeton, A., Batson, J., Chen, B., Jermyn, A., Conerly, T., Turner, N., Anil, C., Denison, C., Askell, A., Lasenby, R., Wu, Y., Kravee, S., Schiefer, N., Maxwell, T., Joseph, N., Hatfield- Doddds Z., Tamkin, A., Nguyen, K., McLean, B., Burke, J. E., Hume, T., Carter, S., Henighan, T., and Olah, C. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023.

---

Mcasuring In-Contex, Computation Complexity via Hidden State Prediction
Brown, T. B. et al. Language models are few-shot learn- ers. In Proc. Advances in Neural Information Processing Systems (NeurIPS), Virtual only, December 2020.
Herrmann, V., Kirsch, L., and Schmidhuber, J. Learning one abstract bit time through self-invented experiments encoded as neural networks. In International Workshop on Active Inference, Ghent, Belgium, 2023.
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers solve math word problems. ArXiv, abs/21 10. 14168, 2021.
Hochreiter, S. and Schmidhuber, J. Long short- memory. Neural computation, pp. 1735-1780, 1997.
Hughes, E., Dennis, M. D., Parker-Holder, J., Behbahani, F., Mavalankar, A., Shi, Y., Schaul, T., and Rocktäschel, T. Position: Open-endedness is essential for artificial superhuman intelligence. Proc. Inl. Conf. on Machine Learning (ICML), Proceedings of Machine Learning Re- search. PMLR, 21-27 Jul 2024.
Deletang, G., Ruoss, A., Duquenne, P.-A., Catt, E., Ge- newein, T., Mattern, C., Grau-Moya, J., Wenliang, L. K., Aitchison, M., Orseau, L., Hutter, M., and Veness, J. Lan- guage modeling is compression. In Proc. Int. Conf. on Machine Learning (ICML), 2024.
Devlin, J., Chang, M., Lee, K., and Toutanova, K. BERT: pre-training of deep bidirectional transformers for lan- guage understanding. In Proc. North American Chapter of the Association for Computational Linguisties on Hu- man Language Technologies (NAACI-HILT), Minneapolis, MN, USA, June 2019.
Itti, L. and Baldi, P. F. Bayesian surprise attracts human attention. In Proc. Advances in Neural Information Pro- cessing Systems (NIPS), Cambridge, MA, 2005.
Jiang. Z.,, Tang, R., Xin, J., and Lin, J. Inserting information bottlenecks for attribution in transformers. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP), Virtual only, 2020.
Dubey, A. et al. The Llama 3 herd of models. ArXiv, abs/2407.21783, 2024.
Johnson, M. J., Duvenaud, D. K., Wiltschko, A. B., Adams, R. P., and Datta, S. R. Composing graphical models with neural networks for structured representations and fast inference. In Proc. Advances in Neural Information Processing Systems (NeurIPS), 2016.
Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodss, Z., Lasenby, R., Drain, D., Chen, C., Grosse, R., McCandlish, S., Kaplan, J., Amodei, D., Wattenberg, M., and Olah, C. Toy models of superposition. Transformer Circuits Thread, 2022.
Jordan, M. I. Attractor dynamics and parallelism in con- nectionist sequential machine. In Proc. Conf. of the Cog- nitive Science Society, Amherst, MA, USA, 1986.
Elman, J. L. Finding structure in time. Cognitive Science, 14(2):179-211, 1990.
Elmoznino, E., Marty, T., Kasetty, T., Gagnon, L., Mittal, S., Fathi, M., Sridhar, D., and Lajoie, G. In-context learning and occam's 's razor. ArXiv, abs/2410.14086.. 2024.
Kingma, P. and Welling, M. Auto-encoding variational bayes. In Int. Conf. on Learning Representations (ICLR), Banff, AB, Canada, April 2014.
Graves, A., Wayne, G., and Danihelka, I. Neural turing machines. ArXiv, abs/1410.5401, 2014.
Kirsch, L., Harrison, J., Sohl-Dickstein, J. N., and Metz, L. General General-purpose in-context learning by meta-learning transformers. ArXiv, abs/2212.04455 2022.
Grünwald, D. The Minimum Description Length Principle. Springer, 2007.
Kolmogorov, A. N. Talk at the information theory sympo- sium in tallinn. Estonia (then USSR), 1974.
Guo, Z. D., Pires, B. Á., Piot, B., Grill, J.-B., Altché, F. Munos, R., and Azar, M. G. Bootstrap latent- predictive representations for multitask reinforcement learning. ArXiv, abs/2004. 14646, 2020.
Koppel, M. Complexity, depth, and sophistication. Complex Systems, 1(6), 1987.
Henderson, and Fehr, F. A VAE for transformers with nonparametric variational information bottleneck. In Int. Conf. on Learning Representations (ICLR), Kigali, Rwanda, May 2023.
Lehman, J. and Stanley, K. O. Exploiting open-endedness to solve problems through the search for novelty. In Proc. Int. Conf. on the Simulation and Synthesis of Living Systems (ALIFE), 2008
Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. Mcasuring mathematical problem solving with the MATH dataset. In Proc. Advances in Neural Information Processing Sys- tems (NeurIPS), December 2021.
Lewkowyzz, A., Andreassen, A., Dohan, D., Dyer, E., Michalewski. H., Ramasesh, V. V., Slone, A., Anil, C., Schlag, I., Gutman-Sola T., Wu, Y., Neyshabur, B., Gur- Ari, G., and Misra, V. Solving quantitative reasoning problems with language models. In Proc. Advances in
10

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, November 2022.
recurrent nets. In Proc. Int. Conf. on Artificial Neural Networks (ICANN), pp. 460-463, 1993.
Li, Y. and Mandt, S. Disentangled sequential autoencoder. In Proc. Int. Conf. on Machine Learning (ICML), 2018.
Schmidhuber, Artificial curiosity based on discover- ing novel algorithmic predictability through coevolution. In Proceedings of the 1999 congress on evolutionary computation-cedt cec99 (cat. no. 99th8406), volume 3, pP- 1612-1618, IEEE, 1999.
Liu, J., Shen, D., Zhang, Y., Dolan, B., Carin, L.., and Chen, W. What makes good in-context examples for GPT-3? In Deep Learning Inside Out Workshop at ACL, May 2022.
Schmidhuber, J. Simple algorithmic theory of subjective beauty, novelty, surprise, interestingness, attention, cu- riosity, creativity, art, science, music, jokes. Journal of the Society of Instrument and Control Engineers, 48(1): 21-32, 2009.
Menon, S., Blei, D., and Vondrick, C. Forget-me-not! Con- trastive critics for mitigating posterior collapse. In Uncer- tainty in Artificial Intelligence, pp. 1360-13700 PMLR, 2022.
Ni, T., Eysenbach, B., Seyedsalehi, E., Ma, M., Gehring, C., Mahajan, A., and Bacon, P.-... Bridging state and history representations: Understanding self-predictive RL. ArXiv, abs/2401.08898, 2024.
Schmidhuber, and Prelinger, D. Discovering predictable classifications. Neural Computation, 5:625 635, 1993.
Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A. C.,, and Bachman, P. Data-efficient reinforcement learning with self-predictive representations. In Int. Conf. Learning Representations (ICLR), 2020.
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context learn ing and induction heads. Transformer Circuits Thread, 2022.
Storck, J., Hochreiter, S., and Schmidhuber, Reinforce- ment driven information acquisition in non-deterministic environments. In Proc. Int. Conf. on Artificial Neural Networks (ICANN), volume 2, 1995.
Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep gen- erative models. In Proc. Int. Conf. on Machine Learning (ICML), volume 32, pp. 1278 1278-1286, Beijing, China, June 2014.
Tishby, N. and Zaslavsky, N. Deep learning and the informa- tion bottleneck principle. 2015 IEEE Information Theory Workshop (I/W), 2015.
van den Oord, A., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. In Neural Information Processing Systems, 2017.
Schlag, I., Irie, K., and Schmidhuber, J. Linear transformers secretly fast weight programmers. In Proc. Int. Conf. on Machine Learning (ICML), 2021.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L.,, and Polosukhin, I. Attention is all you need. Proc. Advances in Neural Informa- tion Processing Systems (NIPS), Long Beach, CA, USA, December 2017.
Schmidhuber, Curious model-building control systems. In Proc. Int. Joint Conf. on Neural Networks, pp. 1458-1463, 1991a.
Schmidhuber, J. A possibility for implementing curios- ity and boredom in model-building neural controllers. In Proc. Int. Conf. on Simulation of Adaptive Behavior: From Animals to Animats, pp. 222-227, 1991b.
Vitányi, P. M. Meaningful information. IEEE Transactions Information Theory, 52(10), 2006.
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Ichter, B., Xia, F,, Chi, E. H., Le, Q. V., and Zhou, D. Chain-of- thought prompting clicits reasoning in large language models. In Proc. Advances in Neural Information Pro- cessing Systems (NeurIPS), New Orleans, LA, USA, November 2022.
Schmidhuber, J. Learning complex, extended sequences using the principle of history compression. Neural Com- putation, 4:234-242. 1992a.
Schmidhuber, J. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1))131-139, 1992b.
Zelikman, E., Wu, Y., Mu, J., and Goodman, N. D. Star: Bootstrapping reasoning with reasoning. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, November 2022.
Schmidhuber, J. Reducing the ratio between learning com- plexity and number of time varying variables fully

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
Zhao, Z., Wallace, E.,, Feng, S., Klein, D., and Singh, S. Calibrate before use: Improving few-shot performance of language models. In Proc. Int. Conf. on Machine Learning (ICML), Virtual only, July 2021.
12

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
A. Details on Experiments with Fully Trained Sequence Models
train both the Transformer and the LSTM model for 30, 000 steps using the Adam optimizer, batch size 16 and gradient norm clipping of 1.0. The learning rate 0.0003, with 500 step linear warm-up from zero and no decay. The standard next token prediction loss and PHi loss are both weighted with I, but for the PHi loss we take the mean of the element-wis wise KL-Divergence for not the sum. To prevent posterior collapse, employ contrastive self-critic loss (Menon et al., 2022) with weighting factor of 0.1.
Transformer Hyperparameters
The model is based the Llama 3.2 architecture (Dubey al., 2024).
Number of layers: 12
Model dimensionality: 768
Number of attention heads: 6
MLP intermediate size: 2048
Position of PHi layer: After layer
PHi Layer:
dimensionality: 768
qui Linear transform for and each
0g: Linear transform
Px One transformer layer like the ones in the rest of the model
LSTM Hyperparameters
standard LSTM architecture with linear input and output projections.
Number of layers: 2
Model dimensionality: 768
Position of PHi layer: After layer
PHi Layer:
dimensionality: 768
qw: Linear transform for μ and a each
ag: Linear transform
Px: One Llama transformer layer with dimensionality of 768, attention heads, and an intermediate MLP size of 2048
Training and Testing Data Generation
We follow Akyürek al. (2024) for the exact setup of the in-context langauge learning tasks. The parameters of the PFAs are randomly sampled from the following ranges:
Number of states: 12
Number of edges per state: 4
Vocabulary size: 4 18
Example length: 10 50
Each sequence consists of 10 20 examples from one PFA. To make training more robust, for half of the examples, we randomly perturb 20% of the tokens. During testing we use no perturbation.
As data for the four different tasks in experiment 3.1.1 we use:
Memorized Sequences task: 10 randomly generated examples
Memorized Programs task: Examples procedurally generated from one 10 fixed PFAs
In Context Language Learning task: Examples procedurally generated from a newly sampled PFA
Random Sequences task: Randomly sampled examples
13

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
Each sequence is uniformly sampled from of the tasks.
For experiment 3.1.2, train transformer with the same hyperparameters and training conditions as described above. However, here we train exclusively on data from Task 3. This is to avoid any biases due to the memorized PFAs from task 2.
B. Details on Experiment with Pre-Trained LLMs
The models involved in all of the experiments have the same hyperparameters and training data. We the pre-trained instruction-tunee Llama 3B model's parameters and only optimize the weights of the PHi layer. They are trained for 10, 000 steps using the Adam optimizer, batch size of and gradient norm clipping of 1.0. The learning rate is 0.0001, with no warm-up or decay. No contrastive self-critic loss is used. One model is trained with the PHi layer positioned after layer 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, or 28, respectively.
PHi Layer:
dimensionality: 3072
qu: Linear transform for y and each
dg: Linear transform
Px: One transformer layer like the ones in the rest of the model
The models are trained on mixture of natural language data from the SlimPajama dataset, the MATH training set and the GSM-8k training set.
Evaluation Data
In section 3.2.1, we use the following data:
Trivial tasks:
Write the multiplication table from 10.
List all the dates Jamuary 2025 with their day of the week.
Convert the numbers 1 50 into binary.
Write the numbers to 50 in Spanish.
Repeat the string 'RJGHDTSL!' 64 times.
Licenses tasks:
Generate the MIT license.
Generate the Apache 2.0 license.
Generate the GNU General Public License 2.0.
Random tasks: Examples from the SlimPajama dataset where the word order is randomly shuffled.
Literature tasks:
The beginning of 'Nigel Browning' by Agnes Giberne.
The beginning of Binney the Beaver' by Lucy Ellen Guernsey.
The beginning of "The Red-Hot Dollar' by H.D. Umbstaetter
The beginning of 'Hooking a Sky Ride' by Dan Morrissey.
The beginning of 'Flying Down a Rainbow' by Homer King Gordon.
The beginning of 'Boots story of the Sierra Paru' by Murray Leinster.
The beginning of "The Radio Cop' by Vic Whitman.
The beginning of 'The Trap' by Murray Leinster.
The beginning of 'The Cradle of the Deep' by Joan Lowell.
The beginning of 'Fix Bayonets!" by John W. Thomason, Jr.
Code tasks: Files from a private python code base that deals mainly with symbolic processing of music data.
14

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
C. Additional Results
C.1. PHi Loss and Copying
To highlight the conceptual issuc of interestingness and redundancy, we extend the experiments from Section 3.1.1 by introducing an additional task: copying. Similar to the random task, each sequence consists of previously unseen random subsequences. However, in this case, every subsequence appears twice within the same context. On its second occurrence, the model should be able predict it almost perfectly via context retrieval. To do that, the information has cross the PHi bottleneck, resulting in high PHi loss. Figures 10 and 11 show the next token loss and the PHi loss for fully trained models across all five tasks, including copying. For the transformer, the PHi loss in the copy task remains significantly lower than in in-context language learning. Still, it dehatable whether the copying task should be regarded as at all interesting. Note that this version of the copying task, with an exact 1-to-1 ratio of new random sequences to copies, maximizes the PHi loss. Increasing the proportion either copies or random data would reduce the PHi loss per token, as the total information passing through the bottleneck is constrained by the lesser of the two components.
3.0
Transformer LSTM
3.0
Transformer
LSTM
2.5
2.5
Relative PHi Loss
2.0
2.0
1.5
1.5
1.0
1.0
0.0
Mem. Seq. Mem. Prog.
ICCL
Random
Copy
Mem. Seq. Mem. Prog.
ICCL
Random
Copy
Figure Token-wise wise PHli loss versus binned next-token losses, stratified by PFA complexity (color). Normalized across each bin. Results from a fully trained Transformer model. Across all next- token prediction losses, complex PEAS result in higher PHi loss.
Figure I1. Token-wise loss versus binned next-token losses, stratified by PFA complexity (color). Normalized across each bin. Results from Llama 3B model with the PFA layer placed after layer 20. Also here, across most next-token prediction losses, more complex PFAS result a higher PHi loss.
15

---

Mcasuring In-Contex, Computation Complexity via Hidden State Prediction
C.2. PFA Task Examples and Normalized Complexity
Figure 12 shows conerete examples of in-context language learning tasks, along with memorized sequence, memorized programs and random data. Figures 13 and 14 show alternative versions of Figures 5 and Figure respectively. Here, the PHi losses normalized across each next token prediction loss bin.
PHi Loss
Next Token Loss
Memorized Sequences
Irit ontipmpprrrrmppriiccillliiiiliii pliirilinprp piliripil
Fllpilicii
(puippppprpmeeiicciiiiiiilmi piliiilinprp piliiipii
entbchhbbbanbrr ennnnbfififeif ennbennbgbanbr ennnnbli
ennbqunbqbęnor ennnnbrinrair ennbqunbgbqntt ennonbri
lbcl gggeoggbbbblb Hessddedooouagubed laafaedaadad
begesspeggbbllbbge Heponiednggutsgghet lqqfoqdaadqdi
Repriedprhbbb r'Shpddnnchccchhhhprrcccc
rckhbhkbl
reprhedprhonb Hkohppdahcckncchhhkkppppncc
rckhbhkbl
trinonppliiitrimppnpp pliirlinpppp onclirriiiipencnniilrriiii
jirnnppiliiniiiipppiii pliriinpppp ineliriiripnenemiiplrliii
abciinnnbbiieeddinaaaaa emnonarridilaaiiiiitmmnmii ebn
Abgnnnntfficciifffttedd emmmbrfreeiieeiilllfinnii ebno
Km hmdrhgmftmmdr kmhnddroghddfonb kmabnomhamateeh
Kmhmddhqmahmmat kmhmargghmarans kmqanbmhqmateq
nirg www.://nngnnnnnninnmmeeeeeeeeeee
nmpmgfmifrgmaett
hrgthmmnggrnmmettnnggggggggi napmafmiramaphh
Ebo iichinnbifiickiiifdddddi chhlitddiifiiffcccdddddaaaddd
enh
Dhgnnmmottraaeiiifaaddd eatelieeaimminnaaaarraaaaal enhr
mmmmmmgafnnnnn imtoggtrrpmammttiiimmmmmmmmmmm
nnfm
anfmnmgmgfinfn omiinpprrgeagmm/iirrmmmmmmmoommo nnfmi
Memorized Programs
X Iha mhhhhhhhh
Jelhahmhhhhhhn khalllhahhllhmhhhee httpl/ahhhhinnhhhhhhhhhaaaaamn
nelammmmmmmmnnnnnn khalthnhmaanm kinmhhhhh
alllmnhhnnnmmmmmmnn klhallhmhhhhhhhnn kihmhhhh
0a0a3aaaa d Cecririii jecteniiiiiiiiiin
esarronmnesseenn qkeqarmmmm aereinettttitreenn
od] opeas
Jnapr/rorj//rerkeelee (roe gojrear
Jepa (11) II PAT
Ihhh
thi
airekrmmil colrggnkoojiinii.. reeginojji) gg5ojrrikg)
Jirgkrm jirgri goirrgrrkrg peirrgrrrnijirrrrrggggnnnnnnnno 399girijk))
forrr goftrn/ girromkkrg firg wwg.:///w/.ooooggggggooooooooooo girgri
gireri
60000
eeec
In-Context Language Learning
nkkmömkdmnkpp hiddmh kr mh km Po mh kps m Tkkm
nkkmomme im kik hkdmhkdmhkkkm mmcc htt.m///ppgpkkkeee kpgmkdm hikmm
mikeeknkik II Mi-kik iikk
alll a
Ikhkk kllk
ag kc khkc
mikclk khke kknkkekkk kc] kejjkk dol kclkkikkke jjjl
oknjaassaeas Son ki abiiabt Jihb oaan Shbifhab ibibibibif
akaiea en Caaahhh hhiihbi oasakinbifhabil bbbibbibbbbbbbbi
Eh hihlhihlhlhhh hahhhhhhhhhhhhhhhhhhhhhhhhh
fihhhhhihi
e hihihihihhhii chihhhhhhhhhhhhhhhhhhhhhhhh
iliihhhhhl
i mookmokkooog podooooogdoe.... momoodbbg--o dllnd go cbm
c kkmmmook Tommmokkokg weoookkaogree woo. doommmgo aodd Edbgod booommm com
n h-ddniinn ndfnpnffnthhd hnthhhhrd 6dth dffnnd
nhd thohdr ndfn n nnfhnhhfhhn :aruthhfhn-fiii thphidhnd Indfhndfhndi
[d [eenilnnnnnnnahee (han' gheg inn linggghe dannogdk
Monnnnnnnnne lihgghgheg nn lyngghegdnnn qdk liken
id
kkkkhhhhkkkjjjjkkkkkkikk [Hkkjkbkkkhal dir ikb nkk
dd kRrkk kkkhikk dkhibx lidkhiki tikhikikkkkin kakhkhid in ikbkkbkdhnnki
hmoccceccec chcccereccceeccccc...
bocccc
hmocssccces
diceccccccceecceee
omeo
pcamgcma Pomamamam poc mammi mmgcmmccmmmec. Ps
p ppaamgccm cma
p:mamamam pcamamoi Imgem.cmm htt.mm/mmmmmcccmmmmmmmm
Random
Pnfric gondkingmkq! hadafhanskagccdddillmmhkkkkcccco
Kingmkal
da rec wwt.hciidccehhhnkkiiiittt gcddnffilmkpppelm Jin]ppk
trkkeerficcmmmeeeaaaajjnnmm gcddnfflikipellm J[fjjjk
ba kph mpammmminndoocccpppkcccccc mpnitenggedrkemgeqqukrrrr
dabi Inephhirpnnnonhho mlccngeddixemmseetoeerrii
en enhdijergj.
Thernmoophhpd
afhhc maammminddsskhiiikkeee.... thermmoophhp
hdhpoffkcankkkhhhhdd
hahpdfkcdnikkjhhhholll
cqn cqnrkbmfor Kobfkrrfrrr (orforro) i hnkkbem//jjboom/ccret rerbehconiigb agpnhihigorol
eqnrkbmfor workkkmffororrmmmmom for jjjjboomi crbohcanligb agpnhihigero
ii bic ddkidirkemcqiah hallssddppnichepdddhhppprr aagrbmo
ll kemeqlah aeljädddonichhejcdhhhpp..k aagrbmi
Ibabodapdmgel phnortthioggliccmmmmmmmmmmmeeeee
apniac doddpameel #koorffiogglcchhmmhmmmmmeeeeeeee
dgag minecpbrrtggrrannhhcccpppoo lmdgkcddrah]]
dgaged www.cooommmmmjjeamao memeeelbrrrraaoaaaabkeeehhhhhh Imdgkeddrahii
ripii galnkfpebrr Rechrchrfyynggkkknmjjorrrrrrrrrr
jibiid doccfmppeebrrr bectrdeftoookkkcmmpnooorrrrriiie
Figure 12. Examples of the four different tasks discussed in Section 3. colored by the PHIi and next token loss from Transformer model trained from seratch. We see clearly that the in-context language learning task leads higher PHi values, especially in the later parts of the sequences.
16

---

Mcasuring In-Contex, Computation Complexity via Hidden State Prediction
10
1.0
1.0
0.8
Relative PHI Loss
0.8
PFA Complexity
PFA Complexte
0.6
0.4
4
10-1
10a
10-2
10
10
Next Token Loss
Next Taken Loss
Figure 13. Token-wise wise PHi loss versus binned next-token losses, stratified PFA complexity (color). Normalized across each bin. Results from a fully trained Transformer model. Across all next- token prediction losses, complex PFAs result in higher PHi loss.
Figure 14. Token-wise PHi loss versus binned next-token losses, stratified by PFA complexity (color). Normalized across each bin. Results from Llama 3B model with the PFA layer placed after layer 20. Also here, more complex PFAS result in higher PHi loss.
next-token prediction losses,
C.3. Correct vs. Erroncous Rationales for the MATH Dataset
These results extend those presented in Section 3.2.3. Figure 15 shows the chance of choosing the correct answer for MATH dataset problems when picking the one with higher PHi loss, both across all answer pairs and for counterintuitive ones specifically. We use a chain-of-thought prompt but, following common practice for this dataset, provide no in-context examples. Once again, we find that responses with high PHi loss are significantly more likely to be correct. Compared to results on GSM-8k, selecting the response with the lowest next token loss is less predictive correctness (only ca. 52%). Figure 16 further breaks down the MATH dataset by difficulty. Notably, for counterintuitive questions. strong correlation emerges between high PHi loss and correctness in difficult problems, whereas this relationship is absent in easier ones. While this finding warrants deeper analysis, we propose the following explanation: Recall that define counterintuitive questions as those where the response with lower next-token loss is incorrect. Such questions may be counterintuitive for various reasons. One possibility is that they contain a tempting but ultimately incorrect simplistic solution. In this case, high PHi loss signaiing greater cognitive effort-may indicate that the model resisting this temptation. For difficult problems, this factor may dominate. Conversely, for easier problems, counterintuitivenss may stem from factual errors, which are not necessarily linked PHi loss.
65%
Chance of Choosine Correct Correct Answer
0.15
Partial Correlation
0.05
50%
0.00
All Questions Counterintuitive
-0.05
45%
18
20
Layer Position Llama Model
Figure 15. Similar to Figure 9, but for the MATH dataset. Chance of choosing the correct answer when selecting the one with the higher PHi loss between correct and wrong option. "Counterintuitive" are only those questions for which the answer with lower next token prediction loss is wrong, In purple, the partial correlation r between PIÍi loss and answer correctness- correctness-controlll controlled for next token prediction loss shown. Answers high PHii loss are clearly more likely to correct.
17

---

Mcasuring In-Contert Computation Complexity via Hidden State Prediction
PHi Layer after Layer 18
PHi Layer after Layer 20
655%
65%
Chance Choosing Correct Answer
60%
535%
0.05
53%
0.05
Corrictiion
50%
0.00
50%
85%
Highest
0005
45%
-0.05
Counterintuitive Question, Highest Piti Loss
4.0%
Question, Lowest
40%
Partial Correlation
00.0
-0.10
35%
33%
Difficully Level
Difficully Level
PHi Layer after Layer 22
PHi Layer after Layer 24
6.5
65%
60%
Choosing Correct
5.5%%
0.05
53%
0.05
50%
85%
-0.05
45%
-0.05
4.0%
-0.10 40%
-0.10
35%
35%
5
Difficulty Level
Difficulty Level
Figure 16. Detailed view of rationales generated for the MATH datasct, separated by the difficulty of the questions. Curves show (green) the chance of choosing the correct answer when selecting the one with highest PHi loss among answer pairs, (orange) among the counterintuitive pairs, (pink) among all pairs when selecting the answer with lowest next token loss, and (purple) the partial correlation between loss and answer correctness, controlled next token loss. We see a strong correlation between the correctness of the rationale and PHi loss, especially for difficult counterintuitive questions. For easy counterintuintuiiive question, this relationship does not exist.
C.4. Histograms for Next Token and PHi Loss
Figure 17 provides more detailed view of the results from Section 3.1.1 in the form of two-dimensional histograms, with next-token loss and PHi loss the axes. Figure 18 presents similar results for models trained also with the copying task. As shown in the top-right histogram, the Transformer model exhibits distinct cluster of copied sequences characterized by low next-token loss but high PHi loss. In contrast, the LSTM model does not display this cluster, likely because it struggles to adequately perform the copying task. Copying is known to be challenging for LSTMs see, for example, Graves et al. (2014).
Figure 19 presents analogous histograms for the Llama 3B model tasks from Section 3.2.1.
18

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
Transformer
Memorized Sequences
Memorized Programs
In-Coneext
Random
40
30
10"h
10.4
10P
10-4
10P
10
10t4
10P
10
10%
10
Toien
Next
Next Takken
LSTM
Sequences
Programs
Language Learning
60
00
30
10-1
10c
10
10
10P
10
10-
10
10-4
10-1
10
Figure 17. Two-dimensional histogram of tokens from the four different tasks presented in Section 3.1.1, PHi loss (y-axis) versus next token loss (x-axis). For memorized sequences, tokens are concentrated on the lower left, whereas for random data, they are clustered the lower right corner. Only in-context language learning yields significant amount of tokens that are high up the y-axis. 10,000 tokens per task from Transformer, and from an LSTM model.
Transformer with Copy In-Context Language Learning
Memorized Sequences
Memorized Programs
Random
Capy
50
30
10
10-²
10
10
10-1
109
10
109
10-2
10-
10-
30-1
Noc: Token Loss
Nex: Token Loss
Noxt Toxon Loss
Next Taken Loss
Next Token Less
LSTM with Copy
Memoriccdd Seq uences
Memorized Programs
In-Context Language Learning
Random
Copy
20
Nest: Toker
Token
Tolen
Figure 18. Similar Figure 17, but including the copying task. For the transformer model, that copied sequences have low next token loss, but relatively high PHi loss.
19

---

Mcasuring In-Context Computation Complexity via Hidden State Prediction
PHi Layer after Liama Layer 18
Trivial
Licenses
Rand
Literature
Code
1200
600
400
10
10-1
10
30a
10m9
10-1
15m4
10P
10
10
10-2
10a
10
10t11
10-4
10a
10
10a
10-2
10a
Neat Tolaen
Nect Talan Lown
Nec: Tolorn Lous
Niett Toben
Next: Tikenn
PHi Layer after Liama Layer 20 Rand om
Trivial
Licenses
Literature
Code
1200
800
000
200
10
10
107
10
10
15
10a
10
10
307
10
10
1.0
10
10
10
101
Nest Thklen Loos
Next Tiiken Loss
Next Thlen Loss
Next Taken Loss
Mext Takeen Loss
PHi Layer after Liama Layer 22
Trivial
Licenses
Random
Literature
Code
1200
800
600
200
10
10-1 Next Tokch Loos
10a5
10-2
100
10
10-2
10a
10a4
10
10a²
100
10:
10-2
10*
Nex: Tokon Loss
Naxt Tollen Loss
Noxt Tokon Loss
Noxt Tokion Loss
PHi Layer after Llama Layer 24
Trivial
Licenses
Random
Literature re
Code
800
402
200
Neat Tolum Loas
Noot:
Max:
Figure 19. Two-dimensional histograms of tokens from the five types of data described Section 3.2. For Literature and Code more token spreading to the upper, high PHi loss part of the histogram. Data shown for the PHi layer located after layers 18, 20, 22 and 24 of the Llama 3B model.
20


Reviews:
[{"review_id": "PMJcl0uxoT", "reviewer": "Reviewer_6ZwR", "summary": "The paper proposes a \"prediction of hidden states\" (PHi) layer, which can be used to quantify the complexity of the computation being performed in a neural model. The layer exists between the activations of a sequence model such as a Transformer. It maps the activations to latent variables. It then computes the KL between the encoding of the next symbol and learned \"prior\" that attempts to predict this encoding given the encodings of previous symbols. This KL divergence therefore quantifies the new information in the next symbol that is not predicted by the prior. \n\nThe experiments show that this KL divergence metric, termed the \"PHi loss\", correlates with various other measures of interest to a greater degree than the overall model's negative log likelihood (NLL) of the next token. For example, while the NLL is high for predicting random strings, \"PHi loss\" is relatively low. Intuitively, this is because the \"PHi loss\" is computed over the latent representations, which are implicitly encouraged to only contain information relevant towards predicting future symbols. Random sequences have high NLL because they have high conditional entropy. Experiments show that \"PHi loss\" correlates with the complexity of in-context learning tasks. Finally, experiments show that \"PHi loss\" over chain of thoughts correlates with accuracy of the final answer for GSM-8k, suggesting this is a desirable property for answer rationales.\n\n## Update after rebuttal\n\nThank you for your response and clarifying the relationship between the information that the posterior and prior condition on. I would recommend seeking to clarify the presentation and intuition around this if possible, in the revised version. I also appreciate further exploring the application of the method to select for rationales leading to correct answers. I will raise my score from 2 to 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Document with additional figures: https://tinyurl.com/yp4ucedn\n\n## 1. Rationales for Mathematical Reasoning\n> Perhaps most intriguingly, the proposed measure when applied to rationales for mathematical reasoning appears to correlate with answer accuracy.\n\nWe agree that this is an intriguing finding. Since the submission, we were able to further confirm and extend it, please see Figures 3&4 in the linked document. We repeated the experiment described in Section 3.2.1 on the MATH dataset and observed very similar results: Here, too—across all tested layers—correct rationales are associated with high PHi loss, both across all question pairs and within the subset of counterintuitive pairs. \nThe MATH dataset further allows us to break down the questions by difficulty. This reveals an interesting pattern within the counterintuitive subset: for easy questions, high PHi loss does not correlate with correct rationales, whereas for difficult questions, there is a strong positive correlation between high PHi loss and correctness. We will include these results in the updated version of the paper.\n\n## 2. Learned Prior and Capacity Constraints\nThe autoregressive prior lacks access to the most recent input token $x_t$. Hence, unless $x_t$ is completely uninformative, the PHi loss will not go to zero. We can see this by re-writing Equation 2 in the paper. Recall that $h_t$ is a function of $x_1, \\dots, x_t$ and $z_t$ is a (stochastic) function of $x_1, \\dots, x_t$. Plugging this in, we can express the PHi loss at time step $t$ as: $$\nL_{\\text{PHi}}(t)=D_{\\mathrm{KL}}\\Bigl( q_\\psi(\\cdot \\mid x_1, \\dots, x_t)\\big\\|p_\\chi(\\cdot \\mid x_1,\\dots,x_{t-1})\\Bigr).\n$$\nThis makes it clear that, in general, prior and posterior are not the same, even if the prior is highly expressive. \nThat said, the reviewer is right that in practice, the prior module $p_\\chi$ not only lacks access to the most recent token, but also has less capacity than the part $B_\\beta$ of the model that computes the target hidden state (the posterior distribution). In this sense, we are measuring the information gain resulting both from the new token *and* from the 'irreducible' computation the model performs in response to previous tokens. Disentangling these two components is an interesting challenge for future work.  \n\n## 3. Train/Test Split and Methodology\n> As there is no \"train and test split\" for most of the experiments...\n\nIt is important to clarify that there *is* a strict split between training and test/evaluation data for all experiments in the paper. The data used to train the models and PHi layers is distinct from the data used in the evaluations (with the exception of the memorization tasks, of course). We will emphasize this more clearly in the paper and provide additional information about train/test split in the appendix. As we describe in responses to other reviews, our method is fairly robust towards different hyperparameters (see response to u77y, paragraph 1 and Figure 5 in the linked document).\n\nWe agree with the reviewer that the results of Section 3.2 are somewhat contingent on the placement of the PHi layer in the Llama 3B model. However, we did our best to be transparent and methodically sound: Experiment 3.2.1 serves as an exploration and allows us to select sensible placements of the PHi layer. In response to Uw5E, paragraph 1, we detail the practical and theoretical considerations of this selection, which we cannot repeat here due to lack of space. All subsequent experiments are conditional only on this one selection process. We will make this clearer in the revised paper.\n\n## 4. PHi Loss and the Minimum Description Length Principle\nUnfortunately, we do not have the space in this reply for a full explanation, but we will clarify this connection in the updated paper. Here, for brevity we adopt the notation of https://arxiv.org/abs/2410.14086, where $p_\\theta$ is the in-context learned model of some data $D$. The MDL says we should minimize $K(p_\\theta) + K(D | p_\\theta)$. \nThe synthesized model $p_\\theta$ has to pass through the information bottleneck in order to be effective. The cumulative PHi loss quantifies the amount of information crossing that bottleneck (i.e., the information contained in the latent sequence), and is thus an upper bound of $K(p_\\theta)$. Our comment above about the prior, in addition to section 2.4 in paper which shows the existence of a compression scheme for the latent sequence, should make it clear that the autoregressive prior *is* accounted for in the PHi loss. Meanwhile, the cumulative NLL loss is an upper bound of $K(p_\\theta) + K(D | p_\\theta),$ and therefore also an upper bound of $K(p_\\theta)$. However—as we argue throughout the paper—it is a significantly weaker bound than the PHi loss.\n\nIf the reviewer would like further clarification, we are very happy to give more details and answer all remaining questions in a future response.", "claims_and_evidence": "TBD (see questions)", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "No theoretical claims.", "experimental_designs_or_analyses": "The experimental settings and analysis seemed reasonable.", "relation_to_broader_scientific_literature": "To the best of my knowledge, the proposed metric is novel. The paper discusses relevant prior work in Section 4.", "essential_references_not_discussed": "Not that I am aware of", "other_strengths_and_weaknesses": "Strengths\n\n* The proposed measure offers a new perspective on quantifying the complexity of the behavior of a neural model such as a Transformer.\n* The proposed measure correlates with description length of tasks in an intuitive way.\n* Perhaps most intriguingly, the proposed measure when applied to rationales for mathematical reasoning appears to correlate with answer accuracy.\n\nWeaknesses\n\n* I may be confused, but I was unclear regarding some aspects of the learned prior. Presumably, if the model used to represent the learned prior was sufficiently expressive, would it be possible for \"PHi loss\" to go to zero during training? Even at inference time, could the prior have equivalent complexity to the overall model, and therefore we would expect low \"PHi loss\" even on complex tasks, since both the overall model predicting the next token and the prior predicting the next hidden state can capture all of the same information (up to the limit of the mutual information between previous tokens and the next token)? In other words, in theory, does the method implicitly require some capacity constraint on the prior in order to be non-vacuous?\n* I worry that some of the results may be contingent on the specific decisions and hyperparameters used for implementing the encoder, prior, etc. As there is no \"train and test split\" for most of the experiments, it is difficult to assess the degree to which some conclusions may be contingent on these choices.\n* It would be helpful to have a bit more clarity on what precisely the authors mean by quantifying \"an upper bound on the complexity of this implicit model... that is generated in-context to predict next tokens\". The connection to work on description lengths, inspired by the MDL principle and Kolmogorov complexity, seemed potentially helpful but I didn't fully understand the connection. The role of the prior, the complexity of which is not accounted for in \"PHi loss\", seems to make the connection unclear. It seems like clarifying this would be quite helpful for the paper and its motivation. It could also help identify relevant methods from prior work to compare the proposed method against.\n\nI will be open to revisiting my judgements if the authors can help clarify these points for me.", "other_comments_or_suggestions": "Just some minor nits:\n* Section 2.3 typo - \"to ensure that ensures that\"\n* Section 6 typo - \"a powerful objective for in applications\"", "questions_for_authors": "See \"Weaknesses\" above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hAB23tmJqM", "reviewer": "Reviewer_u77y", "summary": "I think this paper introduces the PHi (Prediction of Hidden States) layer as a novel way to measure the complexity of computation performed by neural sequence models by examining how predictable their hidden states are. The authors show that this metric correlates better with intuitively \"interesting\" computation than next-token prediction loss.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Document with additional figures: https://tinyurl.com/yp4ucedn\n## 1. Robustness towards Hyperparameters and PHi Layer Placement \n>...the paper lacks rigorous justification for where to place the PHi layer in pre-trained models\n\nWhile it is true that the properties of the PHi layer vary depending on its placement in a pre-trained LLM, our choices are not arbitrary. In our response to Uw5E, paragraph 1, we give several principled reasons—independent of any specific results on interestingness or computational complexity—why the chosen positions are sensible for the PHi layer. We highlight that we selected layers that maximize the variance of PHi loss across different tasks. A more detailed account of these considerations will be included in the updated paper.\n\n\n> How sensitive are the results to the specific architectural choices in the PHi layer?\n\nIn general, our method is quite robust to the choice of hyperparameters and the exact architecture of the PHi layer components. For example, the sizes of the MLPs and the attention module are chosen simply to align with the rest of the model architecture. We show in the paper that our method works with transformers, LSTMs, and pre-trained LLMs.\n\nWe did not perform extensive hyperparameter tuning, but selected reasonable and stable settings and used them consistently. In response to this review, we tested varying the weight of the PHi loss (relative to NLL loss) from 1 to 10, and found that higher weights actually perform better than the one reported in the paper, see Figure 5 in the linked document. If the reviewer is interested, we would be happy to provide a more detailed robustness analysis for specific hyperparameters in a future response.\n\n\n## 2. Subjectivity of Interestingness\n> The evaluation relies heavily on intuitive notions of complexity rather than formal definitions.\n\nThe reviewer raises a valid point: Notions of complexity and interestingness are inherently difficult to formalize. Many definitions—such as Kolmogorov Complexity and Sophistication—are non-computable and cannot be directly used in empirical studies. Some appeal to intuition is therefore difficult to avoid.\n\nHowever, we have attempted to ground our evaluation in formal and objective measures wherever possible. For the PFA-based tasks, we use the description length of the automaton $C(A)$, as described in section 3.1, to quantify complexity. For the complexity of the non-ICLL tasks, we can also give a concrete description length: the complexity of memorized sequences is $\\log N$, where $N$ is the number of memorized sequences. For memorized programs, it is $\\log M$, where $M$ is the number of memorized programs. For random sequences, the complexity is 0, as no inference is possible. \n\nWe agree that the natural language tasks in Section 3.2.1 are categorized in a less formal way, but believe our reasoning is principled and goes beyond intuition. In the MATH dataset experiments, we rely on dataset-provided difficulty levels, which range from simple to difficult. In this context, we take “difficult” to be synonymous with “complex.”\n\nWe also want to emphasize that we never claim to measure the *absolute* or objective interestingness of a sequence—only the *relative* differences between sequences or tasks for a given model.\n\nWe will further clarify these points in the updated paper.\n\n\n## 3. Misunderstanding about Experiments with Large Models\n> The major concern is that it is only tested on small model, can we test this on a BERT/GPT2 or even bigger model such as Llama1B etc,\n\nThe paper includes extensive experiments with large models—specifically, the LLaMA 3.2B model with 3 billion parameters, which is significantly larger than BERT, GPT-2, or LLaMA-1B. Please see Sections 3.2.1, 3.2.2, and 3.2.3. \n\nOur method is also fairly scalable. Training the PHi layer in the pre-trained LLaMA 3B model takes approximately 3 hours on a single consumer-grade GPU (NVIDIA RTX 3090). Upon acceptance, we will release our code, which requires only minor modifications to work with other open-source LLMs.\n\n\n## 4. Additional Points\n> ...baseline comparisons are limited.\n\nWe are not aware of other methods that directly measure the complexity of in-context computation in sequence models. Does the reviewer have any specific baselines in mind? The two alternative approaches discussed in our response to rCDj, paragraph 2, may be of interest.\n\n\n> The application to correctness prediction in self-generated reasoning chains demonstrates practical utility.\n\nWe are glad the reviewer appreciates this. The updated paper will include an expanded experiment that confirms and extends this result on the MATH dataset (see response to 6ZwR, paragraph 1, and Figures 3&4 in the linked document).\n\n\n> ...could better situate itself within the growing body of work on understanding LLM computation\n\nWe have aimed to highlight the major relevant work and would be grateful for any suggestions of related literature we may have missed.", "claims_and_evidence": "the evidence provided through experiments on both smaller models and LLMs is adequate but not completely convincing, as the connection to formal notions of complexity remains mostly theoretical.", "methods_and_evaluation_criteria": "I think the proposed PHi layer is an elegant approach, but its placement in pre-trained models seems somewhat arbitrary, raising questions about the robustness of the method.", "theoretical_claims": "n/a", "experimental_designs_or_analyses": "I think the experiments cover a good range of scenarios, but the evaluation of \"interestingness\" remains subjective, and the baseline comparisons are limited.", "relation_to_broader_scientific_literature": "The paper makes interesting connections to information theory and mechanistic interpretability, but could better situate itself within the growing body of work on understanding LLM computation", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "I think the proposed metric provides a novel lens for quantifying when models perform meaningful computation. The correlations between PHi loss and task complexity across different domains are intriguing. The application to correctness prediction in self-generated reasoning chains demonstrates practical utility.\n\n\nHowever, the paper lacks rigorous justification for where to place the PHi layer in pre-trained models, with results varying significantly based on layer position. The evaluation relies heavily on intuitive notions of complexity rather than formal definitions. The method seems sensitive to hyperparameters and architectural choices, raising questions about generalization. \n\nThe major concern is that it is only tested on small model, can we test this on Llama7/8+B, etc?", "other_comments_or_suggestions": "n/a", "questions_for_authors": "1) How sensitive are the results to the specific architectural choices in the PHi layer? \n\n2) Is there a principled way to determine optimal layer placement in pre-trained models?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AKkCKYJaA0", "reviewer": "Reviewer_rCDj", "summary": "This paper proposes a novel method for probing the hidden representations of neural sequence models: the prediction of the hidden states (PHi) layer. This layer combines an encoder that generates latent variables from the hidden states and an autoregressive LMs that generates latent variables from previous latent variables. The corresponding parameters are trained to minimise the KL divergence between encoder and autoregressive LMs, together with the standard autoregressive objective on the output of the whole model. The authors test their new measure on several tasks, including both cases where the neural network is trained together with the PHi layer from scratch and cases where the PHi layer is inserted into a pre-trained LLM. The tests show how the PHi loss correlates with the computational complexity of the task while the standard next-token prediction loss does not.\n\n## update after rebuttal\n\nThe rebuttal cleared the confusion expressed in the \"Other Strengths and Weaknesses\" section, and I have increased the score accordingly.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Document with additional figures: https://tinyurl.com/yp4ucedn\n\n## 1. Questions about the PHi Layer\n> 1: is the autoregressive part of the PHi layer used to compute the output, or does it just enter to train the encoder?\n\nThe purpose of the autoregressive part (i.e., the causal self-attention layer) is to make the prior $p_\\chi$ more powerful. It plays no role in computing the next token prediction of the model. More details below.\n\n\n> 2: is eq. (4) implying that the first part of the networks does not receive gradients from the negative log-likelihood loss?\n\nNo—the gradients flow back through the reconstructed $h'$, to the encoder $q_\\psi$, the original hidden states $h$ and to the bottom layers $B_\\beta$. No gradients are blocked in the architecture. This ensures that when trained end-to-end, as in Section 3.1, the model is encouraged not only to *predict* its own next hidden states, but also to make those next hidden states more *predictable*.\n\n\n> 3: how is the encoder incentivised to have the latents depend on the previous latents, given that the autoregressive component of PHi is also initialised randomly and trained from scratch? (see also Strengths and Weaknesses section)\n\nThe latents $z$ are sampled from the posterior $q_\\psi(\\cdot | h_t) = q_\\psi(\\cdot | x_1, \\dots, x_t)$. They have to contain the information necessary to predict the next token $x_{t+1}$, otherwise the NLL loss increases. The prior has to predict the posterior distribution but does not have access to the most recent input token $x_t$. However, we make it autoregressive and give it access to the previous latents, since their information content has already been accounted for by previous PHi losses.\n\nIn short: the autoregressive prior $p_\\chi$ is trained to predict the posterior. The posterior $q_\\psi$ is trained to allow accurate next token prediction from the reconstructed state $h'$, while staying as close as possible to what the prior predicts.\n\n\n## 2. Two Simpler Approaches and Why they Do not Work\n> ...if the prior were some fixed distribution.\n\n> ...why not train the prior with a standard next-token prediction objective but in the latent (instead of token) space?\n\nThe reviewer raises two interesting questions which we also explored during the development of our method:\n- Can we simply use an information bottleneck with a fixed, uninformative prior and no hidden state prediction (similar to a vanilla VAE, no autoregressive prior)?\n- Can we use a simpler next hidden state prediction objective? Since the latent or hidden state is continuous, it cannot be the usual categorical cross-entropy, but maybe mean squared error (MSE)?\n\nEach of these simpler approaches has shortcomings that our method avoids:\nIf we use a fixed prior in a pre-trained LLM, it tends to overestimate the information in the hidden state sequence. For example, consider a sequence where the hidden states all lie in the same small subspace. With a fixed prior, we would have to pay the same large price in terms of KL-divergence for every single one of the hidden states. An autoregressive prior, in contrast, quickly adapts to the subspace, reducing the KL for subsequent steps. Our new results include an ablation showing that an information bottleneck without hidden state prediction gives no meaningful results in practice (see Figure 6 in the linked document).\n\nAs for using a straightforward next hidden state prediction loss such as MSE: this encourages the model to scale down the norm of the hidden states to near-zero, potentially leading to machine precision issues. Even if we fix the norm, the model could encode information in tiny perturbations from a default vector, which lead to low MSE but can carry arbitrary information. Without a noisy information bottleneck, the approach lacks a clear information-theoretic interpretation.\n\nOur PHi layer addresses both issues:  \n- The autoregressive prior ensures we pay only for genuinely *new* information not already contained in previous tokens.  \n- The noisy information bottleneck prevents hidden state collapse and provides a clear, principled measure of the information in the latent states.\n\n\n## 3. Statistical Significance of the Results\n\nWe emphasize that we report 95% confidence intervals or p-values for all experiments, and find all reported results to be statistically highly significant. See Figure 2 in the linked document for additional confidence intervals for the partial correlation in the MATH rationale experiment. These were missing in the submitted paper.\nFor a discussion of how we selected the PHi layer location, please refer to our response to reviewer Uw5E, paragraph 1.\n\nIf there are any remaining questions, we are very happy to answer them in a future response.", "claims_and_evidence": "It is far from obvious why the proposed procedure should measure \"a model's ability to predict its own hidden states\", as stated in the abstract. Nevertheless, experimental results agree with the trends claimed by the authors at the end of section 1. The statistical significance of these trends remains open for debate, especially for experiments with LLMs (figures 6 and 7).", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "There are no theoretical claims.", "experimental_designs_or_analyses": "Experiments seem sound although there are aspects of the PHi layer design that I do not understand:\n\n1: is the autoregressive part of the PHi layer used to compute the output, or does it just enter to train the encoder?\n\n2: is eq. (4) implying that the first part of the networks does not receive gradients from the negative log-likelihood loss?\n\n3: how is the encoder incentivised to have the latents depend on the previous latents, given that the autoregressive component of PHi is also initialised randomly and trained from scratch? (see also Strengths and Weaknesses section)", "relation_to_broader_scientific_literature": "This paper could have a significant impact on the literature devoted to understanding the operations performed by LLMs, including mechanistic interpretability and other studies of hidden representations.", "essential_references_not_discussed": "I am not aware of any.", "other_strengths_and_weaknesses": "The problem considered by the paper is timely and relevant and the solution based on the idea of measuring \"how well the model can predict its own future hidden states\" is both interesting and original. However, I cannot completely follow the rationale behind the design of the PHi layer. I would understand the comment after eq. (5), \"the model is incentivized to maintain ...\", if the prior were some fixed distribution. However, the fact that the prior is trained from scratch via Eq. (2) confuses me, as the loss only forces it to agree with the encoder distribution and not with the distribution of post-activations from the training data. For instance, why not train the prior with a standard next-token prediction objective but in the latent (instead of token) space?", "other_comments_or_suggestions": "Typo in section 2.3, third row (\"we need to ensure that ensures\").\n\nTypo in the caption of Figure 5 (\"because of a posterior collapse, most likely due to posterior collapse\").", "questions_for_authors": "I have no further questions. My `weak accept' evaluation stems from the confusion expressed in the \"strengths and weaknesses\" section: if cleared, I would raise the evaluation. --- updated after response.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FqS7dwQtBh", "reviewer": "Reviewer_Uw5E", "summary": "The paper introduces the Prediction of Hidden States (PHi) loss to measure the complexity of computation in neural sequence models. The authors argue that traditional next-token prediction loss does not adequately capture the task complexity. To fix this, they propose evaluating the model’s ability to predict its own future hidden states. The PHi layer is introduced to measure the unpredictability of hidden states, which correlates with task complexity. The method is tested across different tasks and architectures, demonstrating that PHi loss distinguishes between complex and simple computations.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Document with additional figures: https://tinyurl.com/yp4ucedn\n## 1. Location of the PHi Layer in the Model\n> The PHi layer's effectiveness depends on where it is inserted in the model.\n\nWe want to emphasize that in a fully trained model—such as the transformer in Section 3.1—the results are very robust across different placements of the PHi layer. This is supported by new experiments (Figure 5 in the linked document).\n\nWith regard to pre-trained LLMs, the reviewer is correct that the effectiveness of the PHi layer in measuring in-context computation complexity depends on roughly where in the model it is placed. Based on known findings from LLM interpretability research and other empirical considerations, it is not surprising that PHi layers placed about two-thirds of the way through the model tend to perform best. Below we detail our reasoning for choosing placing the PHi layer after layer 18, 20, 22 and 24 of the Llama 3B model in the subsequent experiments.\n- https://arxiv.org/abs/2407.09298 shows that the early and final layers in an LLM serve distinct purposes, making them unsuitable for measuring in-context computation.\n- https://arxiv.org/abs/2410.10912 demonstrates that early layers in LLMs are more brittle and harder to prune—making them more vulnerable to the noise introduced by the PHi layer’s information bottleneck.\n- PHi layers placed early in the model show clear signs of posterior collapse and thus give no clear results.\n- We selected positions based on where the *variance* of the PHi loss across tasks is highest—not where interesting tasks produce the largest losses (although these two  happen to coincide).\n\nThese choices are independent of any specific findings about which tasks are associated with high or low PHi loss. We therefore do not believe our selection introduces bias in the experimental outcomes. We will elaborate on this reasoning in the updated paper.\nA practical benefit of later PHi layer placement is that it reduces computational cost, since we backpropagate through fewer layers of the pre-trained LLM.\n\n## 2. Further Points\n\n> Although promising for small models, computing the PHi loss requires adding a PHi layer and training the model. This could make it difficult to apply on larger-scale LLMs.\n\nIn pre-trained LLMs, only the PHi layer is trained—this requires very modest compute. Training the PHi layer for 10,000 steps in the Llama 3B model, as done for all experiments in Section 3.2, takes between 2.5 and 6 hours on a single consumer-grade GPU (NVIDIA RTX 3090), depending on the PHi layer’s position. We see no reason this approach couldn't scale to larger models.\n\n> Testing PHi loss on encoder-only models (e.g., BERT) and encoder-decoder models (e.g., T5) could provide further insights.\n\nWe agree—measuring in-context computation complexity in these models would be interesting. However, it likely requires conceptual changes to the PHi layer. In decoder-only models, we have a clear compression scheme for the token and latent sequences. In encoder-only models, hidden states can incorporate information from both past and future tokens, complicating the setup.\n\n> How sensitive is PHi loss to model size and dataset scale?\n\nThere *is* a dependency between PHi loss and model size: the KL divergence between two diagonal Gaussians is a sum over dimensions, so the PHi loss scales with the dimensionality of the hidden states. For this reason, PHi loss (in its current form) should only be used to compare tasks within the same model, not across different models. We never claim to compare models in the paper. We also do not interpret the absolute magnitude of the PHi loss. Addressing this limitation is a goal for future work—for example, through the use of quantized latent states.\n\nAs for the PHi layer’s capacity and training data needs, our method is flexible. In the Llama 3B experiments, we trained the PHi layer on roughly 40 million tokens from a standard natural language dataset. Figure 7 in the linked document shows the training curves. The PHi layer's architecture is designed to match the scale of a single LLM layer. If the reviewer finds it helpful, we are happy to provide additional experiments comparing different PHi layer capacities and training data volumes in a future reply.\n\n> Is it possible to combine the PHi loss with other model analysis tool to pinpoint where in-context learning occurs?\n\n**Where in the sequence or dataset?** Although we focus on the aggregated PHi loss in the paper, the method supports token-wise PHi loss measurement. This allows us to pinpoint where in a sequence non-trivial in-context computation occurs (e.g., see Figure 13 in the Appendix).\n\n**Where in the model?** Figures 6&7 in the paper offer some preliminary insights into where interesting computation may be happening within a pre-trained LLM. However, further research is needed to confirm and extend these findings.", "claims_and_evidence": "* Next-token prediction loss is an unreliable indicator of computational complexity.\n    - Evidence: Tasks involving random token sequences yield high next-token prediction loss but have no meaningful computation.\n* Hidden-state predictability (PHi loss) is a better metric for measuring in-context computation complexity.\n    - Evidence: The authors show that PHi loss increases for complex reasoning tasks, in-context language learning, and step-by-step mathematical reasoning.\n* PHi loss correlates with the description length of formal languages learned in-context.\n    - Evidence: The study demonstrates that the PHi loss reflects the complexity of probabilistic finite automata (PFA)-based tasks.", "methods_and_evaluation_criteria": "* Next-token prediction loss is an unreliable indicator of computational complexity.\n    - Evidence: Tasks involving random token sequences yield high next-token prediction loss but have no meaningful computation.\n* Hidden-state predictability (PHi loss) is a better metric for measuring in-context computation complexity.\n    - Evidence: The authors show that PHi loss increases for complex reasoning tasks, in-context language learning, and step-by-step mathematical reasoning.\n* PHi loss correlates with the description length of formal languages learned in-context.\n    - Evidence: The study demonstrates that the PHi loss reflects the complexity of probabilistic finite automata (PFA)-based tasks.", "theoretical_claims": "No theoretical claim was made in this work.", "experimental_designs_or_analyses": "* Insert the PHi layer into different models to evaluate.\n* train-from-scratch models: Transformers and LSTMs\n* Pre-trained LLMs: a frozen Llama 3B model", "relation_to_broader_scientific_literature": "No direct relation to broader scientific literature.", "essential_references_not_discussed": "The paper extensively cites prior work in in-context learning, VAEs, information bottlenecks, and neural sequence models.", "other_strengths_and_weaknesses": "**Strengths**\n\n* PHi loss gives a new way to evaluate the task complexity.\n* The method is tested on multiple architectures and datasets.\n\n**Weaknesses**\n\n* The PHi layer's effectiveness depends on where it is inserted in the model.\n* Although promising for small models, computing the PHi loss requires adding a PHi layer and training the model. This could make it difficult to apply on larger-scale LLMs.", "other_comments_or_suggestions": "* Testing PHi loss on encoder-only models (e.g., BERT) and encoder-decoder models (e.g., T5) could provide further insights.", "questions_for_authors": "* How sensitive is PHi loss to model size and dataset scale?\n* Is it possible to combine the PHi loss with other model analysis tool to pinpoint where in-context learning occurs?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a new approach prediction of hidden states (PHi) to measure the interestingness of neural sequence models' computations. After discussion, reviewers unanimously agreed that this is a novel and interesting tool that can provide new insights. Some concerns were only partially resolved like dependence on where to put the PHi layer, and evaluation on bigger models. But the paper is already a valuable contribution as-is. I recommend accept.
Decision: Accept (poster)