Soft Head Selection for Injecting ICL-Derived Task Embeddings

1. Problem Statement

Large Language Models (LLMs) are typically adapted to downstream tasks via Parameter-Efficient Fine-Tuning (PEFT) (e.g., LoRA) or In-Context Learning (ICL).

• PEFT offers strong performance but requires training and introduces new parameters.
• ICL is training-free but incurs high inference overhead due to long prompts and lacks consistency in performance.
• ICL-driven Embedding-based Adaptation is a newer paradigm that extracts task-specific vectors (embeddings) from intermediate activations of few-shot ICL prompts and injects them into the model during zero-shot inference.

The Core Challenge: Existing embedding-based methods (e.g., Function Vectors, Task Vectors) have failed to consistently outperform standard few-shot ICL or PEFT. They often rely on heuristic search strategies or layer-wise patching that are inefficient and fail to account for the fact that task-relevant information is unevenly distributed across specific attention heads. Previous work largely treated attention heads as a monolithic block or searched for injection points without a principled, differentiable selection mechanism.

2. Methodology: SITE

The authors propose SITE (Soft head-selection for ICL-derived Task Embeddings), a gradient-based method that identifies and injects task embeddings into the most relevant attention heads. The method operates in three stages:

Stage 1: Task Embedding Construction

• The model processes $M$ few-shot ICL prompts (each containing $N$ input-output pairs).
• For every attention head $h$ in every layer $l$, the output activation corresponding to the last token is extracted.
• These activations are averaged across the $M$ prompts to form a task-specific embedding vector $t^{(l,h)}$ for each head.

Stage 2: Soft Head-Selection Optimization

• Instead of hard selection or heuristic search, SITE formulates head selection as a continuous optimization problem.
• A learnable matrix $A$ is introduced, where each entry $\alpha^{(l,h)} \in [0, 1]$ represents the "importance" or injection weight for a specific attention head.
• Injection Mechanism: During optimization (using zero-shot prompts), the original head activation $o^{(l,h)}$ is linearly interpolated with the task embedding $t^{(l,h)}$:
  $$o^{(l,h)} \leftarrow (1 - \alpha^{(l,h)}) \cdot o^{(l,h)} + \alpha^{(l,h)} \cdot t^{(l,h)}$$
• The matrix $A$ is optimized via gradient descent (minimizing next-token prediction loss) while the backbone LLM remains frozen. The parameters $\alpha$ are initialized at 0.5 and constrained via a sigmoid function.

Stage 3: Zero-Shot Inference

• Once optimized, the learned parameters $A$ and task embeddings $\{t^{(l,h)}\}$ are fixed.
• During inference, the injection is applied only once at the last token of the initial input prompt (assuming KV caching).
• This single intervention embeds task-relevant information into the Key-Value (KV) cache, allowing the model to generate the rest of the sequence without further intervention or in-prompt examples.

3. Key Contributions

1. Novel Optimization Framework: Introduced a gradient-based "soft head-selection" mechanism that learns which attention heads to inject task embeddings into, rather than relying on fixed heuristics or expensive search.
2. Superior Performance: Demonstrated that SITE significantly outperforms prior embedding-based methods and few-shot ICL across diverse tasks (generation, reasoning, NLU) while using substantially fewer trainable parameters than PEFT (e.g., ~1K parameters vs. ~3.4M for LoRA).
3. Mechanistic Insights: Through activation patching analysis, the authors revealed:
   • Intra-task: The selected heads effectively capture task-relevant information; patching them boosts performance, while patching unselected heads degrades it.
   • Inter-task: Attention head functionality is highly task-specific. Similar tasks (e.g., English-French and English-German translation) share important heads, while dissimilar tasks do not. This extends mechanistic interpretability from a task-agnostic to a task-specific perspective.
4. Scalability and Generality: Validated on 12 LLMs (4 families, 4B to 70B parameters), showing consistent improvements over 10-shot ICL (average +10-14% accuracy).

4. Experimental Results

• Benchmarks: Evaluated on 57 tasks from the Function Vectors (FV) benchmark, plus ANLI, MMLU-Pro, and Big-Bench Hard (BBH).
• Performance:
  • On the FV benchmark, SITE (with $M=50$) achieved 90.02% accuracy, outperforming 10-shot ICL (76.8%) and approaching PEFT-level performance (LoRA: 86.76%).
  • On reasoning-heavy benchmarks (MMLU-Pro, BBH), SITE performed comparably to PEFT, though slightly lower, suggesting ICL-derived embeddings may have limits in complex reasoning without further adaptation.
• Efficiency:
  • Parameters: SITE requires only 1.02K trainable parameters (the selection matrix $A$) compared to LoRA's 3.4M.
  • Runtime: Inference time is nearly identical to zero-shot baselines and significantly faster than 10-shot ICL or methods requiring multiple injection points.
• Robustness: Performance remained stable across different prompt templates and shot counts used for embedding construction.

5. Significance

• Bridging the Gap: SITE successfully bridges the gap between the training-free nature of ICL and the high performance of PEFT, offering a "best of both worlds" solution with minimal computational cost.
• Mechanistic Understanding: The work provides deep insights into the internal workings of LLMs, proving that task adaptation is not uniform but relies on specific, task-dependent attention heads. This challenges the notion of generic "task vectors" and suggests future adaptation methods should be head-aware.
• Practical Deployment: By requiring only a small amount of labeled data for pre-computation and freezing the backbone model, SITE is highly suitable for scenarios where full fine-tuning is impossible (e.g., proprietary models with API access, though the paper notes current limitations to open-source models requiring internal activation access) or where inference latency and memory are critical constraints.

In conclusion, SITE represents a significant advancement in activation steering, demonstrating that precise, gradient-guided selection of attention heads allows for highly efficient and effective task adaptation without the overhead of traditional fine-tuning.