Compressed Sensing for Capability Localization in Large Language Models

TL;DR

This paper demonstrates that specific capabilities of large language models are localized to small, sparse subsets of attention heads, and introduces a compressed sensing method to identify these heads efficiently, revealing a modular organization.

Contribution

The paper presents a novel compressed sensing approach to identify task-specific attention heads, showing capabilities are localized and modular within Transformer models.

Findings

Zeroing five heads can reduce task performance by 65%

Capabilities are localized to small, sparse head subsets

Models range from 1B to 8B parameters, showing consistent localization

Abstract

Large language models (LLMs) exhibit a wide range of capabilities, including mathematical reasoning, code generation, and linguistic behaviors. We show that many capabilities are highly localized to small subsets of attention heads within Transformer architectures. Zeroing out as few as five task-specific heads can degrade performance by up to $65\%$ 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 8B parameters and a diverse set of capabilities including mathematical abilities and code generation, revealing a modular organization in which specialized capabilities are implemented by 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.