Sparse Attention Decomposition Applied to Circuit Tracing

TL;DR

This work seeks to isolate and identify the features used to effect communication and coordination among attention heads in GPT-2 small, and explores the effectiveness of this approach by tracing portions of the circuits used for the Indirect Object Identification (IOI) task.

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.

Figures (28)

• Figure 1: Orthogonal slice contributions to attention score (a): $A'_{ij} < 0$; (b) $A'_{ij} > 0$.
• Figure 2: Orthogonal slices used when head is firing: (a) AH (3, 0); (b) AH (4, 11); (c) AH (8, 6); (d) AH (9, 9).
• Figure 3: Number of slices used when firing. (left) IOI dataset (right) non-specific dataset.
• Figure 4: Filtering effect of orthogonal slices. Upstream contributions to (a) (10, 0), 'end' token, all slices of $\Omega$; (b) minimal set of slices of $\Omega$; (c) (9,9), 'end' token, all slices of $\Omega$; (d) minimal set of slices of $\Omega$.
• Figure 5: Traced Network, 256 Prompts. Heads are ovals, tokens are boxes.