Source-linked AI summary

Task Context Shapes Component Roles in Language Models

Sovesh Mohapatra, Dani S Bassett

arXiv:2610.13895v1cs.CLcs.AI

TL;DR

Tasks can reuse attention heads, yet it remains unclear how their roles change across tasks. The paper compares fixed components on identical graph facts and measures their influence and answer importance. Task context changes both factual responses and downstream influence, while context-matched response maps reduce prediction error in Gemma and Qwen.

  • Problem

    Although tasks can reuse attention heads, it remains unclear whether the same components respond to different facts or influence later computation differently across tasks.

    Source: §1 · p. 2
  • Method

    The study compares fixed model components across tasks using identical graph facts, measuring their influence and answer importance, and edits internal task state in Gemma and Qwen.

    Source: §1 · p. 2
  • Results

    Task changes alter influence across five models and redirect factual responses in Gemma and Qwen; matched response maps reduce mean squared error by 53.6% in Gemma-2-9B and 5.0% in Qwen3-14B.

    Source: §1 · pp. 2–3, §3.4 · p. 9
  • Takeaways & Limitations

    Shared components’ roles depend on internal task state, and maps matched to that state predict downstream factual effects more accurately within the same prompt.

    Source: §4 · p. 9

    The conditional-response finding is limited to Gemma and Qwen fine-tuned on graph tasks and evaluated under the AND instruction.

    Source: §4 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Tasks can reuse the same attention heads, but it remains unclear whether task context changes which facts those heads respond to and how they influence later computation. Across five models, task changes alter influence patterns; in Gemma and Qwen, internal edits also redirect factual responses while the written instruction stays fixed.

  • Tasks can reuse attention heads and circuit components (Merullo et al., 2024a), but which facts shared components respond to and how their outputs affect later components remains an open question.Source: §1 · pp. 1–2
  • A network perspective focuses on task-dependent interactions among shared components, and the study compares fixed components across tasks using identical graph facts.Source: §1 · pp. 1–2
  • Across five models, changing the task changes influence patterns among the same selected components while graph facts remain fixed.Source: §1 · p. 2
  • In Gemma and Qwen, factual interventions affect answers more strongly when the targeted graph fact is task-relevant, and internal edits redirect head responses and influence.Source: §1 · p. 3
  • Within the same prompt, a response map matched to the current internal context predicts a later head’s response to a factual intervention more accurately than averaging the two maps.The fixed earlier-head perturbation also produces a larger later-head response when its targeted fact can affect the answer.Source: §1 · p. 3

2 METHODS

The methods compare fixed model components across graph-reading tasks, measuring their influence on later computation and importance for answers. Interventions then test task-selective factual effects, edit internal task context, and predict downstream responses.

2.1 TASKS AND MODELS

  • The tasks ask whether A →B and B →C form a path or whether the direct edge A →C is present, using identical graph facts across task instructions.Five open-weight models span four families and 1.5–14 billion parameters; Gemma and Qwen are used for factual interventions, internal edits, and response prediction.Source: §2.1 · p. 3

2.2 INFLUENCE AND IMPORTANCE

  • Influence is the signed later-site response per unit source perturbation, while importance measures the correct-answer margin decrease after replacing a site's activation with its reference mean.The six influence values and four importance effects form profiles whose angle between task means measures changes in relative effect patterns.Source: §2.2 · pp. 3–4

2.3 FACTUAL INTERVENTIONS AND DOWNSTREAM RESTORATION

  • Factual interventions alter residual activations across a graph-fact statement, and task selectivity is the adjusted answer effect under the relevant task minus that under the irrelevant task.Downstream restoration replaces later residual activations with same-prompt values from runs without the intervention; its contribution is the reduction in task selectivity beyond restoration controls.Source: §2.3 · pp. 4–5

2.4 INTERNAL TASK EDITS AND HEAD CONTRIBUTIONS

  • A task-vector edit adds or subtracts the mean reference-activation difference between Read fact 2 and Read fact 1 while keeping the written prompt fixed.The analysis tracks head responses to each fact, compares influence profiles with the alternative-instruction target, and restores H2 or H3 outputs to assess their answer contribution.Source: §2.4 · pp. 5–7

2.5 PREDICTION AND CONDITIONAL RESPONSES

  • For each prompt, a response map predicts H3's output change from a measured H2 displacement; the map for the current context is compared with the average of edited and unedited maps.A separate analysis holds the H2 displacement fixed while changing whether another edge makes the targeted fact relevant to the AND-task answer.Source: §2.5 · p. 7

3 RESULTS

Across tasks, influence shifts among shared components, and early internal edits redirect factual responses and downstream influence while the written instruction remains fixed. Context-matched response maps better predict downstream factual effects, while later-activation restoration reduces task selectivity.

3.1 TASK CHANGES REDISTRIBUTE INFLUENCE

  • Changing instructions redistributes influence: in Falcon-H1-1.5B, H1→H2 shifts from −0.0145 to 0.00101, while H1→H3 rises from 0.0107 to 0.0238.Across the other models, H2→H3 influence increases in Gemma and Llama, while H2’s influence on the later residual site decreases in Qwen and Falcon-H1-7B.Source: §3.1 · p. 7
  • With the residual site included, influence profiles rotate more than importance profiles in all four models; paired angle differences are positive and Bonferroni-adjusted intervals exclude zero.The residual site contributes 97.8–99.95% of the importance profiles’ squared norm.Source: §3.1 · p. 7

3.2 LATER ACTIVATIONS SUPPORT TASK-SELECTIVE FACTUAL EFFECTS

  • Factual interventions have task-selective effects for all three graph facts, reaching 17.24 [16.81, 17.73] logits for A →C in Gemma and 52.46 [50.83, 54.36] logits in Qwen.Downstream restoration reduces task selectivity more than controls for all three facts in both models; for A →C, reductions are 8.60 [7.97, 9.33] logits in Gemma and 51.30 [49.74, 53.08] logits in Qwen.Source: §3.2 · p. 8

3.3 INTERNAL TASK EDITS REDIRECT FACTUAL RESPONSES AND INFLUENCE

  • Under fixed Read fact 1, edits raise Fact 2 response share by 78.9 and 97.1 percentage points at H2 and H3 in Gemma, and 68.7 and 91.8 in Qwen.Across both edit directions, each model switches to the alternative instruction’s answer on all 512 opposite-answer prompts and preserves the correct answer on all 512 same-answer prompts.Source: §3.3 · p. 8
  • After the edit, mean squared distance to the unedited influence profile under the alternative instruction falls by 93.2% in Gemma and 76.7% in Qwen; both are also lower than under random early edits.Individual head restoration produces small control-adjusted increases in the original-answer margin while the early edit remains active (Figure 4D).Source: §3.3 · p. 8

3.4 CONTEXT SHAPES PREDICTION AND DOWNSTREAM RESPONSES

  • Context-matched maps reduce mean squared prediction error by 53.6% in Gemma-2-9B and 5.0% in Qwen3-14B; relevant-to-irrelevant H3 response ratios are 1.40 [1.31, 1.51] and 1.53 [1.34, 1.74].The response-map comparison uses a separate factual intervention; the H3 ratios compare the same H2 displacement under relevant and irrelevant contexts.Source: Figure 5 · p. 8, §3.4 · p. 9

3.3 INTERNAL TASK EDITS REDIRECT FACTUAL RESPONSES AND INFLUENCE

  • With the early edit active, H2/H3 restoration shifts the written-instruction answer margin by 0.203/0.288 logits in Gemma and 0.299/0.661 logits in Qwen.These effects are 1.13%/1.61% of the mean edit-induced margin shift in Gemma and 0.62%/1.37% in Qwen.Source: §3.3 · p. 9

4 DISCUSSION

The discussion frames component roles as dependent on the network state receiving their outputs, and notes that current-state response maps better predict downstream factual effects within the same prompt. It limits these findings to the tested settings and points to transfer across prompts, graphs, and longer reasoning as future work.

4 DISCUSSION

  • A shared component’s role depends on the state of the network receiving its output, and influence and answer importance can change by different amounts across tasks.In Gemma and Qwen, early activation edits shift factual responses and influence while the written instruction remains fixed.Source: §4 · p. 9
  • Task-specific response maps improve within-prompt prediction of downstream factual effects in Gemma and Qwen without additional training.The discussion interprets the finding as evidence that task context helps explain how an intervention propagates despite a fixed component and imposed change.Source: §4 · p. 9
  • Influence values depend on the chosen activation direction and perturbation size; conditional-response results are limited to fine-tuned Gemma and Qwen under AND.The prediction analysis uses a source displacement and response map measured for the same prompt.Source: §4 · p. 9
  • Future work could test response-map transfer to new graphs and instructions without prompt-specific recalibration, and examine feature interactions and longer passages.The proposed extensions would test whether context-dependent interactions predict which information contributes to answers as task demands change.Source: §4 · pp. 9–10

AI USE DISCLOSURE

The authors disclose using generative AI for initial literature review, code assistance, and manuscript readability edits, and report reviewing the AI-assisted work and verifying the code.

AI USE DISCLOSURE

  • The authors report reviewing AI-assisted work, verifying and testing the code, and taking responsibility for the final content.AI tools were used for initial literature review, code assistance, and readability edits.Source: AI USE DISCLOSURE · p. 10

A TASKS, PROMPTS, AND EVALUATION SETS

The appendix specifies controlled graph-task prompts and evaluation sets, then documents how attention heads and activation directions are selected and scaled.

A.1 TASK CONSTRUCTION AND PROMPT EXAMPLES

  • Paired tasks preserve the graph facts and response key while changing the requested operation, including direct-edge, path, and numbered-fact questions.For the internal-edit prompts, fact numbers stay attached to the same edges when statements are reordered.Source: §A.1 · p. 13
  • Across eight graphs, fact-reading tasks require four True and four False answers; Compose path and AND require two True and six False, with both response keys balanced.In the conditional-response example, the answer depends on the first selected edge only when the second is present.Source: §A.1 · p. 14

A.2 MODEL CHECKPOINTS AND INPUT FORMATTING

  • Model formatting differs by family, and answer effects use first-answer-token logits with Yes/No or A/B labels depending on the analysis.Gemma uses one user message; Qwen, Llama, and Falcon-H1 use separate system and user messages, and Qwen disables thinking mode.Source: §A.2 · p. 14

A.3 CALIBRATION AND EVALUATION SETS

  • Evaluation sets use analysis-specific vocabulary groups and prompt combinations across influence, factual-intervention, internal-edit, and conditional-response measurements.Figure 2 adds eight groups for Falcon-H1-1.5B and 16 for other models; the factual-intervention analysis uses 12 groups, and internal-edit analyses use 16.Source: §A.3 · pp. 14–15

A.4 BASELINE TASK PERFORMANCE

  • Baseline correct counts are Gemma 119/128, Qwen 125/128, Llama 116/128, Falcon-H1-1.5B 57/64, and Falcon-H1-7B 128/128.Figure 3 correctness is 2,284/2,304 (99.1%) for Gemma and 2,257/2,304 (98.0%) for Qwen; both models answer all 1,024 Figure 4A/4C/4D and Figure 5A prompts correctly.Source: §A.4 · p. 15
  • Figure 5B baselines are Gemma 767/768 (99.9%) and Qwen 765/768 (99.6%); validation gives Qwen 62/64 True and 192/192 False, while Gemma answers all correctly.The validation counts pool both fact wordings, response keys, and statement orders.Source: §A.4 · p. 15

B.1 COMPONENT LOCATIONS AND SELECTION

  • Head trios are selected across three depth ranges using answer-direction separation and projection-removal damage, with eligibility criteria favoring multiple contributing heads.Candidates are evaluated with leave-one-vocabulary-group-out directions; an eligible trio has positive damage in all four groups and at least two heads with positive mean additional contributions.Source: §B.1 · pp. 15–16

B.2 REFERENCE DIRECTIONS AND ACTIVATION SCALES

  • Reference activations define shared task directions and scales; the task direction removes only the component along one linear answer direction, and one coordinate unit equals one standard deviation.Directions, means, and scales remain the same under both instructions.Source: §B.2 · p. 16

B.3 INTERVENTIONS AND TASK PROFILES

Influence measures a later site's signed response to a source-head perturbation, while importance measures the answer-margin effect of replacing a component.

B.3 INTERVENTIONS AND TASK PROFILES

  • Influence uses finite positive and negative two-unit shifts along each source head's task direction, with intervening computation active.The response is measured in recipient coordinates, so it includes transmission through other components rather than only an isolated direct connection.Source: §B.3 · p. 16
  • Importance replaces a selected head or residual vector with its reference mean and measures the resulting decrease in correct-answer margin.Negative values are retained when replacement increases the correct-answer margin; the profiles contain six influence entries and four importance entries.Source: §B.3 · p. 16
  • The angle between task-mean profiles compares relative patterns, not absolute effect magnitudes.A uniform positive rescaling leaves the angle unchanged; Figure 2B normalizes between-task influence differences within each model to show which site pairs account for the change.Source: §B.3 · pp. 16–17

B.4 CONTRIBUTION OF THE RESIDUAL SITE TO PROFILE ORIENTATION

The residual site dominates the measured importance profiles, while factual activation displacements are constructed from paired prompts that differ in one graph fact.

B.4 CONTRIBUTION OF THE RESIDUAL SITE TO PROFILE ORIENTATION

  • 97.8–99.95% of the four-site importance profile's squared norm comes from the residual site, a measured-profile share rather than a fraction of total model computation.In Gemma, its replacement effect is 8.04 [7.80, 8.28] logits for Compose path and 3.31 [2.54, 4.01] for Read direct edge; the paired decrease is 4.73 [3.98, 5.60] logits.Source: §B.4 · p. 17
  • Head-only profile-angle estimates are descriptive; all four pointwise 95% paired-bootstrap intervals for angle differences include zero.These exploratory profiles exclude the residual site from both influence and importance profiles.Source: §B.4 · p. 17

C.1 FACTUAL ACTIVATION DIFFERENCES

  • Gemma and Qwen factual displacements come from paired reference runs differing in one fact, span the complete fact statement, and are applied under both instructions while the prompt stays fixed.The interventions act at the input to decoder layer 12; controls permute displacement dimensions and flip signs while preserving magnitude.Source: §C.1 · p. 17

C.2 ANSWER EFFECTS AND TASK SELECTIVITY

The analysis compares adjusted factual effects under relevant and irrelevant instructions, then tests how later-state restoration changes task selectivity.

C.2 ANSWER EFFECTS AND TASK SELECTIVITY

  • Answer effects use a True-minus-False logit score oriented by the response key; adjusted effects subtract mean source-control effects, and task selectivity contrasts relevant with irrelevant instructions.Relevance depends on whether the correct answer depends on the targeted edge; for A →B, Compose path is relevant when B →C is present, whereas Read direct edge is independent of A →B.Source: §C.2 · pp. 17–18

C.3 LATER ACTIVATION RESTORATION

  • Later restoration replaces residual activations from an unaltered run while the layer-12 factual intervention remains active, with randomized displacement controls applied in separate runs.Restoration is at layer 23 in Gemma and layer 32 in Qwen, targeting either the task question or response-key span.Source: §C.3 · p. 18
  • 6.87 [6.28, 7.39] logits is the controlled contribution of later restoration to Gemma's A →B task selectivity.Selectivity is 13.14 [12.72, 13.52] logits without downstream intervention, 6.43 [6.06, 6.83] under restoration, and 13.30 [12.88, 13.68] under restoration controls.Source: §C.3 · p. 18

D.1 CONSTRUCTION OF THE TASK EDIT

An early task-vector edit changes which fact the model reads while leaving the written instruction fixed, enabling comparisons of head responses, downstream influence, and answer contributions across conditions.

D.1 CONSTRUCTION OF THE TASK EDIT

  • The task vector is the mean activation difference between fact-reading instructions, added or subtracted across the edited residual span without additional scaling.Random early edits match both displacement and resulting-state magnitudes.Source: §D.1 · pp. 18–19

D.2 RESPONSES TO INDIVIDUAL FACTS

  • Head responses are measured for paired prompts differing in one fact, with fact-specific response shares normalized by the summed energy across all three facts.The fixed-head comparison extends across facts, edit directions, and magnitude-matched random edits.Source: §D.2 · p. 19, Figure S2 · p. 20

D.3 INFLUENCE BEFORE AND AFTER THE EDIT

  • Influence compares later-site responses to controlled source-head shifts along factual directions, with profile distance measuring movement toward the alternative instruction’s unedited profile.The six-pair profiles are averaged before their squared distance is calculated.Source: §D.3 · p. 21

D.4 INDIVIDUAL HEAD CONTRIBUTIONS

  • 17.96 [17.82, 18.09] logits in Gemma and 48.23 [47.91, 48.54] logits in Qwen are paired shifts toward the alternative answer after the task edit.The comparison accounts for random changes at both the early intervention site and the restored head.Source: §D.4 · p. 21

E.1 RESPONSE MAPS AND FACTUAL INTERVENTIONS

  • For each prompt, unedited and task-edited response maps predict H3’s response to a separate factual intervention from the same measured H2 displacement.The factual-intervention response enters the prediction error, while the maps are built from controlled-perturbation measurements.Source: §E.1 · p. 22

E.2 PREDICTION ERRORS AND FIGURE NORMALIZATION

The task-specific response map predicts later-head responses more accurately than the averaged map in both original and fine-tuned models. Additional analyses test prediction across factual contexts and specify the training and resampling procedures.

E.2 PREDICTION ERRORS AND FIGURE NORMALIZATION

  • 53.6% [52.2%, 54.9%] in Gemma and 5.0% [4.0%, 5.9%] in Qwen are mean-error reductions for the task-specific map versus the averaged map.The improvement is concentrated in the two facts requested by the reading instructions; fact-wise summaries are descriptive.Source: §E.2 · p. 22

E.3 PREDICTION IN FINE-TUNED MODELS

  • 50.9% [49.6%, 52.4%] in fine-tuned Gemma and 2.2% [1.4%, 3.1%] in fine-tuned Qwen are percentage reductions versus the averaged map.Both corrected difference intervals exclude zero.Source: §E.3 · p. 23

F.1 FINE-TUNING AND CHECKPOINT SELECTION

  • The fine-tuning curriculum uses six rule–edge-pair combinations across eight graph states, two wordings, two response keys, and two layouts, totaling 12,288 prompts.Checkpoint selection uses different criteria and development data for Gemma and Qwen.Source: §F.1 · p. 23

F.2 FIXED DISPLACEMENTS ACROSS FACTUAL CONTEXTS

  • 0.01347 [0.00731, 0.01963] in Gemma and 0.02320 [0.01323, 0.03317] in Qwen are excess squared H3-response differences over the mean permutation and sign-flip control differences.The control-adjusted effects are reported for an identical H2 displacement across contexts where the targeted fact is or is not answer-relevant.Source: §F.2 · pp. 23–24

G.1 PAIRED AVERAGING AND RESAMPLING

  • Confidence intervals are based on paired resampling that preserves within-group conditions, with vocabulary groups receiving equal weight.Pointwise and simultaneous intervals use distinct procedures across the analyses.Source: §G.1 · p. 24

G.2 INFLUENCE, IMPORTANCE, AND FACTUAL INTERVENTIONS

  • Influence–importance comparisons use paired bootstrap intervals with a four-model Bonferroni correction, while factual-intervention analyses correct across 12 comparisons.The procedures resample vocabulary groups while retaining paired conditions.Source: §G.2 · p. 24

G.3 INTERNAL TASK EDITS AND HEAD RESTORATION

The appendices specify how uncertainty is estimated for internal-edit, prediction, and conditional-response analyses, including their paired resampling and multiple-comparison procedures.

G.3 INTERNAL TASK EDITS AND HEAD RESTORATION

  • Internal-edit analyses use simultaneous intervals across 24 comparisons, while response-share, depth-band, and margin analyses use their specified paired bootstrap procedures.The response-share analysis corrects across four comparisons; the depth-band criterion requires positive lower bounds in all three bands for raw and adjusted scores.Source: §G.3 · pp. 24–25

G.4 RESPONSE PREDICTION

  • Prediction analyses use paired bootstrap procedures for model error reductions and Student-t intervals with Bonferroni correction for fine-tuned model comparisons.The prediction procedures compare task-specific maps with averaged-map and no-response predictors.Source: §G.4 · p. 25

G.5 CONDITIONAL RESPONSES

  • Conditional-response tests use corrected Student-t intervals for raw, adjusted, and fact-specific comparisons, plus paired bootstrap intervals for response ratios.The fact-specific follow-up corrects across six estimates; the ratio analysis corrects across two model ratios.Source: §G.5 · p. 25
Loading 2610.13895v1…