PaperlayerView exact revision on arXiv

Source-linked AI summary

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

William Lugoloobi, Thomas Foster, William Bankes, Chris Russell

arXiv:2602.09924v4cs.CLcs.AIcs.LG

TL;DR

The paper investigates whether LLMs’ pre-generation representations reveal which inputs they will solve, avoiding costly repeated rollouts for success estimation and routing. It trains probes for policy-specific success and difficulty, finding model-specific signals distinct from human difficulty and using them to reduce routing cost. These signals become harder to access with extended reasoning, while the reported routing gains depend on probe reliability.

  • Problem

    Estimating stochastic success rates for routing requires multiple costly rollouts, motivating prediction of model success before generation.

  • Method

    The paper trains linear probes on pre-generation activations to predict human difficulty, model-specific success, and policy-specific outcomes, then uses the estimates for model routing.

  • Results

    70% lower inference cost on MATH lets probe-guided routing match the highest-capability single-model performance, while model difficulty is more predictive than human difficulty.

  • Takeaways & Limitations

    Pre-generation representations can support practical model allocation, but their success signals become harder to access as test-time reasoning increases.

  • Takeaways & Limitations

    The study uses linear probes at one post-instruction position, does not probe during generation or test cross-domain transfer, and uses fixed-k majority voting.

Abstract

from arXiv · show

Running LLMs with extended reasoning on every problem is expensive, but determining which inputs actually require additional compute remains challenging. We investigate whether their own likelihood of success is recoverable from their internal representations before generation, and if this signal can guide more efficient inference. We train linear probes on pre-generation activations to predict policy-specific success on math and coding tasks, substantially outperforming surface features such as question length and TF-IDF. Using E2H-AMC, which provides both human and model performance on identical problems, we show that models encode a model-specific notion of difficulty that is distinct from human difficulty, and that this distinction increases with extended reasoning. Leveraging these probes, we demonstrate that routing queries across a pool of models can exceed the best-performing model whilst reducing inference cost by up to 70\% on MATH, showing that internal representations enable practical efficiency gains even when they diverge from human intuitions about difficulty. Our code is available at: https://github.com/KabakaWilliam/llms_know_difficulty

1 Introduction

The paper asks whether pre-generation activations reveal an LLM’s own likelihood of success and shows that linear probes can extract useful, policy-specific difficulty signals for inference routing.

  • Linear probes extract estimates of an LLM’s own success from pre-generation activations, addressing whether success can be predicted before generation.
  • The study spans mathematics and coding tasks and predicts success under multiple decoding policies.
  • Spearman ρ = 0.83–0.87 for human difficulty and 0.40–0.64 for model difficulty on E2H-AMC, with model difficulty more predictive of actual performance.
  • AUROC > 0.7 for several models under fixed decoding policies, with probe discrimination stable across sampling temperatures and majority-voting thresholds.
  • 70% lower inference cost on MATH lets threshold-based and utility-maximizing routers match the highest-capability single-model performance, with similar gains on AIME and GSM8K.

2 Related Work

Prior work identifies correctness and difficulty signals in model representations, but leaves unresolved whether they reflect human or model-specific difficulty and how to use them for routing.

  • Prior methods estimate correctness through prompted self-assessments, activation directions, or verifiers, but generation overhead and weaker mathematical-reasoning performance limit direct routing use.
  • The paper trains supervised linear classifiers on pre-generation representations instead of extracting unsupervised correctness directions.
  • Existing difficulty studies show linearly decodable signals in pre-generation activations but do not establish whether those signals represent human difficulty, model-specific difficulty, or both.
  • Routing approaches commonly rely on input length, perplexity, heuristic confidence, multiple API calls, or learned latent traits rather than directly predicting model-specific success.

3 Predicting Difficulty

The paper treats difficulty as actor- and policy-dependent, defining human and model-centered targets and probing pre-generation activations to predict them.

  • 3.1 Different Notions of Difficulty: Difficulty depends on the actor facing a task and on the success criterion, motivating separate human- and model-derived definitions.
  • 3.1 Different Notions of Difficulty: IRT difficulty b(q) ∈[0, 1] ranks questions from easiest to hardest using a psychometric scale analogous to Elo.
  • 3.1 Different Notions of Difficulty: Expected success rate measures the probability that a stochastic decoding policy succeeds on question q.
  • 3.1 Different Notions of Difficulty: K Monte Carlo rollouts estimate expected success rate, producing a continuous model-specific difficulty measure under stochastic decoding.
  • 3.1 Different Notions of Difficulty: Binary targets include greedy success, Maj@K majority-vote correctness, and Pass@K success when at least one sampled answer is correct.
  • 3.1 Different Notions of Difficulty: Binary and expected-success metrics are distinct transformations of the answer distribution and therefore provide different prediction targets.
  • 3.2 Activation Probes: Activation probes use residual-stream representations, with MSE for expected success and human IRT difficulty and logistic regression with cross-entropy for binary success.
  • 3.2 Activation Probes: Five-fold cross-validation sweeps layer, regularization weight, and token position among the final non-padding chat-template tokens.

4 Experiments

The experiments evaluate probes of human and model difficulty across mathematical and coding tasks, decoding policies, and reasoning levels. They find that model-specific success signals are useful for prediction but become less recoverable and less aligned with human difficulty as reasoning increases.

  • Experimental Setup: The study evaluates models spanning sizes and reasoning capabilities on mathematics and coding datasets, including E2H-AMC, MATH, GSM8K, AIME, and LiveCodeBench.Multiple answers are generated to construct model-difficulty labels, while E2H-AMC supplies human IRT difficulty for identical questions.
  • Probe Evaluation: Spearman ρ = 0.83–0.87 for human difficulty versus ρ = 0.40–0.64 for model success shows that both signals are linearly represented, but human difficulty is more accessible.Linear probes are evaluated with Spearman correlation, while binary success predictions use AUROC.
  • Reasoning Effects: ρ = 0.58 to ρ = 0.40: GPT-OSS-20B model-success probe performance declines from low to high reasoning despite higher task success.Nonlinear MLP probes do not recover this success-rate loss and degrade more rapidly than linear probes under increased reasoning.
  • Binary Success Prediction: AUROC > 0.7 across most settings, with several results exceeding 0.8, shows that fixed-policy binary success is more discriminable than success-rate prediction.The evaluated policies include Maj@5, greedy, and Pass@5 across math and code tasks.
  • Decoding Policies: 0.84 vs 0.76: greedy decoding generally yields higher probe AUROC than Maj@5 for Qwen2.5-Math-1.5B.For smaller or less capable models, Maj@5 can instead be easier to predict than greedy success.
  • Reasoning Effects: 0.78 to 0.64: increasing GPT-OSS-20B reasoning from low to high decreases Maj@5 AUROC, while MLPs partially recover performance at high reasoning, reaching 0.76 vs 0.64 for linear probes.This pattern indicates reasoning-related representation drift for binary success prediction.
  • Human–Model Difficulty: ∼0.65 to ∼0.45: alignment between probe-predicted model difficulty and human IRT difficulty decreases as GPT-OSS-20B reasoning increases.Probe-based model-difficulty predictions nevertheless outperform human difficulty for predicting Maj@5 failures across reasoning modes.
  • Reasoning Length: Chain-of-thought length increasingly tracks human IRT difficulty while decreasing with empirical and probe-predicted success as reasoning budgets rise.Thus, generation length reflects human-aligned difficulty more than model-relative likelihood of failure in these experiments.

5 Practical Applications: Probe-Guided Routing

Probe-guided routing uses pre-generation success estimates to allocate queries across models with different capabilities and costs. Across benchmarks, it preserves or improves accuracy while adapting model selection and cost to task difficulty, though reliability depends on probe quality and reasoning regime.

  • 5.1 Routing Strategies: Probe-derived success estimates enable performance-cost tradeoffs in utility- and cascade-based routing across heterogeneous model pools.The utility rule combines probe-estimated success probability with normalized expected cost, while separate probes are trained for each model.
  • 5.3 Results: 37% lower cost matches GPT-OSS-20B-high’s 93.3% AIME 2025 accuracy.The router follows the strongest model more closely when benchmark performance varies widely across available models.
  • 5.3 Results: 70% lower cost matches GPT-OSS-20B-high’s 92% MATH accuracy.Experiments use maj@5 accuracy with five generations and vary the cost-weight parameter to trace the frontier.
  • 5.3 Results: On GSM8K, Math-7B is selected as the cost-optimal model at 94.5% accuracy and $0.34, ahead of GPT-OSS-20B-high at 94.4% and $2.4.The router shifts toward efficient models when accuracy saturates across models.
  • 5.4 Discussion: Extended reasoning improves task success but reduces linear accessibility of pre-generation success signals.Longer reasoning increasingly correlates with human difficulty rather than model-relative failure likelihood, weakening probes as routing inputs.
  • 5.4 Discussion: Routing effectiveness is constrained by probe reliability, with larger gaps from oracle utility when discrimination degrades.Even in lower-quality regimes, probes identify cheaper models on saturated tasks and stronger models on harder tasks.

6 Conclusion and Limitations

The paper concludes that pre-generation activations contain extractable, model-specific success signals that can support efficient routing. Its scope is limited by probe degradation under extended reasoning, fixed token-position choices, and the absence of during-generation probing or cross-domain transfer.

  • 6 Conclusion: Pre-generation activations encode model success likelihood, which linear probes extract and use for routing approaching oracle performance.These signals differ from human difficulty, predict model performance more reliably, and generalize across decoding strategies.
  • Limitations: The study focuses on linear probes at one post-instruction position and does not probe during generation or test cross-domain transfer.Probe performance degrades under extended reasoning, and routing uses fixed-k majority voting rather than learned or adaptive k.
  • A Probe and Baseline Formulation: Linear probes use frozen residual-stream activations from selected layers and post-instruction token positions to predict continuous or binary success targets.Ridge regression handles success-rate targets, while regularized logistic regression handles binary correctness.
  • A Probe and Baseline Formulation: Probe configurations are selected on validation data, calibrated with Platt scaling for classification, and evaluated once on held-out test data.The methodology also includes MLP probes to model nonlinear relationships in the representations.
  • A Probe and Baseline Formulation: MLP probes partially recover discrimination lost by linear probes under variable reasoning, reaching AUROC 0.76 versus 0.64 for GPT-OSS-20B-high on Maj@5.For human IRT difficulty and standard binary success prediction, MLPs offer comparable or marginally improved performance.

B Dataset Statistics

This section reports dataset split sizes and probe-selection results, including the layer and token position chosen using validation performance.

  • Train, validation, and test split sizes are reported for each dataset.
  • Figure 5 compares linear and MLP probe validation performance across layers for two QWEN2.5-MATH models on MATH.The linear probe outperforms the MLP probe in both models, while linear performance tends to increase with layer index.
  • The best probe layer and token position are selected on validation, with corresponding validation and test scores reported.Token position is indexed backward from the final token.

D Additional Probe Results

The additional results summarize probe performance across decoding strategies, reasoning modes, model variants, benchmarks, and cross-dataset transfer.

  • AUROC is reported for predicting accuracy under different decoding strategies.
  • Probe performance is compared across GPT-OSS-20B thinking modes using Maj@5 decoding.
  • Benchmark performance is reported across different models.
  • Cross-dataset transfer: Cross-dataset transfer evaluates probes trained on one dataset across all three datasets, with diagonal entries marking in-distribution evaluation.
  • Probe performance is compared across Qwen2.5-7B variants with different degrees of posttraining.
  • All rollouts were performed using VLLM with specified configurations, including Maj@K sampling in the reported hyperparameter tables.

F Additional Routing Results and Data

This section compares routed-model costs and routing strategies on MATH, AIME 25, and GSM8K.

  • Fireworks AI pricing is reported for routed model costs, with GPT-OSS-20B using its own pricing structure.
  • MATH: Table 12 compares routing strategies on MATH.
  • AIME 25: Table 13 compares routing strategies on AIME 25.
  • GSM8K: Table 14 compares routing strategies on GSM8K.

F.1 Full-Cost Routing with an SC-Entropy Baseline

The full-cost analysis accounts for activation-computation costs and compares probe-based routing with an output-based self-consistency entropy baseline.

  • Full-cost probe routing includes input-token costs for obtaining activations from every candidate model, plus the selected model’s generation cost.The routing signal requires one forward pass per model per question.
  • SC-Entropy routing replaces probe-estimated success probability with Shannon entropy over answers from K sampled outputs.
  • The selected model is specified by the routing formulation.
  • SC-Entropy is an expensive output-based reference because its confidence signal becomes available only after generating all K outputs for every candidate model.Its cost includes full generation across the model pool rather than only the selected model’s generation.

F.2 Cascade Routing

Cascade routing uses a probe-estimated success probability to decide whether each query should use a base or stronger model. The strategy can match a stronger model’s accuracy at lower cost and is evaluated across GSM8K, AIME, and MATH routing settings.

  • The cascade routes each input between a base model and a stronger, more expensive model using a threshold-based rule.The base model is Qwen2.5-Math-1.5B and the stronger model is Qwen2.5-Math-7B.
  • The probe estimates the base model’s probability of answering correctly, while τ controls the performance–cost tradeoff.The threshold τ ranges from 0 to 1.
  • The routing results include GSM8K, AIME, and MATH, with model columns reporting the fraction of queries sent to each model.
  • The oracle routes each query to the cheapest model that solves it, escalating only after actual failures and defaulting to the cheapest model if none succeeds.
  • At τ=0.6, the cascade matches GPT-OSS-20B-medium’s 91.2% accuracy while reducing cost by 17%.
Loading 2602.09924v4…