Compressed Sensing for Capability Localization in Large Language Models

2026-02-11

arXiv.org

AI Summary

This work addresses the challenge of efficiently identifying critical components responsible for specific capabilities—such as mathematical reasoning or code generation—in large language models. Introducing compressed sensing to capability analysis, the study leverages the inherent sparsity of attention heads in task representation, enabling precise identification of a minimal subset essential to a given ability with only limited evaluation. Through attention head ablation, sparsity analysis, and cross-model validation across Llama and Qwen architectures (spanning 1B to 14B parameters), the authors demonstrate that disabling as few as five key attention heads can degrade target capabilities by up to 60%, while leaving performance on unrelated tasks largely unaffected. These findings reveal a general principle of functional modularity and sparsity within the Transformer architecture.

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.

Problem

• capability localization
• compressed sensing
• attention heads
• large language models
• modular organization

Innovation

• Compressed Sensing
• Capability Localization
• Sparse Attention Heads
• Transformer Modularity
• Model Interpretability