Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors

TL;DR

The paper tackles the challenge of predicting language model behavior under distribution shifts by leveraging internal causal mechanisms. It introduces a two-stage framework to identify abstract causal models and then predict correctness, supported by two methods: counterfactual simulation and value probing. Across five tasks, causal features prove more robust than non-causal signals, with counterfactual simulation delivering the strongest OOD performance and substantial gains over baselines. The work demonstrates a practical application of causal interpretability for predicting model reliability, highlighting implications for safety and deployability in real-world NLP systems. It also discusses trade-offs, task properties, and the potential for automating high-level causal model discovery to broaden applicability.

Abstract

Interpretability research now offers a variety of techniques for identifying abstract internal mechanisms in neural networks. Can such techniques be used to predict how models will behave on out-of-distribution examples? In this work, we provide a positive answer to this question. Through a diverse set of language modeling tasks--including symbol manipulation, knowledge retrieval, and instruction following--we show that the most robust features for correctness prediction are those that play a distinctive causal role in the model's behavior. Specifically, we propose two methods that leverage causal mechanisms to predict the correctness of model outputs: counterfactual simulation (checking whether key causal variables are realized) and value probing (using the values of those variables to make predictions). Both achieve high AUC-ROC in distribution and outperform methods that rely on causal-agnostic features in out-of-distribution settings, where predicting model behaviors is more crucial. Our work thus highlights a novel and significant application for internal causal analysis of language models.

Figures (6)

• Figure 1: We propose a two-stage framework to predict correctness of model outputs under distribution shifts, illustrated using the counterfactual simulation method on the Indirect Object Identification task wang2023ioi. Stage 1 (Abstract): Identify the abstract causal mechanisms the model uses to solve the task correctly on in-distribution data. Stage 2 (Predict): Predict the model's output correctness on out-of-distribution inputs by checking whether it implements the same mechanisms. Dark blue indicates the key causal variables used for correctness prediction.
• Figure 2: The identified high-level model for the PriceTag task. See high-level models for all tasks in Figure \ref{['fig:highlevel_models']}.
• Figure 3: Robust vs. non-robust features. Many features that are predictive of correctness on in-distribution inputs (purple blocks in panel A), where only very few of them can generalize to out-of-distribution inputs (purple blocks in panel B). These features coincide with features that have causal effects on task predictions, as shown in panel C and D.
• Figure 4: A positive correlation between interchange intervention accuracy and AUC-ROC in in-distribution settings. Interchange intervention accuracy measures the ability to simulate causal effects, and AUC-ROC measures the ability to predict model correctness. Improving the ability to simulate causal effects generally improves the ability to predict the correctness of model outputs. Each point represents one observation. Different shapes represent different dimensions used to train DAS.
• Figure 6: High-level causal models for each task. Non-causal variables are features that are not causally efficacious in the correct high-level causal models, whose values are modified to create OOD settings. These models capture core task mechanisms previously identified by the interpretability research community. Tasks include: IOI, which outputs the non-repeated name wang2023ioi; PriceTag, which determines whether a number falls into a certain interval wu2023interpretabilityatscale; RAVEL, which retrieves country names huang-etal-2024-ravel; MMLU, which selects correct answers hendrycks2021measuringmassivemultitasklanguage; and UnlearnHP, which refuses to answer Harry Potter-related queries arditi2024refusalthaker2024guardrailbaselinesunlearningllms. For MMLU, we use the multiple-choice mechanisms from wiegreffe2025answer, which computes a position variable that is similar to those identified in the IOI and other indexing tasks.