Quantifying LLM Attention-Head Stability: Implications for Circuit Universality

TL;DR

Transformers can exhibit similar final performance across seeds yet different internal computational schemes. The authors quantify attention-head stability across 26 decoder-only architectures by comparing attention score matrices for independently trained refits, revealing a middle-layer stability dip and depth-related divergence. They further find that decoupled weight decay via AdamW improves seed stability in deeper models without sacrificing perplexity, and that the residual stream is notably more stable than individual attention heads, suggesting a robust anchor for universal interpretability. Collectively, the work argues for stability-aware evaluation of LLM circuits to support scalable, auditable safety practices and demonstrates that universal circuits are a nuanced, depth-dependent property influenced by training choices.

Abstract

In mechanistic interpretability, recent work scrutinizes transformer "circuits" - sparse, mono or multi layer sub computations, that may reflect human understandable functions. Yet, these network circuits are rarely acid-tested for their stability across different instances of the same deep learning architecture. Without this, it remains unclear whether reported circuits emerge universally across labs or turn out to be idiosyncratic to a particular estimation instance, potentially limiting confidence in safety-critical settings. Here, we systematically study stability across-refits in increasingly complex transformer language models of various sizes. We quantify, layer by layer, how similarly attention heads learn representations across independently initialized training runs. Our rigorous experiments show that (1) middle-layer heads are the least stable yet the most representationally distinct; (2) deeper models exhibit stronger mid-depth divergence; (3) unstable heads in deeper layers become more functionally important than their peers from the same layer; (4) applying weight decay optimization substantially improves attention-head stability across random model initializations; and (5) the residual stream is comparatively stable. Our findings establish the cross-instance robustness of circuits as an essential yet underappreciated prerequisite for scalable oversight, drawing contours around possible white-box monitorability of AI systems.

Figures (51)

• Figure 1: Experimental workflow for assessing attention head stability. For each of our 26 architectural configurations, we instantiate $N$ GPT-2-small LLMs, each with a different initial random seed. Each LLM is pre-trained from scratch with the OpenWebText or C4 datasets, depending on the architectural configuration. After pre-training, we take $\mathcal{P}$ evaluation prompts and compute attention score matrices for each head in each layer. An anchor model $m$ is selected and compared to a pair model $n$, recursively for each randomly initialized and pre-trained LLM in that architectural subset. Cosine similarity is used to assess the similarity of each attention score matrix from each head, and a "Stability" score is computed from the best matching heads across refits of the architectural subset. See Methods \ref{['methods:attn_head_stab']} for definitions of these terms.
• Figure 2: Only some mid-layer attention heads persist across refits. For each head $h_i$ from layer $l$ in an anchor model, We find its best-matching counterparts in the same layer $l$ across other refits and use $S^{(m)}_{h_{i}}$ from Eq. \ref{['eq:eq4']} as the head's stability (top). This plot thus summarizes inter-seed stability for all the heads of the anchor. Consistently, initial and final layers are more stable than middle layers (bottom).
• Figure 3: Middle layer attention heads exhibit off-layer matching across re-fits. Cross-layer best match alignment heatmap shows strong diagonal dominance with heads typically matching to same-depth layers but broader off-diagonal spread in mid-layers.
• Figure 4: Middle layer attention heads are most unique. Attention heads are most redundant in early and late layers of GPT-like models, while being quite distinct in the intermediate layers.
• Figure 5: Instability of attention heads from middle transformer layers increases with increasing prompt length.