💬 NLP

Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis

This paper introduces the Task Subspace Logit Attribution (TSLA) framework to mechanistically localize In-Context Learning into distinct Task Recognition and Task Learning components by identifying specialized attention heads that align hidden states with task subspaces and rotate them toward correct predictions, respectively.

Original authors: Haolin Yang, Hakaze Cho, Naoya Inoue

Published 2026-04-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haolin Yang, Hakaze Cho, Naoya Inoue

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a Large Language Model (LLM) like a massive, super-smart library that has read almost everything ever written. When you ask it a question with a few examples (like "Here is a sad movie review: 'The acting was terrible.' Label: Negative. Now, what about this review?"), it can figure out the pattern and answer correctly without you teaching it anything new. This is called In-Context Learning (ICL).

But how does the library actually do this? Does it just guess? Does it memorize?

This paper acts like a mechanic pulling back the hood of a Ferrari to see exactly which gears are turning. The authors discovered that the model doesn't just "think" in a blur; it splits the job into two very distinct steps, handled by two different teams of workers inside the model's brain.

The Two Teams: The "Detectives" and the "Translators"

The authors found that the model's internal "attention heads" (think of these as tiny, specialized workers) divide the work into two roles:

1. The Task Recognition (TR) Team: The "Detectives"

What they do: Their job is to look at the examples and say, "Ah! I see what game we are playing here. We are doing sentiment analysis. The possible answers are 'Positive' or 'Negative'."
The Analogy: Imagine you walk into a room and see people wearing red and blue hats. The Detective doesn't care who is wearing what yet; they just realize, "Okay, this is a Red vs. Blue game. The only valid answers are Red or Blue." They ignore the specific sentences and focus entirely on the options available.
Where they live: These workers are usually found in the deeper layers of the model (the later stages of processing). They are the ones who set the stage.

2. The Task Learning (TL) Team: The "Translators"

What they do: Once the Detective says, "It's a Red vs. Blue game," the Translator steps in. Their job is to look at the specific sentence you just wrote and say, "Based on the examples, this specific sentence means 'Red' (Positive)." They learn the specific link between the text and the label.
The Analogy: The Translator looks at the sentence "The movie was boring" and connects it to the "Negative" hat. They are doing the actual matching work.
Where they live: These workers are usually found in the earlier layers of the model. They do the heavy lifting of connecting the dots.

The "Induction Head" Mystery

For a long time, scientists were obsessed with a specific type of worker called the Induction Head. They thought these were the "magic" workers that made the whole thing work.

This paper solves a big mystery: The Induction Heads are actually just the Detectives (TR Team) in disguise.

They are great at spotting the pattern of "Label A, Label B, Label C."
They are not the ones doing the fine-tuned translation of the specific sentence.
If you remove the Induction Heads, the model forgets what game it's playing (Task Recognition fails), and it can't answer at all.

The "Geometry" of Thinking

The authors didn't just count heads; they looked at the "shape" of the model's thoughts (mathematically called "hidden states").

The Detective's Move: They push the model's thought into a specific "room" (a subspace) where only the valid answers live. If the thought is outside this room, the model is confused. The Detective pulls it inside.
The Translator's Move: Once the thought is inside the "room," the Translator spins it around until it points exactly at the correct answer (like a compass needle finding North).

Why This Matters: The "Zero-Shot" Problem

Have you ever asked an AI a question without examples, and it gave a generic, boring answer? That's because the Detectives are asleep. The model knows the words, but it hasn't realized, "Oh, this is a specific game with specific rules."

The paper shows that if you artificially "wake up" the Detective team (by injecting their patterns into the model), the model suddenly becomes much better at guessing the right answer, even without examples.

The Big Takeaway

In the past, we thought the model was a black box that just "learned" everything at once. This paper proves that:

Recognition comes first: The model must first realize what kind of task it is (Detective).
Learning comes second: Then it figures out the specific rules for this instance (Translator).
They are separate: You can break one without breaking the other. If you break the Detective, the model forgets the game. If you break the Translator, the model knows the game but can't pick the right answer.

It's like a restaurant: The Detective is the host who sees the menu and says, "We are serving Italian tonight." The Translator is the waiter who takes your order and connects "Spaghetti" to the kitchen. You need both to have a good meal, but they do very different jobs!

1. Problem Statement

Large Language Models (LLMs) exhibit In-Context Learning (ICL), the ability to solve new tasks from demonstrations without gradient updates. While ICL is powerful, its mechanistic underpinnings remain debated. Two dominant research paradigms exist but suffer from limitations:

Introspective (Component-Level): Focuses on specific attention heads (e.g., Induction Heads) and uses ablation to measure performance drops. However, it often fails to explain how these heads function or distinguish between different functional components of ICL.
Holistic (Functional): Decomposes ICL into Task Recognition (TR) (identifying the label space) and Task Learning (TL) (learning the text-to-label mapping) via input perturbations. However, it cannot trace these abstract functions back to concrete model components (specific attention heads).

The Gap: There is a lack of a unified framework that maps the functional decomposition of ICL (TR vs. TL) to specific, localized attention heads and explains their distinct geometric roles in the model's hidden state space.

2. Methodology

The authors propose a novel framework based on Task Subspace Logit Attribution (TSLA) to identify and analyze attention heads specialized for TR and TL.

A. Task Subspace Logit Attribution (TSLA)

Unlike previous methods that rely on direct logit attribution to specific tokens (which can be sensitive to label surface forms and ignore competition), TSLA uses geometric projections:

Task Recognition (TR) Score: Measures the projection of a head's output onto the task subspace spanned by the unembedding vectors of all demonstration labels ( $Y$ $Y$ ).
- Formula: $\| \text{Proj}_{W_U^Y} a_{N,k}^l \|^2$ .
- Intuition: TR heads concentrate their output within the subspace of potential task labels, effectively "activating" the label space.
Task Learning (TL) Score: Measures the alignment of the head's output (projected onto the task subspace) with the average direction between the correct label ( $y^*$ $y^{*}$ ) and incorrect labels ( $Y \setminus \{y^*\}$ $Y ∖ {y^{*}}$ ).
- Formula: $\frac{\mathbb{E}_{y' \in Y \setminus \{y^*\}} [a_{N,k}^{l,\top} (W_U^{y^*} - W_U^{y'})]}{\| \text{Proj}_{W_U^Y} a_{N,k}^l \|^2}$ .
- Intuition: TL heads rotate the hidden state within the task subspace toward the correct label, distinguishing it from competitors.

B. Experimental Validation

The authors validate the framework using:

Correlation & Overlap Analysis: Comparing TSLA-identified heads with Induction Heads (IHs) and across datasets.
Fine-Grained Ablation: Selectively removing top TR or TL heads to observe effects on accuracy and the Task Recognition Ratio (proportion of predictions within the correct label set).
Input Perturbations: Shuffling demonstration text (disrupting TL) or relabeling demonstrations (disrupting TR) to test causal independence.
Steering & Geometric Analysis: Injecting head outputs as "Task Vectors" into zero-shot prompts and analyzing hidden state evolution (Subspace Alignment vs. Logit Difference).

3. Key Contributions

Novel Identification Method (TSLA): A theoretically grounded method to localize TR and TL heads that overcomes the limitations of direct logit attribution by utilizing task-subspace geometry.
Mechanistic Decomposition: Demonstrates that TR and TL are realized by distinct, separable sets of attention heads.
- TR Heads: Align hidden states with the task subspace (enabling label recognition).
- TL Heads: Rotate hidden states within the subspace toward the correct label (enabling prediction).
Reconciliation of Induction Heads: Shows that Induction Heads (IHs) are primarily a subset of TR heads. Their role is to recognize the label space rather than perform the specific text-to-label mapping, explaining why ablating IHs causes task recognition failure.
Geometric Interpretation: Provides a unified geometric account where ICL is a two-step process: (1) TR heads project the state into the task subspace, and (2) TL heads perform a rotation within that subspace to select the correct label.

4. Key Results

Separability: Ablating TR heads collapses the Task Recognition Ratio (predictions fall outside the label set) and accuracy. Ablating TL heads reduces accuracy but leaves the TR ratio largely intact (predictions remain within the label set but are random).
Induction Heads as TR: Top Induction Heads show high overlap and correlation with TR heads, but low overlap with TL heads. IHs are crucial for recognizing the label space, not for learning the specific mapping.
Layer Distribution: TR heads are predominantly located in deeper layers of the Transformer, while TL heads and IHs appear in earlier/middle layers. This suggests a processing order where task learning (mapping) begins early, but final task recognition (subspace alignment) is solidified in deeper layers.
Attention Patterns:
- TR Heads: Attend heavily to demonstration labels (e.g., "positive", "negative") to aggregate label space information.
- TL Heads: Attend heavily to query tokens to extract semantic meaning for mapping.
Task Vector Steering:
- In classification (fixed label space), injecting TR-based task vectors restores zero-shot performance by fixing task recognition.
- In open-ended generation (infinite label space), TL-based task vectors are more effective, as they capture the text-to-output mapping rather than a fixed label set.
Cross-Task Generalization: TR heads are highly consistent across different datasets (task-invariant), whereas TL heads vary significantly depending on the specific task mapping.

5. Significance

This paper provides a unified and interpretable account of ICL by bridging the gap between component-level circuit analysis and holistic functional decomposition.

Theoretical Impact: It clarifies the mechanistic role of Induction Heads, resolving conflicting literature by positioning them as TR drivers.
Practical Implications: The findings suggest that improving zero-shot performance in classification tasks requires enhancing Task Recognition (activating TR heads), while improving generation tasks requires better Task Learning (activating TL heads).
Methodological Advance: The TSLA framework offers a robust tool for dissecting LLM behaviors, moving beyond simple ablation to understanding the geometric transformations (alignment vs. rotation) that drive model predictions.

In summary, the authors successfully localize the "what" (TR vs. TL) and the "how" (subspace alignment vs. rotation) of In-Context Learning, revealing that LLMs execute ICL through a coordinated, geometrically distinct interplay of specialized attention heads.