Extracting Rule-based Descriptions of Attention Features in Transformers

TL;DR

This work targets mechanistic interpretability of transformer attention by shifting from exemplar-based feature descriptions to symbolic rule-based descriptions. It defines three rule types—skip-gram, absence, and counting—and develops a pipeline to automatically extract these rules from sparse autoencoder representations of attention heads, demonstrated on GPT-2 small. Empirically, many attention features can be well approximated by around 100 skip-gram rules, with absence rules appearing early and counting rules emerging through feature competition, including distractor suppression scenarios. The results lay groundwork for a scalable, interpretable framework to describe how input patterns influence transformer outputs, offering a concrete direction for future work to extend rule types, enhance feature decompositions, and integrate rules across model components to understand full model behavior.

Abstract

Mechanistic interpretability strives to explain model behavior in terms of bottom-up primitives. The leading paradigm is to express hidden states as a sparse linear combination of basis vectors, called features. However, this only identifies which text sequences (exemplars) activate which features; the actual interpretation of features requires subjective inspection of these exemplars. This paper advocates for a different solution: rule-based descriptions that match token patterns in the input and correspondingly increase or decrease the likelihood of specific output tokens. Specifically, we extract rule-based descriptions of SAE features trained on the outputs of attention layers. While prior work treats the attention layers as an opaque box, we describe how it may naturally be expressed in terms of interactions between input and output features, of which we study three types: (1) skip-gram rules of the form "[Canadian city]... speaks --> English", (2) absence rules of the form "[Montreal]... speaks -/-> English," and (3) counting rules that toggle only when the count of a word exceeds a certain value or the count of another word. Absence and counting rules are not readily discovered by inspection of exemplars, where manual and automatic descriptions often identify misleading or incomplete explanations. We then describe a simple approach to extract these types of rules automatically from a transformer, and apply it to GPT-2 small. We find that a majority of features may be described well with around 100 skip-gram rules, though absence rules are abundant even as early as the first layer (in over a fourth of features). We also isolate a few examples of counting rules. This paper lays the groundwork for future research into rule-based descriptions of features by defining them, showing how they may be extracted, and providing a preliminary taxonomy of some of the behaviors they represent.

Figures (10)

• Figure 1: SAE features are usually explained via examplars, which are subjective and hard to interpret. We approximate them with rules that promote or suppress the feature.
• Figure 2: Given an attention layer in a transformer language model (left), our goal is to explain each output feature as an explicit function of input features (center), and provide a global description of these functions in terms of formal rules (right).
• Figure 3: Attention rules can take different forms, depending on how the input features interact. Query-key interaction may promote the generation of specific outputs, leading to skip-gram rules (left). Analogously, suppression leads to absence rules (center). Competition between the two types may lead to counting rules (right).
• Figure 4: Precision and recall for predicting binarized activation values using skip-gram rules, as a function of the maximum number of terms (meaning pairs of key and query features) per output feature. In Fig. \ref{['fig:skipgram_precision_recall']}, terms are selected either according to the magnitude of their weights (Weight-based) or by calculating a gradient-based importance score using a small training set (Importance-based; see Sec. \ref{['sec:experiments']} for more details). Results are averaged over 100 features for each head in GPT-2 small. Importance-based scoring achieves higher recall with fewer terms, but with lower precision. Fig. \ref{['fig:skipgram_f1_by_layer']} shows the F1 scores by layer, selecting terms according to importance. Higher layers generally have worse approximation scores and require more terms.
• Figure 5: Sequences that activate a layer-0 attention feature in GPT-2 small (\ref{['fig:skipgram_examples_dfa']}) and the highest-scoring pairs of key and query input features associated with this feature (\ref{['fig:skipgram_examples_rules']}). The three examples can be explained by the presence of a single skip-gram pattern, ...[If]...[then]; see Section \ref{['sec:skipgram_results']}.