Internal Causal Mechanisms Robustly Predict Language Model Out-of-Distribution Behaviors
This addresses the challenge of improving model reliability and interpretability for AI researchers and practitioners, though it is incremental in applying existing causal analysis techniques to new prediction tasks.
The paper tackled the problem of predicting language model behaviors on out-of-distribution examples by leveraging internal causal mechanisms, achieving high AUC-ROC scores that outperform causal-agnostic methods.
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.