Task Context Shapes Component Roles in Language Models

Sovesh Mohapatra; Dani S Bassett

cs.CLOct 9, 2026

Task Context Shapes Component Roles in Language Models

Authors: Sovesh Mohapatra, Dani S Bassett

Abstract

Large language models reuse the same attention heads across tasks, but identifying these shared heads does not explain what they do differently in each task. A natural source of inspiration is task-dependent interactions among shared brain regions, which motivate a network view of language-model computation. Here, we compare fixed model components across tasks using identical graph facts, statements specifying which directed connections are present. We measure how interventions affect attention heads and the activation state in later layers (influence), as well as the answer (importance). Task changes redistribute the influence in five open-weight models spanning four families and 1.5-14 billion parameters. In Gemma-2-9B and Qwen3-14B, an activation edit changes an early internal activation vector and switches which fact the model reads while the written instruction stays fixed. The same heads then respond to different graph facts and influence later components differently. For each prompt, we construct two linear response maps, one before the edit and one after it. Each map predicts a later head's response from a measured change in an earlier head. For a separate factual intervention, the map for the corresponding condition reduces mean squared prediction error by 53.6% in Gemma and 5.0% in Qwen compared with the average of the two maps. These findings show that task context changes both which graph facts the shared components respond to and how their outputs affect later computation. A network account may help explain how models use shared parameters to perform different computations. Measurements of task-dependent influence could also guide interventions intended to alter one behavior while preserving others.

Figures & tables

Figure 1: Task-dependent computation with shared components. A. Same components, different interactions. Brain networks motivate the hypothesis that shared model components support different tasks through changes in their interactions. H1–H3 denote fixed attention heads in layer order; Res. denotes a later residual-stream site. B. Measure task-dependent influence. Identical graph facts yield True for Compose path and False for Read direct edge : the path $A\to B\to C$ exists, whereas the direct edge $A\to C$ is absent. Under each instruction, source perturbations $\pm\delta$ produce responses at later sites that quantify influence, with intervening computation active. C. Change the internal task state. An early activation edit changes $h$ to $h^{\prime}=h+\Delta_{\mathrm{task}}$ while the prompt and model weights remain fixed. The vector $\Delta_{\mathrm{task}}$ is the mean activation under Read fact 2 minus the mean activation under Read fact 1 on reference prompts. R1 and R2 denote two fixed attention heads. The schematic illustrates an answer change from True (Fact 1) to False (Fact 2), while the written instruction remains Read fact 1 .

Figure 2: Task changes redistribute influence among shared components. A. Influence changes with the task. Falcon-H1-1.5B shows different influence patterns under Compose path and Read direct edge with identical graph facts. H1–H3 denote fixed attention heads in layer order; Res. denotes a later residual-stream site. Arrow labels give mean influence in units of $10^{-3}$ , and widths indicate magnitude on a shared scale. Solid and dashed arrows denote positive and negative responses, respectively. B. Where influence changes. Color shows the difference in influence ( Read direct edge minus Compose path ) in four additional models. Each column is normalized to unit Euclidean norm, showing where the task-dependent change is concentrated. C. Influence profiles rotate more. Points show the angle between task profiles for the six influence values (filled) and four component-importance values (open). Whiskers show 95% confidence intervals from 10,000 paired bootstrap resamples of vocabulary groups.

Figure 3: Later activations support task-selective factual effects. A. Task-selective factual effects. Cells show task selectivity in Gemma-2-9B and Qwen3-14B: the control-adjusted answer effect under the relevant instruction minus that under the irrelevant instruction. The edges $A\to B$ and $B\to C$ are relevant to Compose path , whereas $A\to C$ is relevant to Read direct edge . B. Downstream contribution. The factual intervention remains active while selected later residual activations are restored to their values from the same prompt without the intervention. Points show the reduction in task selectivity under restoration minus the mean reduction under restoration controls. Brackets in A and whiskers in B show pointwise 95% confidence intervals from 10,000 paired bootstrap resamples of vocabulary groups.

Figure 4: Internal task edits redirect factual responses and influence. Released Gemma-2-9B and Qwen3-14B weights remain fixed during all interventions. A. Same heads, different facts. Cells show mean response shares at fixed heads H2 and H3 before and after the early edit, with the written instruction held at Read fact 1 . Shares are normalized over squared head-output responses to changes in all three facts. B. Shifts grow with depth. Dots show each head’s mean factual-reassignment score, and lines show layer medians across all 672 Gemma heads and 1,600 Qwen heads. Dashed lines mark the edit layer. C. Influence shifts toward the target. Squared distances compare the six-pair influence profile with the unedited profile under the alternative instruction. Thin lines represent vocabulary groups, and thick lines show means. D. Individual head restoration. Points show control-adjusted increases in the original-answer margin, in logits, after individual H2 or H3 restoration while the early edit remains active. Whiskers show approximate simultaneous 95% confidence intervals.

Figure 5: Context shapes prediction and downstream responses. A. Task-specific prediction. Response maps in released Gemma-2-9B and Qwen3-14B predict the H3 output change caused by a separate factual intervention at H2. Errors are divided by each model’s mean error for predicting no H3 change, shown by the dashed line at 1. Thin lines join paired vocabulary means, and thick lines show overall means. Percentages give reductions in mean squared error relative to the averaged map. B. Relevance strengthens responses. Gemma and Qwen fine-tuned on graph tasks receive the same H2 displacement in the relevant and irrelevant contexts under AND. Both axes show squared H3 responses divided by the model’s mean response in the irrelevant context. Circles represent vocabulary means, diamonds show overall means, and the dashed diagonal marks equal responses.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure S1: Source controls and downstream restoration. A. Factual effects follow task relevance. Points show signed answer effects of factual interventions and two source controls under Compose path and Read direct edge , before source-control subtraction. Positive effects favor True for an absent-to-present intervention and False for a present-to-absent intervention. B. Restoration reduces task selectivity. Task selectivity is shown without a downstream intervention, under restoration of the original response-key activations, and under two restoration controls. Every estimate in panel B includes the source-control adjustment, recalculated for that downstream condition. The factual intervention remains active during restoration. Each source or restoration control combines a random permutation of the displacement’s dimensions with random sign flips, preserving its magnitude at each intervened position. Both panels use the models, prompts, and intervention sites from Figure 3 . Points show means over 12 vocabulary groups, with both wordings averaged within each group. Whiskers show pointwise 95% confidence intervals from 10,000 paired bootstrap resamples of vocabulary groups.

Model	Head	Fact 1	Fact 2	Fact 3
Gemma	H2	$81.638\to 8.427$	$6.459\to 83.450$	$3.773\to 5.242$
Gemma	H3	$16.421\to 0.141$	$0.208\to 16.674$	$0.120\to 0.138$
Qwen	H2	$23.356\to 2.930$	$3.812\to 22.692$	$2.665\to 2.250$
Qwen	H3	$63.412\to 2.071$	$1.556\to 66.308$	$1.138\to 1.942$

Appendix

Table S1: Absolute factual-response energies. Entries show unedited

\to

task-edited mean squared head-output responses under Read fact 1 , averaged over 16 vocabulary groups.

Figure S2: Factual responses under both internal edit directions. Cells show each graph fact’s share of squared head-output responses. Each share uses the sum over all three facts as its denominator and is calculated within a vocabulary group and written instruction before averaging over the 16 groups from Figure 4 A. The task edit targets Fact 2 under Read fact 1 and Fact 1 under Read fact 2 , while the written instruction remains fixed. The two random early edits match the task edit’s displacement magnitude and resulting-state magnitude. The same H2 and H3 heads are used throughout: H26.09 and H32.07 in Gemma, and H24.17 and H32.13 in Qwen, with zero-based layer and head indices.

Explore similar work

Sep 28, 2026cs.CL

A mechanistic study of language model introspection

Large language models (LLMs) can sometimes report perturbations to their internal activations---even when the input provides no evidence that an intervention occurred. How do models detect and localize such internal changes? We study this question using a controlled task that keeps the input text fixed. We either inject a concept vector into the hidden state at one of ten token positions or apply no intervention. The model is asked to identify the perturbed position or report that no intervention occurred. Across three model families, we identify two small groups of attention heads with distinct roles in introspective reporting. Middle-layer gate heads influence whether the model reports a change, while router heads in a later layer help select the position to report. Interventions on gate heads can suppress position reports even when router heads supply location information. We further examine why reporting accuracy varies across concepts. Concept vectors that are localized more accurately produce stronger attention-score and output responses in gate heads, which is associated with better alignment of the induced key and value changes in their QK and OV computations. Together, these findings identify attention-head mechanisms supporting introspective detection and localization.

Jiahong Zou, Xiangkun Sun, Lingkai Kong +1

Shandong University · Tsinghua University · Northeastern University +1

Attention Head Analysis LLM Interpretability Language Model Introspection Metacognition in Language Models

Ask PDF

Apr 23, 2026cs.CL

Shared Lexical Task Representations Explain Behavioral Variability In LLMs

One of the most common complaints about large language models (LLMs) is their prompt sensitivity -- that is, the fact that their ability to perform a task or provide a correct answer to a question can depend unpredictably on the way the question is posed. We investigate this variation by comparing two very different but commonly-used styles of prompting: instruction-based prompts, which describe the task in natural language, and example-based prompts, which provide in-context few-shot demonstration pairs to illustrate the task. We find that, despite large variation in performance as a function of the prompt, the model engages some common underlying mechanisms across different prompts of a task. Specifically, we identify task-specific attention heads whose outputs literally describe the task -- which we dub lexical task heads -- and show that these heads are shared across prompting styles and trigger subsequent answer production. We further find that behavioral variation between prompts can be explained by the degree to which these heads are activated, and that failures are at least sometimes due to competing task representations that dilute the signal of the target task. Our results together present an increasingly clear picture of how LLMs' internal representations can explain behavior that otherwise seems idiosyncratic to users and developers.

Zhuonan Yang, Jacob Xiaochen Li, Francisco Piedrahita Velez +5

Brown University · MIT · Northeastern

Prompt Sensitivity LLM Prompting LLM Interpretability Few-Shot Prompting Mechanistic Interpretability Attention Mechanisms

Ask PDF

Sep 28, 2026cs.LG

How Language Models Differ in Redistributing Attention-Head Activity Under Serial Demand

The way a model distributes activity over each layer's attention heads offers a coarse view of how it routes information through depth; how this changes with the task is part of what a mechanistic account must explain. Holding prompt length fixed, we vary how many serial steps a task demands and measure, in every layer of 17 open-weight models, whether activity concentrates on a few heads or spreads across many as demand rises. Both occur: in most models, layers just before mid-depth concentrate activity and later layers spread it. Models differ in where and how strongly this happens. The Qwen2.5 base models from 0.5B to 7B, for example, spread less than the average model in every task and concentrate activity in parts of their second half, where Llama models from 1B to 8B and OLMo-2 spread; the contrast largely holds between Llama-3.1-70B and Qwen2.5-72B, which have the same number of layers and heads. These differences are reproducible, and post-trained models keep much of their base model's pattern. An ablation study suggests that, within a task, models whose activity is more concentrated on their top heads also depend more on those heads for the answer. Concentration and spreading across layers thus offer a new way to compare models, by how they route information through depth. Code is available at https://github.com/johnnyjli/serial-demand-heads.

Johnny Jingze Li, Abdulla Kuleib, Kalyan Basu +1

Carnegie Mellon University · University of California, San Diego · Qualtrics LLC

Attention Head Analysis LLM Interpretability Attention Head Specialization

Ask PDF