Title: Towards Understanding the Nature of Attention with Low-Rank Sparse Decomposition
Primary area: interpretability and explainable AI
Keywords: Mechanistic Interpretability, Attention Superposition, Sparse Dictionary Learning, Circuit Analysis

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 \textit{attention superposition} to understand attention-mediated interaction between
features in different token positions. Lorsa helps find cleaner and finer-grained versions of previously discovered MHSA behaviors like induction
heads, successor heads, attention sink, and a comprehensive family of arithmetic-specific Lorsa heads. Interestingly, we identify a novel
head type called \emph{subtoken induction heads} that function at character level rather than token level.
Automated interpretability analysis indicates
that Lorsa achieves parity with SAE in interpretability while Lorsa exhibits superior circuit discovery properties.
We also conduct extensive experiments on architectural design ablation, correlation to original MHSA heads and error analysis.
Our early attempt to fully sparsify a toy Transformer succeeds to reveal clean global circuits. Eventually, we hope Lorsa would help
us greatly understand attention computation and enable full sparsification of model computation along with its MLP counterparts.
Lorsa is open-sourced at https://anonymous.4open.science/r/Lorsa-5686/.

TL;DR: We propose a method to interpret transformer attention blocks by decomposing them into interpretable units.