Research map / Large language models / Reasoning and chain of thought
Language Models Can Predict Their Own Behavior
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.
Trains linear probes on input-token hidden states of Llama3.1-8B to predict eventual LM behaviors (format failures, jailbreak compliance, confidence, CoT answers) and calibrates them with conformal prediction to enable early exit.
Problem
Detecting LM behaviors such as alignment failures, instruction-format failures, and low confidence currently requires generating the full output; the paper asks whether these behaviors can be predicted from internal representations before any output token is generated.
Approach
Linear probes trained on hidden states of the final input token (no generated tokens) predict eventual output behavior; conformal prediction calibrates a confidence threshold so the probe acts only when its estimate is reliable, otherwise deferring to the LM. Applied as early-exit for CoT classification, this skips generation on confident instances.
Key results
- On WildJailbreak, the conformal probe early-warning system reduces jailbreak success from 30% to 2.7% (91% reduction), versus about 20% with BERT.
- Selective CoT acceleration with probes reduces inference cost by 65% on average across 27 datasets, with worst-case accuracy loss of 1.34% and negative average accuracy loss (-0.46%).
- Probes use about 0.0025% of the parameters of fine-tuned BERT yet outperform it on most format-following settings; conformal probes at alpha=0.9 reach about 89.2% average consistency with 10% coverage.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use conformal probes to override the LM with abstention on predicted safety failures | Llama3.1-8B · WildJailbreak | jailbreak success rate: 30% to 2.7% | Helped |
| Use conformal probes instead of BERT for early warning | Safety failure detection · WildJailbreak | consistency: 83.4 -> 94.0 | Helped |
| Exit CoT generation when a conformal probe is confident | Llama3.1-8B CoT text classification · 27 text classification datasets | inference cost reduction: 65% | Helped |
| Exit CoT generation when a conformal probe is confident | Llama3.1-8B CoT text classification · 27 text classification datasets | accuracy loss: -0.46% | Helped |
| Train and calibrate probes on other datasets rather than the test dataset | Llama3.1-8B CoT early exit · OOD MCQ and Sentiment datasets | inference cost reduction: 81% | Helped |
| Increase the number of probe training datapoints up to 500 (ablation) | Conformal probes · MMLU | consistency and coverage: 500 | Helped |
| Use later rather than early LM layers for probing (ablation) | Conformal probes · MMLU and other datasets | estimation consistency | Helped |
| Vary the probing layer between middle and late layers (ablation) | Conformal probes · MMLU and other datasets | estimation consistency | Mixed |
| Increase the conformal confidence threshold α (ablation) | Conformal probes · MMLU and other datasets | coverage | Hurt |
| Increase the conformal confidence threshold α (ablation) | Conformal probes · MMLU and other datasets | estimation consistency | Helped |
| Increase the LM size (ablation) | Llama3 probe system · Multiple tasks | probe consistency | Helped |
| Probe later rather than initial input tokens (ablation) | CoT prediction probes · CSQA | estimation consistency | Helped |
Surprising findings
- Accuracy of the CoT classifier often increases when the probe interrupts generation early (negative average accuracy loss).
- Probes built from early token positions in the input perform near random; only later tokens carry accessible signal.
Limitations
- Probes perform near random when predicting the LM's incorrect answers, since they cannot access knowledge the model lacks.
- Consistency falls as expected output length grows (negative correlation of -6.1%).
- Layer choice is task-specific with no consistent best layer; early layers perform poorly.
Experimental setup
Datasets: NaturalQA MSMarco TriviaQA SelfAware KnownUnknown WildJailbreak MMLU CosmoQA PiQA ARC MedMCQA CommonsenseQA
Benchmarks: SelfAware KnownUnknown WildJailbreak MMLU OpenBookQA 27 text classification datasets
Metrics: estimation consistency coverage precision recall accuracy loss inference cost reduction
Baselines: BERT-large fine-tuned on prompt text random baseline vanilla CoT
Builds on: conformal prediction linear probing BERT fine-tuning
Closest papers by meaning
- Test-Time Alignment for Large Language Models via Textual Model Predictive Control (ICLR 2026 · similarity 0.89)
- Language Model Behavioral Phases are Consistent Across Architecture, Training Data, and Scale (NeurIPS 2025 · similarity 0.89)
- Observational Scaling Laws and the Predictability of Langauge Model Performance (NeurIPS 2024 · similarity 0.88)
- A Probabilistic Perspective on Unlearning and Alignment for Large Language Models (ICLR 2025 · similarity 0.87)
- Detecting High-Stakes Interactions with Activation Probes (NeurIPS 2025 · similarity 0.87)
- Statistical Hypothesis Testing for Auditing Robustness in Language Models (ICML 2025 · similarity 0.87)
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.