LLMs Encode Their Failures: Predicting Success from Pre-Generation Activations

1. Problem Statement

Large Language Models (LLMs) achieve high performance on complex reasoning tasks (math, coding) but at a significant computational cost, especially when using extended reasoning (e.g., Chain-of-Thought) or sampling-based methods (e.g., Majority Voting).

• The Challenge: Determining a priori which inputs require expensive inference (high compute) and which can be solved cheaply is difficult. Current routing systems rely on indirect proxies (input length, perplexity) or require costly multiple rollouts to estimate success probability.
• The Core Question: Can an LLM's likelihood of success be predicted from its internal state before any output is generated, and can this signal be used to route queries efficiently?

2. Methodology

The authors propose a framework using linear probes trained on pre-generation activations to predict model success.

A. Data and Benchmarks

• Datasets: MATH, GSM8K, AIME, LiveCodeBench, and E2H-AMC (Easy2HardBench AMC subset).
• E2H-AMC Significance: This dataset is crucial as it provides both human IRT (Item Response Theory) difficulty labels and allows for the estimation of model-specific difficulty via Monte Carlo rollouts on the same questions.
• Models: Tested across various architectures including Qwen2.5-Math (1.5B, 7B), GPT-OSS-20B (with varying reasoning budgets), and DeepSeek-R1 variants.

B. Prediction Targets

The study distinguishes between two types of difficulty:

1. Human Difficulty: Based on psychometric IRT scores ($b(q)$).
2. Model Difficulty:
   • Expected Success Rate: The probability of success estimated via $K$ stochastic rollouts ($\hat{s}_{MC}$).
   • Binary Success: Whether a specific deterministic policy (e.g., Greedy or Majority Voting @ $K$) succeeds.

C. The Probe Mechanism

• Input: Residual stream activations extracted at the post-instruction position (the final tokens before generation begins).
• Model: Simple linear classifiers (logistic regression for binary, linear regression for continuous) trained to predict the target (success/failure or difficulty score).
• Baselines: Compared against surface features like question length and TF-IDF vectors.

D. Routing Strategies

Two probe-guided routing policies were evaluated:

1. Cascade Routing: A threshold-based rule where queries with low predicted success probability ($\hat{p} < \tau$) are escalated from a smaller/cheaper model to a larger/more expensive model.
2. Utility-Based Routing: Selects the optimal model from a heterogeneous pool by maximizing a utility function: $\hat{p}_i(x) - \lambda \hat{c}_i$, balancing predicted success against inference cost.

3. Key Contributions

1. Decoupling Human vs. Model Difficulty: The paper demonstrates that LLMs encode two distinct signals in pre-generation activations:

   • Human Difficulty: Highly linearly decodable ($\rho \approx 0.83-0.87$).
   • Model Difficulty: Less linearly accessible ($\rho \approx 0.40-0.64$) but more predictive of actual model performance.
   • Crucial Finding: As models employ extended reasoning, their internal notion of difficulty diverges from human difficulty. Models solve "human-hard" problems efficiently, yet their internal representations still track human difficulty via longer generation lengths, obscuring their own success signals.

2. Predictive Power of Pre-Generation Activations:

   • Linear probes significantly outperform surface features (TF-IDF, length) in predicting success.
   • Probes achieve strong discrimination (AUROC > 0.7, often > 0.8) for binary success under fixed decoding policies (Greedy, Maj@K).
   • Degradation with Reasoning: Probe reliability decreases as reasoning budgets increase (e.g., AUROC drops from 0.78 to 0.64 for GPT-OSS-20B as reasoning goes from Low to High), suggesting extended Chain-of-Thought makes difficulty signals less linearly separable pre-generation.

3. Efficient Model Routing:

   • Probe-guided routing achieves Pareto improvements, matching the accuracy of the strongest single model while reducing inference costs by 17% to 70%.
   • On the MATH benchmark, the router matched the best model's performance at 70% lower cost.
   • The system adapts to dataset difficulty: it routes to expensive models on hard tasks (AIME) but selects efficient models when performance plateaus (GSM8K).

4. Key Results

• Difficulty Divergence: In GPT-OSS-20B, as reasoning capability increases, the correlation between probe-predicted model difficulty and human IRT difficulty drops (from $\sim0.65$ to $\sim0.45$).
• Chain-of-Thought Length: Extended reasoning causes models to allocate more tokens to problems humans find difficult, even when the model is likely to succeed. This creates a "human-aligned" length signal that decouples from the model's actual uncertainty.
• Cost Savings:
  • MATH: 70% cost reduction while matching top-tier accuracy.
  • AIME 2025: 37% cost reduction matching the strongest model.
  • GSM8K: The router correctly identifies that expensive high-reasoning models are unnecessary, selecting cheaper models with comparable accuracy.
• Bottleneck: The primary limitation of routing is not the routing algorithm itself, but the reliability of the probe. When probes degrade (due to extended reasoning), routing performance gaps widen relative to an "oracle" (perfect knowledge) policy.

5. Significance and Implications

• Practical Efficiency: This work provides a method to reduce the cost of LLM inference by up to 70% without sacrificing accuracy, making high-reasoning models more economically viable.
• Theoretical Insight: It reveals a fundamental shift in how LLMs represent difficulty. As models become more capable, their internal "difficulty" metric becomes less aligned with human intuition and harder to extract linearly before generation.
• Limitations & Future Work:
  • Probes degrade with extended reasoning, suggesting difficulty signals may become non-linear or shift to later generation stages.
  • Current probes are position-specific (post-instruction); future work should explore intermediate positions or non-linear probes.
  • Cross-domain transfer (e.g., Math to Code) remains unexplored.

Conclusion: The paper establishes that LLMs "know" their own failures before generating text. By leveraging these internal signals via linear probes, we can build adaptive inference systems that dynamically allocate compute, achieving near-oracle efficiency even when human and model notions of difficulty diverge.