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

TL;DR

This paper demonstrates that internal causal mechanisms in language models can be used to predict their out-of-distribution behaviors accurately, leveraging causal features for improved interpretability and robustness.

Contribution

It introduces two causal-based methods, counterfactual simulation and value probing, to predict model correctness on out-of-distribution data, advancing interpretability techniques.

Findings

Causal features are highly predictive of model correctness.

Proposed methods outperform causal-agnostic baselines in OOD settings.

Internal causal mechanisms can reliably forecast model behavior beyond training distribution.

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.