Compressed Sensing for Capability Localization in Large Language Models

Large language models (LLMs) exhibit a wide range of capabilities, including mathematical reasoning, code generation, and linguistic behaviors. We show that Transformer architectures contain small subsets of attention heads that are necessary for certain capabilities. Zeroing out as few as five task-specific heads can degrade performance by up to $60\\%$ on standard benchmarks measuring the capability of interest, while largely preserving performance on unrelated tasks. We introduce a compressed sensing-based method that exploits the sparsity of these heads to identify them via strategic knockouts and a small number of model evaluations. We validate these findings across Llama and Qwen models ranging from 1B to 14B parameters and a diverse set of capabilities including mathematical abilities and code generation, revealing a modular organization in which specialized capabilities are dependent on sparse, functionally distinct components. Overall, our results suggest that capability localization is a general organizational principle of Transformer language models, with implications for interpretability, model editing, and AI safety. Code is released at https://github.com/locuslab/llm-components.

Large language models can do many different things, including solving math problems, and writing code. But how does a single model handle all of these abilities? How is the processing of these various abilities organized within the model? We found that for many capabilities, a surprisingly small number of components inside the model are critical for performance. Just five out of more than a thousand "attention heads" can be responsible for a model's ability to do math. Disable those five, and the model's math performance collapses, but its ability to answer general knowledge questions or generate fluent text stays intact. Finding these critical components could require testing them one at a time, which is slow. We borrowed a technique called compressed sensing, which was originally designed to reconstruct signals from very few measurements, and used it to identify these components much more efficiently. Our method works across different models of varying sizes and several capabilities, and could be used to help researchers better understand how language models work.