# Measuring In-Context Computation Complexity via Hidden State Prediction

arXiv: https://arxiv.org/abs/2503.13431 · HTML: https://mlsanity.com/paper/2503.13431/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 54 references that are in the map.
- First posted on arXiv 17 Mar 2025.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Calibrate Before Use: Improving Few-Shot Performance of Language Models](https://mlsanity.com/paper/2102.09690/) (2021)
- [Solving Quantitative Reasoning Problems with Language Models](https://mlsanity.com/paper/2206.14858/) (2022)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [STaR: Bootstrapping Reasoning With Reasoning](https://mlsanity.com/paper/2203.14465/) (2022)
- [Deep Learning and the Information Bottleneck Principle](https://mlsanity.com/paper/1503.02406/) (2015)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Linear Transformers Are Secretly Fast Weight Programmers](https://mlsanity.com/paper/2102.11174/) (2021)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Language Modeling Is Compression](https://mlsanity.com/paper/2309.10668/) (2023)
- [Disentangled Sequential Autoencoder](https://mlsanity.com/paper/1803.02991/) (2018)
- [Composing graphical models with neural networks for structured representations and fast inference](https://mlsanity.com/paper/1603.06277/) (2016)
- [General-Purpose In-Context Learning by Meta-Learning Transformers](https://mlsanity.com/paper/2212.04458/) (2022)
- [Bootstrap Latent-Predictive Representations for Multitask Reinforcement Learning](https://mlsanity.com/paper/2004.14646/) (2020)
- [In-Context Language Learning: Architectures and Algorithms](https://mlsanity.com/paper/2401.12973/) (2024)
- [Open-Endedness is Essential for Artificial Superhuman Intelligence](https://mlsanity.com/paper/2406.04268/) (2024)
- [The Description Length of Deep Learning Models](https://mlsanity.com/paper/1802.07044/) (2018)
- [Bridging State and History Representations: Understanding Self-Predictive RL](https://mlsanity.com/paper/2401.08898/) (2024)
- [Forget-me-not! Contrastive Critics for Mitigating Posterior Collapse](https://mlsanity.com/paper/2207.09535/) (2022)
- [Learning One Abstract Bit at a Time Through Self-Invented Experiments Encoded as Neural Networks](https://mlsanity.com/paper/2212.14374/) (2022)

## Built on by (most cited first, 3 of 3)

- [Thinking Machines: Mathematical Reasoning in the Age of LLMs](https://mlsanity.com/paper/2508.00459/) (2025)
- [Interestingness as an Inductive Heuristic for Future Compression Progress](https://mlsanity.com/paper/2605.14831/) (2026)
- [Multiple Token Divergence: Measuring and Steering In-Context Computation Density](https://mlsanity.com/paper/2512.22944/) (2025)

## Abstract (arXiv)

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.