Soft Head Selection for Injecting ICL-Derived Task Embeddings

1. Problemstellung

Große Sprachmodelle (LLMs) werden traditionell entweder durch Parameter-Effizientes Fine-Tuning (PEFT) oder In-Context Learning (ICL) an nachgelagerte Aufgaben angepasst.

• PEFT (z. B. LoRA) liefert starke Ergebnisse, erfordert jedoch das Training von zusätzlichen Parametern.
• ICL (Few-Shot) ist training-frei, führt aber zu langen Prompts und erhöhtem Inferenzaufwand.

Ein neuer Ansatz, die ICL-getriebene Embedding-basierte Anpassung, versucht, aufwendige Prompts zu vermeiden, indem er Aufgaben-spezifische Embeddings aus den Zwischenaktivierungen des Modells (basierend auf Few-Shot-Prompts) extrahiert und diese während der Inferenz injiziert.
Das Problem: Bisherige Methoden dieser Kategorie haben in der Praxis keine konsistenten Leistungssteigerungen gegenüber PEFT oder Few-Shot-ICL gezeigt. Ein Hauptgrund ist, dass sie oft nicht wissen, wo und wie diese Embeddings effektiv injiziert werden sollen. Die Annahme, dass Aufgabeninformationen gleichmäßig über alle Aufmerksamkeitsschichten verteilt sind, hat sich als falsch erwiesen; die Relevanz bestimmter „Attention Heads" variiert stark je nach Aufgabe.

2. Methodik: SITE (Soft Head Selection for ICL-derived Task Embeddings)

Die Autoren schlagen SITE vor, eine Methode, die Gradienten-basierte Optimierung nutzt, um die relevantesten Attention Heads für die Injektion von Aufgaben-Embeddings zu identifizieren. Der Prozess besteht aus drei Stufen:

• Stufe 1: Konstruktion von Task Embeddings
  Aus Few-Shot-Prompts (z. B. 50 Prompts mit je 10 Beispielen) werden die Aktivierungen des letzten Tokens für jeden Attention Head in jeder Schicht extrahiert. Diese werden gemittelt, um einen stabilen „Task Embedding"-Vektor pro Head zu erhalten.

• Stufe 2: Optimierung der Soft Head-Selection-Parameter
  Dies ist der Kern der Innovation. Statt festzulegen, welche Heads injiziert werden, führt SITE eine kontinuierliche Optimierung durch:

  • Es wird eine lernbare Matrix $A$ eingeführt, wobei jeder Eintrag $\alpha^{(l,h)}$ (zwischen 0 und 1) die Stärke der Injektion für einen spezifischen Head $h$ in Schicht $l$ steuert.
  • Die ursprüngliche Aktivierung $o$ wird linear mit dem Task Embedding $t$ interpoliert:
    $$o' \leftarrow (1 - \alpha) \cdot o + \alpha \cdot t$$
  • Das LLM bleibt eingefroren; nur die Parameter $\alpha$ werden über Gradientenabstieg (Cross-Entropy-Verlust) optimiert, um die Leistung bei Zero-Shot-Inferenz zu maximieren.
  • Die Parameter werden als Sigmoid-Funktion eines unbeschränkten Skalars parametrisiert, um den Wertebereich [0, 1] zu gewährleisten.

• Stufe 3: Zero-Shot Inferenz
  Während der eigentlichen Inferenz werden die gelernten $\alpha$-Werte und die Task Embeddings verwendet. Die Injektion erfolgt nur einmal am letzten Token des Eingabe-Prompts (unter Nutzung von KV-Caching), um die verbleibenden Tokens ohne weitere Eingriffe zu generieren.

3. Hauptbeiträge

1. SITE-Algorithmus: Eine neue, effiziente Methode zur Aufgabenanpassung, die Few-Shot-ICL signifikant übertrifft und die Leistung von PEFT (wie LoRA) erreicht, jedoch mit deutlich weniger trainierbaren Parametern.
2. Mechanistische Einblicke: Durch die Analyse der optimierten $\alpha$-Parameter zeigen die Autoren, dass Attention Heads aufgaben-spezifische Funktionen haben. Ähnliche Aufgaben teilen sich oft wichtige Heads, während unähnliche Aufgaben unterschiedliche Heads nutzen.
3. Skalierbarkeit und Generalisierung: Die Methode wurde auf 12 verschiedene LLMs (von 4B bis 70B Parametern, verschiedene Familien wie Llama, Mistral, Qwen, Gemma) erfolgreich getestet.

4. Ergebnisse

Die Evaluation umfasste 4 Benchmarks (FV, ANLI, MMLU-Pro, Big-Bench Hard) und 12 Modelle:

• Leistung: SITE übertrifft Few-Shot-ICL (10-Shot) auf dem FV-Benchmark um ca. 13,26 % und liegt im Durchschnitt über allen Benchmarks um 7,45 % höher.
• Vergleich mit PEFT: Auf Aufgaben wie FV und ANLI erreicht SITE bessere Ergebnisse als LoRA und (IA)³, bei vergleichbaren Ergebnissen auf komplexen Reasoning-Aufgaben (MMLU-Pro, BBH).
• Effizienz:
  • Parameter: SITE benötigt nur 1,02K trainierbare Parameter (im Vergleich zu ~3,4M bei LoRA).
  • Zeit: Die Trainingszeit (Optimierung der $\alpha$) ist vergleichbar mit PEFT, aber die Inferenzzeit ist nahezu identisch mit Zero-Shot-ICL, da keine langen Prompts benötigt werden.
• Aktivierungs-Patching-Analyse:
  • Intra-Task: Das Ersetzen der Aktivierungen der von SITE ausgewählten „High-$\alpha$" Heads mit den Task Embeddings verbessert die Zero-Shot-Leistung drastisch. Das Ersetzen von „Low-$\alpha$" Heads verschlechtert sie.
  • Inter-Task: Das Verwenden von Head-Selektionsparametern einer ähnlichen Aufgabe führt zu guten Ergebnissen, während Parameter einer unähnlichen Aufgabe die Leistung stark beeinträchtigen. Dies bestätigt die Aufgaben-Spezifität der Attention Heads.

5. Bedeutung und Fazit

SITE stellt einen wichtigen Schritt in der mechanistischen Interpretierbarkeit und effizienten Anpassung von LLMs dar.

• Paradigmenwechsel: Es zeigt, dass Aufgabenwissen nicht nur in den Gewichten (PEFT) oder im Prompt (ICL), sondern gezielt in spezifischen Teilen der Modellarchitektur (Attention Heads) kodiert werden kann.
• Praktische Relevanz: Die Methode bietet eine extrem effiziente Alternative zu PEFT, die besonders für Szenarien geeignet ist, in denen Speicherplatz oder Rechenzeit für das Fine-Tuning begrenzt sind, aber eine hohe Anpassungsfähigkeit benötigt wird.
• Zukunftsperspektive: Die Arbeit legt den Grundstein für ein besseres Verständnis der Rolle von Attention Heads bei spezifischen Aufgaben und eröffnet Wege für kontrollierbare Inferenz-Anpassungen ohne Zugriff auf die Modellgewichte (sofern die Aktivierungen zugänglich sind).

Einschränkungen: Die Methode erfordert einen kleinen Satz gelabelter Daten zur Konstruktion der Embeddings und den Zugriff auf interne Modellaktivierungen (was sie auf Open-Source-Modelle beschränkt). Zudem wurde sie noch nicht umfassend auf Multi-Step-Reasoning oder sehr lange Kontexte getestet.