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

Zhengfu He, Junxuan Wang, Rui Lin, Xuyang Ge, Wentao Shu, Qiong Tang, Junping Zhang, Xipeng Qiu

Keywords: Sparse Autoencoders, Interpretability tooling and software, Circuit analysis

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

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. 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 *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/>.

Submission Number: 263