Causal Head Gating: A Framework for Interpreting Roles of Attention Heads in Transformers

TL;DR

This paper introduces causal head gating (CHG), a scalable, causal interpretability method for analyzing attention heads in transformer models, revealing their roles and interactions across diverse tasks.

Contribution

CHG is a novel, dataset-agnostic approach that assigns causal roles to attention heads, enabling causal insights into transformer mechanisms without relying on hypotheses or prompt templates.

Findings

LLMs contain multiple sparse, task-specific sub-circuits.

Head roles depend on interactions, indicating low modularity.

Instruction following and in-context learning rely on separable mechanisms.

Abstract

We present causal head gating (CHG), a scalable method for interpreting the functional roles of attention heads in transformer models. CHG learns soft gates over heads and assigns them a causal taxonomy - facilitating, interfering, or irrelevant - based on their impact on task performance. Unlike prior approaches in mechanistic interpretability, which are hypothesis-driven and require prompt templates or target labels, CHG applies directly to any dataset using standard next-token prediction. We evaluate CHG across multiple large language models (LLMs) in the Llama 3 model family and diverse tasks, including syntax, commonsense, and mathematical reasoning, and show that CHG scores yield causal, not merely correlational, insight validated via ablation and causal mediation analyses. We also introduce contrastive CHG, a variant that isolates sub-circuits for specific task components. Our findings reveal that LLMs contain multiple sparse task-sufficient sub-circuits, that individual head roles depend on interactions with others (low modularity), and that instruction following and in-context learning rely on separable mechanisms.