Differentiation and Specialization of Attention Heads via the Refined Local Learning Coefficient

Thank you for the comments and clear indication of changes that would improve your score.

Larger scale. We note that the underlying technique of LLC estimation (non-refined) has been validated in deep linear networks up to 100M parameters $1$ . This means there is no technical obstacle to estimating refined LLCs at least at this scale, and the only current obstacles to scaling to much larger models is hyperparameter selection and computational cost, which are the subject of ongoing work. As mentioned in the top-level comment, we hope to include some preliminary results on Pythia models in a revision.

Importance. The final paragraph of the introduction has been rewritten with this suggestion in mind, drawing a closer connection to the developmental perspective emphasised in the final section.

Questions:

1. We are not sure (though we expect the same analysis to lead to similar observations in other settings). This would make for interesting follow-up work.
2. Equations:
   1. No, if $f\_w$ mapped directly to probabilities, then you would not need a softmax.
   2. This is actually an expression for the theoretical learning coefficient, so there is no need for a hat (see $1$ for more).
   3. Thank you for pointing this out. That is confusing! We have changed it to \epsilon in the revised version.

$1$ Z. Furman and E. Lau “Estimating the local learning coefficient at scale” 2024; as cited in the paper.

Following your suggestion we tested our methodology on Pythia-70M and you can find the results in Appendix H. We hope this addresses your question about scaling to larger models.

We thank the reviewer for their careful reading of the paper and detailed comments.

Meaning of terms. The description in the introduction has been changed to “which measures the complexity of a component of the model (for example an attention-head or whole layer) with respect to a given data distribution, which may differ from the training distribution. For example, we can measure the rLLC for a particular attention-head in a transformer trained on the Pile with respect to the distribution which conditions on a token sequence representing code.” We hope that clarifies also the sense in which the next paragraph refers to focusing on the special case where the component is an attention head.

Development on p.2. Thanks for this suggestion, in the new revision, we have attempted to explain more clearly how we intend this term to be read and provide some more context for why we adopt the developmental perspective at the end of p.2.

Thank you for your comments and questions.

Larger models. Here are some details on the computational cost for running these experiments, which we have added to the revision (appendix F.2): running an SGLD chain is approximately as expensive as running SGD training for the same number of steps. We estimate that the cost of producing rLLC curves like those in Fig 1 for all heads scales like C*H*log(N)*T, where C is a roughly constant O(100) number of SGLD samples, H is the number of heads, N is the total number of training steps, and T is the cost of a single training pass.

There is no essential difference between a “GPT-like” architecture and our models, except for scale and the inclusion of MLP layers. LLC estimation for transformer models with MLP layers was done in $1$ and we don’t expect any fundamental obstacle in performing rLLC estimates for MLP layers just as we have done in the present paper for attention heads. As mentioned in our top-level comment, we will share preliminary results for larger Pythia models in a coming revision.

Architecture optimization and training processes. We acknowledge that the focus of our paper is scientific understanding, and we have not thought about how these results might be used to inform the design of better architectures or training processes.

Comparison with existing interpretability metrics. Thank you for the references. To the best of our knowledge Michaud et al. does not introduce any interpretability metrics. Xiao et al seems very interesting, and we thank the referee for bringing it to our attention. We prioritised comparing the rLLC against two techniques commonly in use in the literature, namely Hessian-based complexity metrics and ablation interpretability analysis.

In Appendix D, we compared the rLLC as a complexity metric against the Hessian trace, the Fisher Information Matrix (FIM) Trace, and the Hessian Rank. We found that “the $r$ LLC appears to present a clearer picture of the model’s development than the Hessian trace” and that the FIM Trace looks similar to the Hessian trace but noisier. Our Hessian Rank results were inconclusive due to difficulties in estimation.

In Appendix E, we compared the use of the rLLC as an interpretability tool against zero, mean, and resampling ablations, as well as path-patching. We found that these four ablation techniques taken together allow distinguishing heads and stages to a similar extent to the rLLC, with some heads more easily distinguished using the rLLC, and some more easily using ablations. Note, however, that these two techniques are complementary, as the rLLC operates in weight space and ablations in activation space. Finally, we want to reiterate that ablations are somewhat more computationally efficient, yet are less theoretically grounded (reflected in the fact that one has to choose which ablation to run).

$1$ Hoogland et al, 2024; cited as in the paper.

We appreciate the comments and suggestions for improvement.

Bigger transformers with more layers: the rLLC can be measured for transformers at any scale, albeit with more computational cost and the hyperparameter selection process may become more difficult at larger scale. The inclusion of MLP layers or deeper networks does not present a serious obstacle. In this paper we concentrate on smaller models, where it is easier to have a comprehensive understanding of the ground truth of the model’s behaviour, in order to have something to compare our rLLC measurements to. We expect the same techniques to apply to larger models, and as mentioned in the top-level comment, hope to include some preliminary results in a revision during the discussion period.

Different architectures: the rLLC methodology can indeed be straightforwardly extended to study the development of other architectural components, including convolutional layers. While the current paper focuses on attention heads to maintain a clear scope and narrative, we agree this would be interesting future work.

K-means clustering: The K-means clustering was performed on the rLLC trajectories themselves, treating each head's concatenated Pile and Github rLLC curves as a point in a 2T-dimensional space (where T is the number of checkpoints). We found that using both datasets' trajectories helped better distinguish the functional types of heads (using individual datasets also works, but leads to more varied outcomes over different clustering algorithms). For transparency, we have included detailed clustering methodology in Appendix C.2, including comparisons with other clustering approaches like DTW K-means and Shape-Based K-means that are specifically designed for time-series data. The clusters are robust across these different methods and also generalize to other training runs (Appendix F).

Beyond interpretability: we haven’t considered any of these yet, but the connection to model pruning seems particularly promising to us.