Can Capability Enhancement be Localized?

Transformer-based language models have shown strong capabilities across many tasks and domains. However, prior work shows that, for a given task, not all attention heads matter equally. Specifically, they suggest that many $\mathit{\text{redundant}}$ heads can be safely removed with little-to-no effect, while a small set of $\mathit{\text{important}}$ heads, when ablated, severely damages the model’s performance on the given task. In this work, we investigate a complementary aspect, called $\mathit{\text{Capability}}\mathit{\text{Enhancement}}\mathit{\text{Localization}}$ (CEL): $\mathit{\text{given}}\mathit{\text{a}}\mathit{\text{model}}\mathit{\text{and}}\mathit{\text{a}}\mathit{\text{task,}}\mathit{\text{does}}\mathit{\text{there}}\mathit{\text{exist}}\mathit{\text{a}}\mathit{\text{small}}\mathit{\text{set}}\mathit{\text{of}}\mathit{\text{heads}}\mathit{\text{whose}}\mathit{\text{removal}}\mathit{\text{improves}}\mathit{\text{the}}\mathit{\text{task}}\mathit{\text{performance?}}$ We first investigate this question $\mathit{\text{empirically}}$ across models from different families. Using a simple two-stage routine that screens candidate heads and evaluates their joint ablations, we identify sparse head combinations whose removal improves task-specific performance without further fine-tuning. Motivated by these observations, we ask $\mathit{\text{theoretically}}\mathit{\text{whether}}\mathit{\text{such}}\mathit{\text{a}}\mathit{\text{phenomenon}}\mathit{\text{can}}\mathit{\text{naturally}}\mathit{\text{arise}}\mathit{\text{from}}\mathit{\text{training}}\mathit{\text{dynamics,}}\mathit{\text{and}}\mathit{\text{under}}\mathit{\text{what}}\mathit{\text{conditions}}$. To this end, we construct a simplified multi-head softmax attention model and two tasks with conflicting features, and show that gradient descent from small random initialization provably concentrates the conflicting features into particular heads. We further show that ablating such heads improves performance on one task while harming the other, even though training converges to a global optimum of the joint objective. This establishes a possible training-induced theoretical explanation for CEL.