How Task Context Reshapes Component Roles in Language Models

The Mystery of Shared Attention Heads

There is a repeatedly observed phenomenon in large language models: the same set of attention heads gets reused across different tasks. Yet identifying these shared components doesn't explain what different roles they play in each task. In other words, we know which parts are reused, but we don't know whether they serve the same function in different contexts.

This research draws inspiration from neuroscience — shared brain regions in the human brain generate different interaction patterns depending on the task. The researchers transferred this thinking to language models, proposing a "network perspective" for understanding internal computation: rather than examining individual components in isolation, look at how they connect and influence each other within a specific task context.

An attention head is the basic unit of the multi-head self-attention mechanism in the Transformer architecture. Each attention head independently learns a set of Query, Key, and Value projection matrices, responsible for capturing dependencies between different positions in a sequence. Large language models typically deploy dozens of attention heads per layer across dozens of layers, meaning the total number of attention heads across the entire model can reach into the hundreds or even thousands.

Early interpretability research found that certain attention heads take on relatively fixed "functional roles" — for example, specifically tracking subject-verb agreement, copying earlier tokens, or retrieving entity attributes. However, in recent years, researchers conducting cross-task analyses have repeatedly found that the same set of heads activates on unrelated tasks, which contradicts the intuition of "functional specialization." This phenomenon, known as "multi-task reuse" of attention heads, is one of the core unsolved problems in mechanistic interpretability today.

Using Graph Facts for Controlled Experiments

To compare the behavior of fixed model components across different tasks, the study uses "graph facts" as a unified experimental vehicle. Graph facts are statements that explicitly specify which directed connections exist. By using graph facts with identical structure but different content, the researchers were able to observe differences in how the same set of components responds to different inputs under strictly controlled conditions.

The experiment introduces two key measurement dimensions:

• Influence: how an intervention changes the activations of an attention head and subsequent layers.
• Importance: how an intervention ultimately affects the model's answer.

This dual-metric design separates "internal mechanism changes" from "external output changes," allowing the analysis to both observe how components propagate signals internally and measure the actual contribution of that propagation to the final result.

Consistent Findings Across Five Open-Source Models

The study covers five open-weight models spanning four model families, with parameter counts ranging from 1.5B to 14B. The results show that changing the task redistributes the influence of components in these models. This means the roles of shared components are not fixed — they dynamically adjust with task context.

On Gemma-2-9B and Qwen3-14B, the researchers conducted more fine-grained activation editing experiments. By modifying early internal activation vectors while keeping written instructions unchanged, they successfully switched which fact the model was "reading." The same set of attention heads then responded to different graph facts and influenced downstream components in different ways.

This result is quite revealing: simply altering an early activation state can cause the model to attend to completely different information under identical instructions — all through the same shared set of heads. The role switch happens above the parameters, within the context.

The Predictive Power of Linear Response Maps

To quantify these task-dependent interactions, the researchers constructed two linear response maps for each prompt — one before the activation edit and one after. Each map predicts the response of later attention heads based on measured changes in earlier ones.

Key validation came from an independent factual intervention: the corresponding conditional response map, compared to the average of both maps, reduced mean squared prediction error by 53.6% on Gemma and 5.0% on Qwen. The large gap between the two models indicates that different architectures vary significantly in their sensitivity to task context. Gemma-2-9B's component role switching is cleaner and more linearly predictable, while Qwen3-14B's is considerably more complex.

A linear response map is an analytical tool inspired by the concept of "Green's functions" in physics and neuroscience: given a small perturbation to an early component, a first-order linear approximation is used to predict the change in response of downstream components. The core assumption is that within a local neighborhood, causal transmission relationships between components can be approximated by a matrix. The advantage of this method is that it can estimate the structure of the entire "influence network" directly from a small number of intervention experiments, without requiring exhaustive pairwise testing of every component pair.

The reduction in mean squared error (MSE) reflects the additional predictive value of the conditional response map over the unconditional average map. The 53.6% reduction on Gemma-2-9B means that task context shapes component interactions in this model very "cleanly," and the linear assumption holds well here. The mere 5.0% reduction on Qwen3-14B suggests its internal interactions may involve stronger nonlinearities or more diffuse task encoding, limiting the explanatory power of the linear response framework.

Implications for Interpretability and Model Intervention

These findings collectively point to one conclusion: task context simultaneously changes two things — which graph facts the shared components respond to, and how their outputs influence subsequent computation. Simple "component identification" is far from sufficient to characterize model behavior; a contextualized, network-based analytical framework is essential.

From a practical standpoint, this work has two layers of value. First, the network perspective may help explain a problem that has long puzzled researchers: how models execute fundamentally different computations using the same shared parameters. Second, measuring task-dependent influence can guide model interventions — modifying a specific behavior while minimizing disruption to others. This has potential applications in model editing, safety alignment, and behavioral control.

Of course, the enormous gap between Gemma and Qwen — 53.6% versus 5.0% — also reminds us that the generalizability of such linear response frameworks still needs validation across more models and tasks. But the methodology itself — treating attention heads as measurable network nodes rather than isolated parts — offers a promising path worth pursuing in mechanistic interpretability research.