TSLA: Task Subspace Logit Attribution

• The paper introduces TSLA to decompose attention heads into Task Recognition and Task Learning by projecting head outputs onto a demonstration-derived subspace.
• It derives two scalar scores per head—a TR score for subspace alignment and a TL score for discriminative rotation toward the correct label—providing actionable insights for dissecting in-context learning.
• Empirical validation through ablation studies, correlation with induction heads, and geometric steering confirms TSLA’s capability to pinpoint causal roles in transformer models.

Searching arXiv for the specified paper to ground the article in the current literature. Task Subspace Logit Attribution (TSLA) is a principled framework for attributing the contributions of individual attention heads to the two core components of in-context learning (ICL)—Task Recognition (TR) and Task Learning (TL)—by asking: “How much does each head move the model’s hidden state into the subspace spanned by the task’s label embeddings, and in which direction within that subspace?” By projecting each head’s output onto a “task subspace” defined by the unembedding vectors of the demonstration labels, TSLA yields two scalar scores per head: a TR score measuring subspace alignment, and a TL score measuring discriminative rotation toward the correct label. In "Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis" (Yang et al., 29 Sep 2025), TSLA is presented as a way to reconcile the component-level analysis of attention heads with the holistic decomposition of ICL into TR and TL, providing a unified and interpretable account of how LLMs execute ICL across diverse tasks and settings.

1. Conceptual role in in-context learning analysis

TSLA is introduced in the context of two dominant perspectives on ICL. One is the component-level analysis of attention heads; the other is the holistic decomposition of ICL into Task Recognition and Task Learning. The framework is designed to localize these two functions at the level of individual attention heads and to demonstrate their distinct yet complementary roles.

Within this framing, TR and TL are treated as separable components. TR concerns whether the model recognizes the relevant label space defined by the demonstrations. TL concerns whether, once restricted to that label space, the model discriminates the correct label from its competitors. TSLA operationalizes this distinction by assigning each attention head a score for subspace alignment and a score for discriminative directionality.

A central implication of this construction is that attention-head-level analysis and functional decomposition need not be competing descriptions. TSLA is explicitly presented as a framework that reconciles them. This suggests that the paper’s contribution is not only a new attribution method, but also a reorganization of prior mechanistic observations into a common geometric language.

2. Mathematical definition

TSLA begins from a circuit decomposition of the final-token hidden state. Let $x_1,\ldots,x_N$ be the full ICL prompt, consisting of $m$ demonstrations plus query, and let $h_N^L$ be its final hidden state at layer $L$. The residual-stream updates are written as

$$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$

where each head output $a_{N,k}^l \in \mathbb{R}^d$ contributes additively to logits via $W_U^\top a_{N,k}^l$, with $W_U \in \mathbb{R}^{d \times |V|}$ the unembedding matrix.

The task subspace is defined from the demonstration labels. Suppose the demonstration labels come from a label set $\mathcal{Y}=\{y_1,\ldots,y_m\}$. Denote by $W_U^\mathcal{Y}\in\mathbb{R}^{d\times r}$ the submatrix of $m$0 whose columns are the unembeddings of these labels, where $m$1. The projector onto that span is

$m$2

For a head $m$3 at the final token, let $m$4. Its Task-Recognition score is the norm of its projection into the task subspace:

$m$5

This measures how much that head’s logit contribution lies in directions that boost or suppress any of the demonstration labels.

The Task-Learning score asks whether the head preferentially pushes toward the correct label $m$6 rather than its competitors. Writing $m$7 and $m$8 for each incorrect label unembedding, the definition is

$m$9

Because $h_N^L$0 lives in the task subspace, this ratio is the cosine of the angle between the projected head output and the mean discriminant direction toward $h_N^L$1. It lies in $h_N^L$2, where large positive values indicate that the head strongly amplifies the correct label relative to its competitors.

3. Algorithmic procedure and head selection

The algorithmic procedure for TSLA is specified for a pretrained Transformer and a batch of ICL prompts.

1. Forward-run each prompt and collect the final-token head outputs $h_N^L$3 for every head $h_N^L$4.
2. Assemble $h_N^L$5 from the demonstration labels’ unembedding vectors and compute $h_N^L$6.
3. For each head output $h_N^L$7:
   • Compute $h_N^L$8 and set $h_N^L$9.
   • Compute the numerator $L$0 and set $L$1.
4. Sum both scores across a held-out set of prompts to smooth out noise.
5. Rank heads by $L$2 to select top TR heads, and by $L$3 to select top TL heads.

This procedure turns the TR/TL decomposition into a ranking problem over heads. The resulting head sets are then used for correlation analysis, ablation studies, and steering interventions. A plausible implication is that TSLA is intended not merely as a descriptive score, but as a selection mechanism for causal tests of hypothesized function.

4. Empirical localization of TR and TL heads

The empirical validation of TSLA proceeds through three forms of evidence: correlation with induction heads, ablation studies, and input perturbations (Yang et al., 29 Sep 2025).

For correlation with induction heads, the paper compares the top $L$4 of TR heads, TL heads, and induction heads, where induction heads are ranked by their IH-pattern score. TR heads and induction heads have very high Jaccard overlap, much greater than random, and strong Spearman and Kendall correlations in their rankings. TL heads barely overlap with induction heads or TR heads. The stated interpretation is that induction heads are a subset of TR heads mechanistically, primarily responsible for recognizing the label space rather than discriminating the correct label.

The ablation studies zero out the top $L$5 of TR heads and TL heads separately and evaluate two metrics on the same ICL prompts: Accuracy, defined as overall fraction of correct query predictions, and Task-Recognition ratio, defined as fraction of predictions that fall within the demonstration label set. Ablating TR heads collapses the TR ratio, close to random labels outside $L$6, and thus accuracy. Ablating TL heads leaves TR ratio near $L$7 but drives accuracy toward random guessing over $L$8. Ablating induction heads mimics TR ablation.

Input perturbations further separate the two functions. When demonstration texts are shuffled to destroy the text-to-label mapping, thereby disrupting TL, ablating TL heads has no effect because the mapping is already broken, but TR ablation still kills performance. When demonstrations are relabeled with arbitrary tokens to break label semantics, thereby disrupting TR, ablating TR heads no longer matters because the TR block is already destroyed, but TL ablation still lowers accuracy. These experiments are presented as confirming that TR and TL heads implement separable functions.

5. Geometric interpretation through steering

The steering experiments provide a geometric interpretation of the roles assigned by TSLA (Yang et al., 29 Sep 2025). The summed outputs of top TR heads, denoted $L$9, and top TL heads, denoted $$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$0, are extracted from $$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$1-shot prompts and injected at an intermediate layer, for example layer $$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$2, of zero-shot inputs.

Two geometry metrics are then measured on the hidden state $$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$3. The first is subspace alignment,

$$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$4

and the second is logit difference,

$$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$5

The reported results distinguish the two head classes sharply. Adding $$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$6 substantially increases subspace alignment, aligning $$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$7 with the task subspace, but has small effect on logit difference. Adding $$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$8 rotates $$h_N^L = h_N^0 + \sum_{l=1}^L \sum_{k=1}^K a_{N,k}^l + \sum_{l=1}^L m_N^l,$$9 within the task subspace toward the correct label, boosting logit difference without noticeably changing overall subspace alignment.

Layer-wise tracking of these metrics correlates strongly with the per-layer contributions of TR and TL heads. The interpretation given is that TR heads are the causal drivers of alignment to the task subspace, while TL heads are the causal drivers of rotation within that subspace toward the correct label. This geometric account makes the distinction between “recognizing the label set” and “selecting the correct label” operational at the hidden-state level.

6. Relation to induction heads, task vectors, and broader interpretation

TSLA is explicitly used to unify prior mechanistic concepts (Yang et al., 29 Sep 2025). For induction heads, the framework states that induction heads almost exactly coincide with TR heads and that their primary function is to push hidden states into the label subspace. In this view, earlier toy-model accounts of induction heads as simple copy operations are subsumed by a subspace-alignment interpretation that extends to arbitrary label semantics.

For task vectors, the paper connects TSLA to findings that averaging the outputs of certain heads across demonstrations produces task vectors that, when injected at zero-shot time, recover ICL performance. TSLA is presented as specifying which heads to pick, distinguishing TR from TL, and as supplying a geometric explanation: TR-based task vectors restore subspace alignment, which is described as critical for classification, while TL-based vectors provide the fine-grained rotation needed for generation tasks with open label spaces.

A common point of tension in mechanistic interpretability is whether ICL should be understood primarily through discrete circuits such as induction heads or through more global functional decompositions such as TR and TL. TSLA is presented as a reconciliation of these views rather than a replacement of one by the other. This suggests that the main significance of the framework lies in showing how head-level specialization can instantiate a higher-level decomposition of ICL.

In summary, TSLA provides a clear mathematical decomposition of each head’s logit attribution into task-recognition and task-learning components; a simple algorithm to rank and select TR and TL heads via projections onto the demonstration-label subspace; empirical validation through correlation with induction heads, targeted ablative interventions, and steering interventions; and a unified geometric narrative in which TR heads align hidden states to the task subspace and TL heads rotate them within that subspace toward the correct label.