PaperlayerView exact revision on arXiv

Source-linked AI summary

Measuring In-Context Computation Complexity via Hidden State Prediction

Vincent Herrmann, Róbert Csordás, Jürgen Schmidhuber

arXiv:2503.13431v1cs.LG

TL;DR

The paper addresses how to detect interesting computation in neural sequence models when next-token loss is an inadequate indicator. It introduces hidden-state predictability through an architecture-agnostic PHi information bottleneck and learned predictive prior. PHi loss correlates with formal-language complexity, mathematical reasoning difficulty, and the correctness of self-generated reasoning chains.

  • Problem

    Detecting when neural sequence models perform interesting computation is difficult because next-token loss does not reliably reflect underlying computational complexity.

  • Method

    The PHi layer bottlenecks and autoregressively predicts hidden states, while a learned prior measures novel information through posterior-prior KL divergence.

  • Results

    PHi loss correlates with formal-language description length, mathematical problem complexity, and the correctness of self-generated reasoning chains.

  • Takeaways & Limitations

    Hidden-state unpredictability provides a principled tool for detecting interesting in-context behaviors across neural sequence-modeling tasks.

  • Takeaways & Limitations

    PHi loss can increase for repeated random sequences, so its precise connection to formal notions of interestingness remains unresolved.

Abstract

from arXiv · show

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.

1. Introduction

The paper argues that next-token loss poorly measures interesting computation and introduces hidden-state predictability, using the PHi layer to quantify information gained during in-context processing.

  • Empirical evidence: Interesting tasks such as in-context learning and modeling code or literature show high PHi loss, whereas memorized or random structure shows low PHi loss.Figure 1 contrasts PHi loss with next-token loss across specialized Transformers and pretrained LLMs.
  • Motivation: Next-token loss can misrepresent computation: random data may be hard to predict yet require no useful processing, while trivial memorized sequences remain easy.The authors therefore seek a metric sensitive to the model’s internal computation rather than token unpredictability alone.
  • Method: PHi loss measures how well a sequence model predicts its own future hidden states instead of future tokens.The PHi layer imposes an information bottleneck and predicts hidden representations autoregressively.
  • Method: The metric quantifies the complexity of the model’s in-context program through hidden-state prediction loss and a learned autoregressive prior over latent codes.The prior predicts each latent code from preceding latent codes, enabling the information gained at each step to be measured.
  • Method: The PHi layer is architecture-agnostic, supports joint training or post-hoc insertion, and can be used with Transformers, RNNs, and pretrained LLMs.Its bottleneck encourages hidden states to retain information relevant for generating future tokens.
  • Empirical evidence: Across formal languages, mathematical reasoning, and generated reasoning chains, PHi loss tracks task complexity or solution correctness beyond next-token loss.The reported evidence covers language complexity, problem hardness, and the correctness of self-generated solutions.

2. Method

The method inserts a PHi layer between bottom and top sequence-model layers to bottleneck hidden states, predict them autoregressively, and quantify novel information in the model’s computation.

  • 2.1. Prediction of Hidden States (PHi): The PHi layer is inserted between bottom and top layers, bottlenecks each hidden state, and predicts future hidden states autoregressively.The bottom layers produce hidden representations, while the top layers use transformed representations for next-token prediction.
  • 2.1. Prediction of Hidden States (PHi): A learned encoder maps each hidden state to a latent posterior, while an autoregressive learned prior predicts the current latent from preceding latent variables.The latent variables are sampled from the posterior, and the decoder reconstructs the transformed hidden states.
  • 2.1. Prediction of Hidden States (PHi): The PHi layer can be trained jointly with the sequence model or fitted post-hoc while keeping a pretrained model fixed.This supports insertion into the autoregressive modeling pipeline without requiring a single training regime.
  • 2.2. Quantifying “Interestingness”: The PHi loss is the KL divergence between posterior and prior, measuring novel information at each timestep that is not predictable from the past.For whole sequences, the method averages this quantity over tokens to obtain an overall interestingness measure.
  • 2.3. Training Objective: Training combines next-token negative log-likelihood with PHi loss, preserving token-prediction capability while discouraging latent information that is unnecessary for future predictions.This creates an information bottleneck: latent states retain information needed for next-token prediction without deviating unnecessarily from the learned prior.
  • 2.4. The Connection Between LNLL and LPHi: The hidden-state description length quantifies information extracted from the input and used to predict the sequence remainder, complementing token-sequence entropy coding.The latent sequence can be encoded using the learned predictive prior, followed by token prediction without explicitly encoding the past tokens.

3. Experiments

Across fully trained sequence models and pretrained Llama, PHi loss distinguishes non-trivial in-context computation more reliably than next-token loss and tracks task complexity. Higher PHi loss also correlates with mathematical difficulty and helps identify correct self-generated reasoning chains.

  • 3.1.1. BORING VS. INTERESTING TASKS: PHi loss is high for in-context language learning but low for memorized, lookup-based, and random tasks, consistently across Transformer and LSTM models.Next-token loss instead ranks memorized tasks lowest, random sequences highest, and in-context learning intermediate, which does not identify the underlying computation.
  • 3.1.2. SIMPLE VS. COMPLEX TASKS: PFA complexity predicts higher PHi loss even after controlling for next-token loss, with sequence-level partial correlation r = 0.37 and 95% confidence interval [0.32, 0.43].The analysis uses 1000 instances of unfamiliar PFAs and groups complexity into 10 levels.
  • 3.2. Experiments with Large Language Models: In pretrained Llama 3B, placing PHi in early layers can cause posterior collapse, whereas later placement minimally affects next-token accuracy and differentiates tasks most strongly around layers 18-24.The later-layer setting is used for the remaining experiments.
  • 3.2.2. SIMPLE VS. COMPLEX TASKS IN LLMS: PHi loss increases with mathematical problem difficulty, with partial correlation r = 0.079 and 95% confidence interval [0.07, 0.09] across tested layers.The relationship is measured on step-by-step solutions from MATH problems labeled from Level 1 to Level 5.
  • 3.2.3. CORRECT VS. ERRONEOUS RATIONALES: Selecting the rationale with higher PHi loss significantly increases the chance of choosing the correct mathematical answer, including when lower next-token loss favors the wrong answer.The comparison separates hidden-state interestingness from rationale coherence, which is associated with lower next-token loss.

4. Related Work

The paper situates PHi among hidden-state prediction, variational autoencoders, information bottlenecks, mechanistic interpretability, in-context language learning, and model-complexity measures.

  • Hidden-State Prediction: Earlier hidden-state prediction work used hierarchical RNNs or predictable representations, but lacked the explicit predictability incentive introduced here.Schmidhuber (1992a) predicted another RNN’s hidden state, while Schmidhuber & Prelinger (1993) encouraged informative and predictable representations.
  • Variational Models: PHi relates to VAEs through latent-variable reparameterization while replacing fixed priors with structured autoregressive prediction.Structured priors have also been proposed in VAE research.
  • Interpretability: Information bottlenecks and autoencoders are established tools in mechanistic interpretability, including token attribution and sparse feature decomposition.The cited work applies these tools to BERT attribution, superposition, and language-model representations.
  • In-Context Language Learning: The experiments adapt in-context language learning with probabilistic finite automata, where models infer automaton structure from transitions presented in context.Aky¨urek et al. (2024) introduced this setup, which the paper adopts for evaluation.
  • Complexity Measures: Unlike prequential coding, which measures joint encoding of a network and its training data, this work quantifies the complexity of an implicit in-context program.The paper connects this two-part description perspective to sophistication while noting the formal connection remains unresolved.

5. Limitations & Future Work

The paper identifies conceptual, scaling, placement, aggregation, and evaluation limitations while proposing exploration and self-supervised applications for PHi loss.

  • Conceptual Scope: PHi loss can increase for repeated random sequences because redundancy must pass through the bottleneck, even though redundancy may not constitute interestingness.Its precise relation to formal concepts such as sophistication remains open.
  • Measurement Scale: The hidden states require many more bits than discrete tokens, complicating direct comparison between PHi loss and next-token loss.The paper attributes this difference most likely to hidden-state dimensionality and suggests quantized bottlenecks as one possible remedy.
  • Model Integration: Pretrained-model applications use heuristic bottleneck placement, and optimal positions plus training large models with the bottleneck remain open questions.The authors also suggest training from scratch or fine-tuning so models can accommodate the information bottleneck.
  • Aggregation: Mean aggregation works well for fixed-length sequences, but variable-length data may require alternatives such as summation, top-k averaging, or thresholding.The paper also calls for further study of correlation between next-token and PHi losses when forming aggregate scores.
  • Future Work: Because interestingness and task complexity are difficult to define, evaluation would benefit from datasets with explicit complexity gradations.The paper proposes PHi loss as a possible intrinsic reward for exploration or self-supervised objective when external feedback is unavailable.

6. Conclusion

The paper introduces PHi loss as an information-theoretic measure of in-context computation, implemented through a self-predictive bottleneck and evaluated across task-complexity settings.

  • Conclusion: PHi loss measures the complexity or “interestingness” of in-context computation by quantifying non-trivial structure encoded in neural sequence models’ hidden states.The metric is implemented by augmenting architectures such as Transformers or RNNs with a self-predictive information bottleneck.
  • Conclusion: The method can be trained from scratch or added post-hoc to pretrained models, including Llama, providing flexibility in deployment.Experiments report alignment with task complexity, including the description length of probabilistic finite automata.
  • Conclusion: The experiments find that PHi loss correlates with meaningful task-complexity notions and intuitive conceptions of complexity across varied tasks.The paper suggests it could serve as an objective for applications lacking direct external feedback.

A. Details on Experiments with Fully Trained Sequence Models

The fully trained sequence-model experiments use Transformer and LSTM architectures, probabilistic finite-automaton tasks, perturbations for robustness, and separate training to avoid memorization bias.

  • Model Configurations: The experiments use a 12-layer, 768-dimensional Transformer and a 2-layer, 768-dimensional LSTM, each with one middle PHi layer.The Transformer uses six attention heads and 2048-dimensional MLPs; the PHi layer is placed after layer 6, versus after layer 1 in the LSTM.
  • PHi Configuration: The PHi components use 768-dimensional latents, linear posterior and activation transforms, and an autoregressive predictor implemented as one model layer.The Transformer setup uses a Llama transformer layer for the predictor.
  • PFA Data: PFA parameters vary over 3–12 states, 1–4 edges per state, vocabulary sizes of 4–18, and example lengths of 10–50.Training sequences contain 10–20 examples, with 20% of tokens randomly perturbed in half the examples; testing uses no perturbation.
  • Task Construction: The four tasks contrast memorized subsequences, memorized programs, newly sampled in-context languages, and random sequences.Sequences are uniformly sampled across tasks, with ten randomly generated examples for memorized sequences and fixed or newly sampled PFAs for program tasks.
  • Training Design: The second experiment trains exclusively on newly sampled in-context-language data to avoid bias from memorized PFAs in the memorized-program task.This isolates the target task from possible training exposure to fixed automata.

B. Details on Experiment with Pre-Trained LLMs

The pre-trained LLM experiments fix the Llama 3B model and train only PHi-layer weights under controlled optimization settings, using mixed language, math, and reasoning data. The section evaluates PHi across several task types, including trivial, license, random, and literature tasks.

  • Experimental setup: The experiments fix the pre-trained instruction-tuned Llama 3B parameters and optimize only the PHi-layer weights.Training uses Adam for 10,000 steps with batch size 2, gradient clipping at 1.0, and learning rate 0.0001 without warm-up or decay.
  • Experimental setup: The PHi layer is evaluated at multiple positions in the Llama model, including after layers 2, 4, 6, 8, and 10.
  • Training data: The training mixture combines SlimPajama natural-language data with the MATH and GSM-8k training sets.
  • Task examples: The task suite includes trivial sequence-generation tasks, memorized license texts, randomly shuffled data, and literature tasks.Examples include multiplication tables, calendar dates, binary conversion, Spanish number names, repeated strings, and software licenses.

C. Additional Results

Additional experiments test whether PHi remains informative when repeated context makes sequences predictable and when next-token losses are matched. Across Transformer and Llama settings, more complex PFAs generally produce higher PHi loss.

  • Copying task: The copying task repeats previously unseen random subsequences so that the second occurrence can be predicted through context retrieval.This tests whether information must cross the PHi bottleneck even when the repeated sequence becomes predictable from prior context.
  • Normalized complexity: More complex PFAs produce higher PHi loss across next-token-loss bins in the fully trained Transformer experiment.The losses are normalized within each next-token-loss bin, isolating the relationship between PFA complexity and PHi loss.
  • Normalized complexity: The same higher-PHi-loss pattern appears across most next-token-loss bins when the PFA layer is inserted after layer 20 in Llama 3B.

C.2. PFA Task Examples and Normalized Complexity

The appendix illustrates four task families and normalizes PHi losses within next-token-loss bins to compare computational complexity independently of token-prediction difficulty. In-context language learning shows especially high PHi values later in sequences.

  • PFA task examples: The examples compare in-context language learning with memorized sequences, memorized programs, and random data.
  • Normalized complexity: PHi losses are normalized across each next-token prediction-loss bin in the alternative views of the complexity figures.This normalization supports comparisons at comparable next-token-loss levels.
  • PFA task examples: In-context language learning produces higher PHi values, especially in later sequence positions, than memorized sequences, memorized programs, and random data.
  • Normalized complexity: The figures plot PHi loss against next-token loss while stratifying results by PFA complexity.

C.3. Correct vs. Erroneous Rationales for the MATH Dataset

For MATH reasoning chains, responses with higher PHi loss are more likely to be correct, including when lower next-token loss favors the wrong answer. This relationship is strongest for difficult counterintuitive questions, while lower next-token loss is only about 52% predictive overall.

  • Correctness prediction: The analysis extends the earlier results on correctness prediction for generated rationales.
  • Correctness prediction: High-PHi-loss responses are significantly more likely to be correct on MATH problems, whereas choosing the lower-next-token-loss response is only about 52% predictive.The comparison uses chain-of-thought prompts without in-context examples.
  • Difficulty analysis: For difficult counterintuitive MATH questions, high PHi loss correlates strongly with correctness, unlike for easier questions.Counterintuitive questions are those where the response with lower next-token loss is incorrect.
  • Difficulty analysis: Figure 16 separates MATH rationales by question difficulty and compares high-PHi selection, counterintuitive-pair selection, low-next-token-loss selection, and partial correlation.

C.4. Histograms for Next Token and PHi Loss

These histograms compare next-token loss with PHi loss across synthetic tasks and Llama 3B data. They highlight cases where low next-token loss coexists with high PHi loss, especially for copying and language-like data.

  • Copied sequences form a Transformer cluster with low next-token loss but high PHi loss, unlike the LSTM result.This cluster appears in the top-right histogram; the passage does not complete the explanation for the LSTM difference.
  • In the four-task comparison, memorized sequences cluster at low values for both losses, while random data clusters at low PHi loss and high next-token loss.
  • Only in-context language learning produces a substantial concentration of tokens with high PHi loss in the four-task histograms.
  • For Llama 3B, Literature and Code show more tokens spreading into the high-PHi-loss region across layers 18, 20, 22, and 24.
Loading 2503.13431v1…