Towards Understanding the Nature of Attention with Low-Rank Sparse Decomposition

TL;DR

This work tackles the challenge of interpreting Transformer attention by addressing attention superposition, where atomic attention units are entangled across many MHSA heads. It introduces Low-Rank Sparse Attention (Lorsa), an overcomplete, sparsity-constrained replacement for MHSA that yields thousands of 1D read/write units (OV circuits) and selects Top-K active heads per token to reconstruct MHSA outputs. Lorsa recovers known attention mechanisms (e.g., induction heads, name movers) and reveals new interpretable behaviors, including arithmetic-specific heads and thematic anchors, while achieving interpretability comparable to Sparse Autoencoders (SAEs) and enabling cross-head attribution of Q/K/V activity. The approach enables finer-grained circuit discovery and provides tools for automated interpretability assessment, though challenges remain in achieving fully independent QK circuits and in understanding cross-layer interactions, with potential implications for in-context learning analysis and model biology.

Abstract

We propose Low-Rank Sparse Attention (Lorsa), a sparse replacement model of Transformer attention layers to disentangle original Multi Head Self Attention (MHSA) into individually comprehensible components. Lorsa is designed to address the challenge of attention superposition to understand attention-mediated interaction between features in different token positions. We show that Lorsa heads find cleaner and finer-grained versions of previously discovered MHSA behaviors like induction heads, successor heads and attention sink behavior (i.e., heavily attending to the first token). Lorsa and Sparse Autoencoder (SAE) are both sparse dictionary learning methods applied to different Transformer components, and lead to consistent findings in many ways. For instance, we discover a comprehensive family of arithmetic-specific Lorsa heads, each corresponding to an atomic operation in Llama-3.1-8B. Automated interpretability analysis indicates that Lorsa achieves parity with SAE in interpretability while Lorsa exhibits superior circuit discovery properties, especially for features computed collectively by multiple MHSA heads. We also conduct extensive experiments on architectural design ablation, Lorsa scaling law and error analysis.

Figures (15)

• Figure 1: (A) Low-Rank Sparse Attention (Lorsa) comprises thousands of sparsely activated attention heads with 1D outputs, designed to extract interpretable attention units from the original Multi Head Self Attention (MHSA). (B) Lorsa serves as a replacement model for Transformer attention, substituting sparse interpretable components for attention modules. (C) Each Lorsa head explains an atomic feature-feature interaction across token positions, which was originally a part of an MHSA head or spread across multiple heads, i.e. put in attention superposition.
• Figure 2: Visualization dashboard for a "you"-specific induction Lorsa head. We provide an example interpretation of each item below.
• Figure 3: Automated interpretability scores of Lorsa heads and SAE features. Each distribution is estimated with 100 heads / features. The average score of each group is represented by a horizontal dash line. We highlight distributions with larger mean value suggested by t-tests with $\alpha=0.05$.
• Figure 4: Examples of Lorsa heads re-discovering previously reported heads. Lorsa.5.3378: The Acronym Head attends to the parentheses and preceding text to predict the abbreviation. Lorsa.6.2814: Successor Head attends to the previous number token and predicts the next number. Lorsa.8.5963: Copy Suppression Head attends to the previous subject and suppresses its copy. Lorsa.10.4066: Attention Sink Head simply attends to the '<|beginoftext|>' token.
• Figure 5: For the prompt "$36+62=$", Lorsa moves two operands to the last position with 3 heads each. The first operand (36) is attended in terms of $z$ pattern by an "op1$\in \numrange{27}{43}$", an "op1$\%\ 10 \in [4, 5, 6]$" and an "op1$\%\ 10 \in [6, 7, 8]$" head, which uniquely determines "op1 = 36". The same applies to op2.