Attention Pattern Discovery at Scale

Jonathan Katzy, Razvan Mihai Popescu, Erik Mekkes, Arie van Deursen, Maliheh Izadi

Keywords: Applications of interpretability, Automated interpretability, Vision transformers

TL;DR: We mine attention patterns at a large scale, encode them using a novel model, and use this to differentiate correct from incorrect predictions of an LLM.

Abstract:

Language models have scaled rapidly, yet methods for explaining their outputs are lagging behind. Most modern methods focus on a fine-grained explanation of individual components of language models. This is resource-intensive and does not scale well to describe the behavior of language models as a whole.
To enable high-level explanations of model behavior, in this study, we analyze and track attention patterns across multiple predictions.
We introduce Attention Pattern Masked AutoEncoder (AP-MAE), a vision-transformer–based approach that encodes and reconstructs large language model attention patterns at scale. By treating attention patterns as images, AP-MAE enables efficient mining of consistent structures across a large number of predictions.

Our experiments on StarCoder2 models (3B–15B) show that AP-MAE (i) reconstructs masked attention with high fidelity, (ii) generalizes across unseen model sizes with minimal degradation, and (iii) predicts whether a token will be correct, without access to ground truth, with up to 70% accuracy. We further discover recurring attention patterns demonstrating that attention patterns are structured rather than random noise. These results suggest that attention maps can serve as a scalable signal for interpretability, and that AP-MAE provides a transferable foundation for analyzing diverse large language models. We release code and models to support future work in large-scale interpretability.

Submission Number: 267