Sparse Attention Decomposition Applied to Circuit Tracing

TL;DR

This paper introduces a method to analyze GPT-2 small's attention mechanisms by identifying sparsely encoded features in singular vectors, revealing communication pathways and circuit redundancies during the IOI task.

Contribution

It presents a novel approach using sparse attention decomposition to isolate and identify features used for communication among attention heads in GPT-2 small.

Findings

Sparse signals enable clear separation of communication features.

Identified redundant paths in GPT-2 attention circuits.

Enhanced understanding of attention head interactions during IOI.

Abstract

Many papers have shown that attention heads work in conjunction with each other to perform complex tasks. It's frequently assumed that communication between attention heads is via the addition of specific features to token residuals. In this work we seek to isolate and identify the features used to effect communication and coordination among attention heads in GPT-2 small. Our key leverage on the problem is to show that these features are very often sparsely coded in the singular vectors of attention head matrices. We characterize the dimensionality and occurrence of these signals across the attention heads in GPT-2 small when used for the Indirect Object Identification (IOI) task. The sparse encoding of signals, as provided by attention head singular vectors, allows for efficient separation of signals from the residual background and straightforward identification of communication paths between attention heads. We explore the effectiveness of this approach by tracing portions of the circuits used in the IOI task. Our traces reveal considerable detail not present in previous studies, shedding light on the nature of redundant paths present in GPT-2. And our traces go beyond previous work by identifying features used to communicate between attention heads when performing IOI.