💬 NLP

Task Vectors, Learned Not Extracted: Performance Gains and Mechanistic Insight

This paper introduces Learned Task Vectors (LTVs) as a superior, directly trained alternative to extracted task vectors for in-context learning, while providing mechanistic insights into how these vectors steer predictions through specific attention circuits and linear propagation patterns within Transformer models.

Original authors: Haolin Yang, Hakaze Cho, Kaize Ding, Naoya Inoue

Published 2026-04-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haolin Yang, Hakaze Cho, Kaize Ding, 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 you have a super-smart robot (a Large Language Model, or LLM) that knows a lot of facts but doesn't know how to do a specific new job, like writing a poem or solving a math problem.

Usually, to teach it, you'd have to retrain the whole robot, which is slow and expensive. But this robot has a magic trick called In-Context Learning (ICL). If you show it a few examples right before asking the question (like showing it three examples of "Happy" vs. "Sad" movie reviews), it suddenly understands the task and gets it right.

The Big Question: How does the robot do this? Where does it "store" the lesson it just learned from your examples?

The Old Way: Finding the "Ghost" in the Machine

Previous researchers thought the robot was compressing those examples into a tiny, invisible "instruction manual" hidden inside its brain. They called this a Task Vector (TV).

However, finding this manual was like trying to catch a ghost. They had to dig through the robot's messy internal data, use complex math to filter out the noise, and hope they found the right "instruction." It was slow, unreliable, and nobody really understood how the instruction actually worked once found.

The New Way: "Learned" Task Vectors (LTVs)

This paper introduces a better approach: Learned Task Vectors (LTVs).

Instead of trying to fish the instruction out of the robot's messy brain, the authors say: "Let's just write the instruction from scratch."

They treat the Task Vector like a tiny, custom-made key. They use a computer to "train" this key directly, testing millions of versions until they find the perfect one that unlocks the robot's ability to do the task.

The Analogy: Imagine the old method was trying to copy a secret recipe by tasting the final dish and guessing the ingredients. The new method is just writing the recipe down perfectly on a piece of paper and handing it to the chef.
The Result: These "Learned" keys work better, faster, and can be used in different parts of the robot's brain, not just one specific spot.

How Does the "Key" Actually Work? (The Mechanism)

The authors didn't just stop at making a better key; they opened the robot up to see how the key turns the gears. They found two fascinating things:

1. The "Specialized Gear" (Low-Level Insight)

The robot's brain is made of billions of tiny gears (called attention heads). The authors discovered that the Task Vector doesn't talk to all the gears. It only talks to a very small, specific group of "Key Gears."

The Analogy: Imagine a massive orchestra. When the conductor (the Task Vector) wants the song to change, they don't shout at every musician. They only signal a few specific violinists and drummers. Those few musicians then change the sound of the whole orchestra.
The Discovery: If you remove these "Key Gears," the Task Vector stops working. If you remove random gears, nothing happens.

2. The "Rotating and Stretching" Dance (High-Level Insight)

As the Task Vector moves through the robot's layers (from the front of the brain to the back), it changes shape. The authors found this process is surprisingly simple and linear.

Early Layers (The Rotation): When the vector enters the robot early on, the brain rotates it.
- Analogy: Imagine you have a compass needle pointing North. The robot spins it so it points East. It's re-orienting the instruction to face the right direction for the specific task.
Late Layers (The Stretching): By the time the vector reaches the end of the brain, it stops rotating. Instead, the brain just stretches it (makes it bigger or smaller).
- Analogy: Now that the compass is pointing East, the robot just turns up the volume on that direction. It amplifies the signal to make the final decision clear.

Why This Matters

This paper is a big deal for two reasons:

Practical Power: We now have a way to create "super-instructions" that make AI models smarter and more flexible without needing to retrain the whole model. It's like giving a general-purpose robot a specific, high-quality manual for any job you need.
Understanding the "Black Box": For years, AI has been a "black box"—we know it works, but we don't know why. This paper pulls back the curtain. It shows us that the robot's "learning" is actually a very structured dance of rotating ideas into the right direction and amplifying them at the end.

In short: The authors stopped trying to guess how the robot learns by looking at its messy notes. Instead, they wrote the notes themselves, found the perfect ones, and discovered that the robot's brain works like a simple, elegant machine that just needs the right key to turn the right gears.

1. Problem Statement

Large Language Models (LLMs) exhibit In-Context Learning (ICL), where they perform new tasks based on demonstrations provided in the input prompt without weight updates. Recent research posits that LLMs compress these demonstrations into Task Vectors (TVs)—compact vector representations injected into hidden states to steer predictions.

However, existing approaches suffer from two critical limitations:

Inefficient Extraction: Current methods (e.g., Vanilla TV, Function Vectors) rely on extracting vectors from model outputs or hidden states using opaque, cumbersome optimization or filtering. These methods are highly sensitive to the quality of the model's internal representations and often yield suboptimal vectors.
Lack of Mechanistic Understanding: Prior work demonstrates that injected TVs improve performance but fails to explain how. The specific low-level components (e.g., attention heads vs. MLPs) and high-level propagation dynamics (e.g., linear vs. nonlinear transformations) through which TVs influence the model remain a "black box."

2. Methodology

A. Learned Task Vectors (LTVs)

Instead of extracting vectors from ICL hidden states, the authors propose directly training Task Vectors.

Training Objective: A vector $\theta$ is added to specific hidden states (at layers $L$ and token positions $P$ ) of a zero-shot query. The vector is optimized via gradient descent to minimize the negative log-likelihood of the correct label:
$\min_{\theta} -\log p(y_q | x_q, \theta, L, P)$
Flexibility: Unlike baselines that inject only at the last token of a specific layer, LTVs can be trained and injected at arbitrary layers and token positions, including within ICL prompts themselves.
Efficiency: This approach finds the optimal TV unconstrained by the model's pre-existing representation quality.

B. Mechanistic Analysis Framework

The authors analyze the mechanism of TV effectiveness through two lenses:

Low-Level Interactions: Investigating which model components interact with the injected TV. They focus on Attention Head OV Circuits (Value embedding $W_V$ $W_{V}$ and Output projection $W_O$ $W_{O}$ ).
- Reconstruction Experiment: They reconstruct the TV effect by summing the transformed TV ( $W_O^\top W_V \theta$ ) across all attention heads in subsequent layers and injecting this aggregate back into the residual stream.
- Ablation: They identify "Key Heads" based on saliency scores and ablate them to measure performance degradation.
High-Level Propagation: Analyzing how the TV evolves as it passes through subsequent layers.
- They hypothesize that the composite layer update acts linearly on the TV.
- They decompose the linear transformation matrix $W_{TV}$ into Rotation ( $Q$ ) and Stretch ( $\Sigma$ ) components (Polar Decomposition) to understand how early vs. late injected TVs are processed.

3. Key Contributions

Superior Performance via Learning: The paper introduces LTVs, which consistently outperform extracted baselines (Vanilla TV and Function Vectors) across various models (Llama, Qwen, Yi) and tasks. LTVs achieve accuracy matching or exceeding few-shot ICL levels, even when injected at late layers where baselines fail.
Flexibility and Scalability: LTVs are not confined to the last token or a single layer. They remain effective when injected at multiple positions, multiple layers, or even within ICL prompts.
Low-Level Mechanism (OV Circuits): The study identifies that Attention Head OV circuits are the primary conduit for TV effects. Specifically, a small subset of "Key Heads" (those with high saliency) drives the performance. These heads exhibit less "attention sink" behavior and focus more on final token positions compared to random heads.
High-Level Mechanism (Rotation-Stretch Dynamics):
- Linearity: Despite Transformer nonlinearities, TV propagation is largely linear.
- Early Layers: Injected vectors undergo significant rotation to align with task-relevant subspaces (label unembeddings).
- Late Layers: The rotation effect diminishes, and the mechanism shifts to scaling (stretching) the magnitude of the vector to boost logits.
Label Space Dependence: The paper demonstrates that cross-task generalization of TVs is fundamentally constrained by the overlap in task label spaces (e.g., an SNLI TV works on RTE because both use {true, false}, but not on Capital tasks).

4. Experimental Results

Performance: On Llama3.1-8B, LTVs injected at the middle layer achieved 83.49% accuracy on a complex generation task (Myopic dataset), significantly outperforming Vanilla TV (37.80%) and Function Vectors (37.30%). LTVs also surpassed ICL performance in several settings.
PEFT Potential: With only $d$ trainable parameters (where $d$ is the hidden dimension), LTVs outperformed Prefix Tuning and LoRA in SST-2 classification while maintaining lower inference FLOPs and memory costs.
Mechanistic Validation:
- OV Reconstruction: Reconstructing the TV effect solely through OV circuits recovered ~83% of the original LTV's performance, confirming their dominance over MLPs.
- Linear Hypothesis: A linearly reconstructed TV (fitted via regression on perturbed inputs) matched the original TV's performance across most layers, validating the linear propagation hypothesis.
- Rotation vs. Stretch: Analysis showed that early-layer injections require strong rotation to align with task directions, while late-layer injections are already aligned and primarily require magnitude scaling.

5. Significance

Practical Tool: LTVs provide a robust, parameter-efficient method for steering LLMs, offering a viable alternative to traditional fine-tuning or prompt engineering.
Theoretical Insight: The paper moves beyond "black-box" injection to provide a principled mechanistic explanation of ICL. It reveals that ICL relies on a linear geometric transformation (rotation followed by scaling) mediated by specific attention head circuits.
Interpretability: By identifying "Key Heads" and the linear nature of TV propagation, the work demystifies how LLMs utilize context, offering new avenues for model steering, safety alignment, and understanding emergent capabilities.

In summary, this work shifts the paradigm from extracting task vectors from imperfect model states to learning optimal vectors, while simultaneously decoding the geometric and circuit-level mechanisms that allow these vectors to steer large language models.