Research map / Probabilistic and scientific modeling / Causal discovery
Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
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.
Shows that internal causal variables localized via interchange interventions (DAS) predict LLM correctness more robustly under distribution shift than confidence scores or causal-agnostic probes, via counterfactual simulation and value probing.
Problem
Predicting whether a language model's outputs are correct under distribution shifts, where confidence scores and probes on non-causal internal features are unreliable.
Approach
Use high-level causal mechanisms localized in the residual stream (via DAS) to predict correctness: counterfactual simulation checks whether the model still computes key causal variables under background-variable perturbations, and value probing trains a classifier on causal variable values and uses its confidence.
Key results
- Counterfactual simulation (output mapping) reaches AUC-ROC 0.856 on PriceTag OOD (change currency format) versus 0.544 for the best confidence-score variant on the same shift.
- Counterfactual simulation performs best on 8 of 10 OOD settings, improving average AUC-ROC by 13.84% over prior baselines per the authors.
- In-distribution, causal-feature methods are comparable to or slightly better than confidence scores (e.g., IOI AUC-ROC 0.999 for counterfactual simulation vs 0.935 for confidence).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use counterfactual simulation with output mapping instead of the strongest non-counterfactual baseline | Correctness prediction with Llama-3-8B-Instruct · OOD tasks (8 of 10) | AUC-ROC | Helped |
| Use causal internal features rather than non-causal features | Correctness prediction · OOD settings across five tasks | AUC-ROC | Helped |
| Use confidence scores under distribution shift | Correctness prediction · OOD settings | AUC-ROC | Hurt |
| Use value probing instead of other causal methods | Correctness prediction · OOD settings | AUC-ROC | Hurt |
| Scale model size for confidence-score prediction | RAVEL · In-distribution | AUC-ROC | No clear effect |
| Use confidence scores instead of counterfactual simulation as task accuracy declines | Correctness prediction · OOD settings | AUC-ROC | Hurt |
Surprising findings
- Confidence scores become notably less predictive as task accuracy declines under distribution shift, and the optimal temperature does not transfer across ID and OOD settings.
- Confidence scores on RAVEL do not improve much when scaling from 8B to 70B models.
Limitations
- Relies on linear localization (DAS) and assumes linear representations of causal variables.
- Requires prior knowledge to propose high-level causal models; MMLU causal models are very incomplete and causal methods do not beat confidence on MMLU.
- Evaluated mainly on Llama-3-8B-Instruct and five tasks, with OOD settings constructed by the authors.
Experimental setup
Datasets: IOI PriceTag RAVEL MMLU UnlearnHP
Benchmarks: IOI PriceTag RAVEL MMLU UnlearnHP
Metrics: AUC-ROC
Baselines: Confidence Score (first N tokens) Confidence Score (answer tokens) Correctness Probing (prompt last token) Correctness Probing (background variables)
Builds on: Distributed Alignment Search causal abstraction temperature scaling
Closest papers by meaning
- Can Large Language Models Infer Causation from Correlation? (ICLR 2024 · similarity 0.88)
- Rethinking the Development of Large Language Models from the Causal Perspective: A Legal Text Prediction Case Study (AAAI 2024 · similarity 0.87)
- Causality can systematically address the monsters under the bench(marks) (NeurIPS 2026 · similarity 0.87)
- Causal Concept Graph Models: Beyond Causal Opacity in Deep Learning (ICLR 2025 · similarity 0.86)
- Training Language Models to Explain Their Own Computations (NeurIPS 2026 · similarity 0.86)
- Towards Best Practices of Activation Patching in Language Models: Metrics and Methods (ICLR 2024 · 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.