Research map / Architectures and efficiency / Attention and transformers
Measuring In-Context Computation Complexity via Hidden State Prediction
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes the PHi layer, an information bottleneck with a learned autoregressive prior that measures the KL-based description length of a model's hidden states as a metric of in-context computation complexity.
Problem
Next-token prediction loss does not reliably indicate whether a sequence model performs interesting or complex in-context computation, since low loss can come from trivial sequences and high loss from irrelevant randomness.
Approach
Insert a PHi layer (VAE-style information bottleneck with a learned autoregressive prior over latents) into the residual pathway of a sequence model and use the per-token KL divergence between posterior and prior, i.e. novel information in hidden states not predictable from the past, as the interestingness metric; train jointly with NLL loss or post-hoc on a frozen pretrained LLM.
Key results
- On in-context PFA learning, mean sequence PHi loss partially correlates with PFA complexity (r=0.37, 95% CI [0.32, 0.43]) after controlling for next-token loss.
- On MATH with Llama 3B, PHi loss partially correlates with problem difficulty (r=0.079, 95% CI [0.07, 0.09]) across tested layers.
- On GSM8K self-generated rationales, choosing the answer with higher PHi loss gives a significantly increased chance of correctness, while choosing lower next-token loss gives 71%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Measure hidden-state predictability with PHi loss rather than relying on next-token loss | Transformer in PFA language learning · PFA language learning | partial correlation with PFA complexity: r=0.37, with a 95% confidence interval of [0.32, 0.43] | Helped |
| Measure hidden-state predictability with PHi loss | Llama 3.2 3B on MATH solutions · MATH | partial correlation with problem difficulty: r=0.079, with a 95% confidence interval of [0.07, 0.09] | Helped |
| Select the rationale with higher PHi loss instead of choosing at random | Llama 3.2 3B generated GSM8K rationales · GSM8K | correct-answer selection rate | Helped |
| Select the rationale with lower next-token loss | Llama 3.2 3B generated GSM8K rationales · GSM8K | correct-answer selection rate: 71% | Helped |
| Insert the PHi layer in early rather than later Llama layers (ablation) | Llama 3.2 3B · Trivial, license, random, literature, and code sequences | next-token prediction loss | Hurt |
| Place the PHi layer after layers 18–24 rather than in early layers (ablation) | Llama 3.2 3B · Trivial, license, random, literature, and code sequences | next-token prediction loss | Helped |
| Compare PHi loss across task types | Transformer and LSTM trained from scratch · Memorized subsequences, memorized programs, PFA language learning, and random examples | PHi loss | Mixed |
| Compare PHi loss across task types | Llama 3.2 3B · Trivial, memorized licenses, random words, literature, and code | PHi loss | Helped |
| Add repeated unseen subsequences that must be copied from context (ablation) | Fully trained Transformer · Copying task | PHi loss | Hurt |
Surprising findings
- Choosing the rationale with lower mean next-token loss yields 71% correctness on GSM8K pairs, higher than the PHi-based choice, although PHi and next-token loss are generally positively correlated.
- Inserting the PHi layer in early layers of Llama 3B causes posterior collapse with large next-token loss degradation, whereas later layers avoid this.
Limitations
- Hidden states need many more bits than discrete tokens, so PHi and NLL losses are on different scales; the metric is not directly comparable.
- Redundancy such as copying increases PHi loss, which is not necessarily interesting computation.
- Bottleneck placement is chosen heuristically for pretrained models; aggregation by mean is simplistic for variable-length sequences.
Experimental setup
Datasets: SlimPajama MATH GSM8K synthetic probabilistic finite automata data
Benchmarks: MATH GSM8K
Metrics: PHi loss next token loss partial correlation pairwise selection accuracy
Baselines: next token prediction loss
Builds on: variational autoencoder reparameterization trick in-context language learning prequential coding
Closest papers by meaning
- Auto-Regressive Next-Token Predictors are Universal Learners (ICML 2024 · similarity 0.88)
- Multiple Token Divergence: Measuring and Steering In-Context Computation Density (ICLR 2026 · similarity 0.86)
- Limits of Deep Learning: Sequence Modeling through the Lens of Complexity Theory (ICLR 2025 · similarity 0.86)
- Unifying Attention Heads and Task Vectors via Hidden State Geometry in In-Context Learning (NeurIPS 2025 · similarity 0.86)
- Transformers Represent Belief State Geometry in their Residual Stream (NeurIPS 2024 · similarity 0.86)
- Learning to (Learn at Test Time): RNNs with Expressive Hidden States (ICML 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.