Precomputed figures and the code used to generate them are in this GitHub repository. To run the analysis on other models, edit the models.py file.

TL;DR

Empirical Spectral Distribution (ESD) of the correlation matrix of $W_{QK}$ in dense (non-MoE) text-to-text models exhibit certain trends. These trends can be used as a criterion to create a taxonomy of attention layers, with potential to speed up automated circuit discovery and model diffing. Here I discuss 2 approaches, one is model-agnostic the other is model-specific, neither approaches require forward passes and are task agnostic.

Introduction

Most circuit discovery techniques treat model weights as a passive backdrop for activation-level experiments. A key property of all trained neural networks is that the weight matrices are highly structured. Identifying the relationships between weight matrices is the key to prune the search space and optimizing the search algorithm.

For example, in the Automated Circuit Discovery paper^[1] (Conmy et al., 2023), the connections between the nodes of the computational graph are patched sequentially which makes it difficult to scale. Since only a tiny fraction of components contribute to any specific task^[2][3] most patching steps are unnecessary for identifying the relevant circuit. What if once we identified that a head $h$ of layer $L$ was not important, we could prioritize/deprioritize other heads in the same taxonomy?

To this end I started looking for ways in which I could classify the components of a transformer network. Martin et al. (2021) showed that Heavy Tailed-Random Matrix Theory could be used to identify distinct 5+1 training phases that a weight matrix of an MLP goes through. Inspired by this, I propose 2 classification methods for attention heads based on the ESD of QK circuits ($W_{QK}=W_Q{W}_K^T$) and demonstrate the efficacy of these methods across model sizes and model families^[4].

Why QK?

An attention head factors into:

1. QK circuit (decision making): Selects where to read/write information. Can suppress writing altogether by dumping attention to <bos> or punctuations.
2. OV circuit (transport/translator): Applies a linear map to the the residual stream, largely "dumb" without QK weighting source positions.

To create a taxonomy of attention heads QK was the natural choice.

Setup and Notation

$W_Q$ and $W_K$ have the shape $(d_{model},d_{head})$, therefore shape of $W_Q{W}_K^T$ is $(d_{model},d_{model})$ which is a large matrix e.g. Llama 3.1 405B has $d_{model}=16384$. Computing the SVD values (${\sigma }_i$) of such large matrices is computationally expensive $O\left(n^3\right)$ and sensitive to finite precision errors. 

Instead we compute a smaller matrix $M$ as follows:

1. Compute the SVD of  $W_Q$ and $W_K$.
   $W_Q=U_Q{\Sigma }_Q{V}_Q^T$ and $W_K=U_K{\Sigma }_K{V}_K^T$
    $U$,$V$ are orthonormal matrices and $\Sigma =diag\left({\sigma }_i\right)$
2. Therefore 
   $W_Q{W}_K^T=\left(U_Q{\Sigma }_Q{V}_Q^T\right)\left(V_K{\Sigma }_K{U}_K^T\right)=U_Q\left({\Sigma }_Q{V}_Q^T{V}_K{\Sigma }_K\right)U_K^T=U_{Q}M{U}_K^T$ 
   where M =${\Sigma }_Q{V}_Q^T{V}_K{\Sigma }_K$
3. Since $U_Q$ and $U_K$ are orthonormal matrices, they do not impact the singular values.
   Therefore ${\sigma }_i\left(W_{QK}\right)={\sigma }_i\left(M\right)$^[5]

$M$ is a much smaller $(d_{head},d_{head})$ matrix which is faster to compute and more stable.

The singular values of $W_{QK}$ are related to the eigenvalues of its correlation matrix as

 ${\sigma }_i(W_{QK}{)}^2={\lambda }_i(W_{QK}{W}_{QK}^T)$

The Empirical Spectral Distribution (ESD) is the distribution of the ${\lambda }_i$ denoted by $\rho \left(\lambda \right)$

Classification A: $(\alpha ,\theta )$ of ESD

Classification scheme A aims to create an attention head taxonomy based on the shape and scale of their ESD, this approach is model agnostic. To compare ESD across attention heads, the ESD is normalized by dividing by $\frac{1}{n}\sum {\lambda }_i$ and the counts are normalized so that AUC=1. Using the normalized counts as features, the attention heads classified into clusters with k-means clustering^[6]. 

The average distribution within each class follows a gamma distribution (see section below) so each class can be identified by the associated shape parameter $\alpha$ and scale parameter $\theta$. The goodness-of-fit between the  calculated average ESD and the fitted Gamma distribution is measured using  Jensen-Shannon Divergence (JSD). Across all models studied, the JSD typically lies in the range (0.0-0.1) indicating a very good fit.

Below are the classified attention heads of Qwen3 family of models (0.6B - 32B). 

• Blue bars are the ESD
• Dotted red line is the fitted gamma distribution
• The calculated JSD is in the subplot title
• The error bars show the standard deviation
• "Class_#" are dummy names

Why Gamma distribution?

If ${\sigma }_i$ were normally distributed we would expect to see a ${\chi }^2$ distribution which is a special case of the gamma distribution. The fact that the mean ESD follow the gamma distribution closely (despite ${\sigma }_i$ not being normally distributed) does not have a trivial explanation. The exact mechanism not known and is empirical at present. 

Next step

Look for functional similarity between classes e.g. check if classes with low values of alpha are "writer" heads i.e. have high direct logit attribution (DLA) scores and high alpha are "aggregator" head i.e. consolidate information at specific positions in the residual stream or vice versa.

Classification B: Stats of ESD

Per‑layer aggregates (across heads) of ESD statistics (mean, standard deviation, skew, excess kurtosis) show a layer dependent trend i.e. $x_t\perp /x_{t-1}$. The goal of this post is to show that trends exist, so I report the BDS statistic and p-value^[7]. For the actual classification criterion I provide a boxplot^[8] and a conceptual description.

This method of classification has some obvious drawbacks compared to classification A, like being model specific, able to group only a subset of layers, and a lack of general metric to measure the goodness-of-fit. On the other hand, identifying functional roles for these groups may be more straightforward, since the observed patterns cannot be explained by architectural design choices and instead likely reflect emergent internal representations.

Case Studies

Excess Kurtosis Plain of Llama 3.1 405B-Instruct

The "plain" (Layers 33-48) is a region of abnormally low kurtosis values compared to the other layers. For gamma distribution, excess kurtosis ($k$) is equal to $\frac{6}{\alpha }$ so the attention heads in the "plain" have more rounded peak and thin tailed ESD whereas other layers have a power-law like shape ($y\propto e^{-t}$). Some possible interpretations are

1. The layers are not fully trained. (Martin et al., 2021)
2. The layers are generalist layers i.e. perform multiple functions depending on context.

SD Peaks of Qwen3 14B-Instruct and 32B-Instruct

The mean standard deviation of 14B shows a single peak in the middle layers of the model (Layers 18-19), the 32B model has two peaks at Layers 23 and 41. Layers can be classified by the section in which they lie.  

SD Slopes of Mistral 7B-Instruct

The mean standard deviation shows an almost monotonic decline (Spearman's r2 ≈ 0.98) for all 3 versions of the model, such a strong trend are not seen in other models of comparable size.

Mean flatlands of Gemma2 27B-it

The average of mean of ESD is almost constant $\sigma =0.03$. There is also an oscillatory pattern in the standard deviation of ESD which is probably caused by Gemma2 models using interleaved global and local attention layers.    

Related work

In the post "The Singular Value Decompositions of Transformer Weight Matrices are Highly Interpretable", the authors demonstrate that right singular vectors of $W_{OV}$ are interpretable using logit-lens.

Closing thoughts

The approach is complementary to the one taken in On the Biology of a Large Language Model. The authors of the paper describe recent work in interpretability as developing a better "microscope" to explain the inner workings of models. The approach proposed here is then akin to studying organs and organ systems.

I believe this approach is worth exploring, if not as a independent line of research then as a complementary path towards developing  efficient interpretability techniques. 

References

Conmy, A., Mavor-Parker, A., Lynch, A., Heimersheim, S., & Garriga-Alonso, A. (2023). Towards automated circuit discovery for mechanistic interpretability. Advances in Neural Information Processing Systems, 36, 16318-16352.

Martin, C. H., & Mahoney, M. W. (2021). Implicit self-regularization in deep neural networks: Evidence from random matrix theory and implications for learning. Journal of Machine Learning Research, 22(165), 1-73.

Men, X., Xu, M., Zhang, Q., Wang, B., Lin, H., Lu, Y., ... & Chen, W. (2024). Shortgpt: Layers in large language models are more redundant than you expect. arXiv preprint arXiv:2403.03853.

1. ^{^}

   The authors frame it as a proof of concept that partial automation of the mechanistic interpretability workflow is possible.

2. ^{^}

   The paper mentions 68 out of 32000 (0.21%) edges contribute to GPT-2 Small's ability to compute Greater-than operation .

3. ^{^}

   Men et al. (2024) show that 25% of the layers of Llama2-13B could be dropped with minimal degradation in performance.

4. ^{^}

   All models from Qwen3, Gemma-2, Mistral, and Llama 3.x that are dense and text-to-text. Yes, including Llama 3.1-405B.

5. ^{^}

   Technically $W_Q{W}_K^T$ has total $d_{model}$ singular values of which $d_{model}-d_{head}$ are always 0.

6. ^{^}

   No. of bins and clusters are hyperparameters. No. of bins = 24 and number of clusters = 6 is chosen for better visualization. 

7. ^{^}

   The Brock-Dechert-Scheinkman (BDS) statistic is tests whether a time series is i.i.d. or not. Positive values indicates the clustering in phase space negative values indicate repulsion. |BDS| ≥ 3 is considered strong evidence of dependence.

8. ^{^}

   For some models, Layers 0 and 1 distort the plots so I omit them across all models for consistency.