Isolating Attention Heads with Path Patching

This article provides a direct examination of how path patching is used in artificial intelligence interpretability to locate the specific attention heads responsible for Indirect Object Identification (IOI) in Transformer models. By systematically replacing internal activations between clean and corrupted text prompts, researchers can map the exact neural circuits that allow a language model to deduce the correct recipient of an action. What follows is a straightforward breakdown of the IOI task, the mechanics of path patching, and the distinct groups of attention heads uncovered by this technique.

The Indirect Object Identification (IOI) Task

To understand how a Transformer processes language, researchers use the Indirect Object Identification (IOI) task. An IOI sentence involves two subjects and an action, such as: "When Mary and John went to the store, John gave a drink to..."

In this scenario, "John" is the repeated subject, and "Mary" is the indirect object. A well-trained Transformer will accurately predict "Mary" as the next word. Interpretability researchers aim to find the exact internal components—specifically the attention heads—that calculate this answer rather than simply guessing.

How Path Patching Works

Standard activation patching tests the importance of a single attention head by overwriting its output from a "clean" prompt with the output from a "corrupted" prompt. However, this affects every downstream layer in the network, making it difficult to tell if the head influences the final output directly or merely triggers another head later on. Path patching solves this by isolating specific connections (edges) rather than just individual components (nodes).

Path patching operates through the following steps:

  1. The Clean Run: The model processes the standard IOI sentence ("When Mary and John... John gave a drink to"), and the internal activations of all attention heads are recorded.
  2. The Corrupted Run: The model processes a modified sentence where the clues are removed, using three distinct names ("When Alice and Bob... Charlie gave a drink to"). The model can no longer reliably predict "Mary" or "John". These activations are also recorded.
  3. Targeted Intervention: Researchers select a "sender" head and a "receiver" head (or the final output layer). They run the clean prompt again, but they replace the specific communication pathway between the sender and the receiver with the activation from the corrupted run. All other pathways in the network remain frozen to their clean states.
  4. Measuring the Logit Difference: Researchers measure how much this isolated change harms the model's ability to predict "Mary" over "John". If the model's accuracy drops significantly, that specific pathway between the two heads is proven to be vital for the IOI task.

Isolating the IOI Circuit

By applying path patching backward from the final output of the model to the earliest layers, researchers can trace the flow of information step-by-step. This process filters out hundreds of irrelevant attention heads and isolates the exact circuit responsible for IOI.

Path patching reveals three primary classes of attention heads working in sequence:

  • Duplicate Token Heads: Located in the early layers, these heads identify that the name "John" has appeared twice in the sentence.
  • S-Inhibition Heads: Located in the middle layers, these heads receive information from the Duplicate Token Heads. Path patching shows that their specific role is to actively suppress the repeated name ("John") from being attended to by later layers.
  • Name Mover Heads: Located in the late layers, these heads attend to names in the sentence and copy them directly to the final output. Because the S-Inhibition Heads have blocked the pathway for "John", the Name Mover Heads only copy the remaining valid name ("Mary") to the output, successfully completing the sentence.

Through path patching, researchers move beyond knowing where information is stored in a Transformer to knowing how it moves. By freezing the surrounding network and corrupting only single pathways at a time, path patching proves how Duplicate Token Heads, S-Inhibition Heads, and Name Mover Heads communicate to solve complex grammatical problems.