Soft Head Selection for Injecting ICL-Derived Task Embeddings

Probleemstelling

Grote Taalmodellen (LLMs) worden traditioneel aangepast aan downstream-taken via twee methoden:

1. Parameter-Efficient Fine-Tuning (PEFT): Bijv. LoRA of Adapters. Dit levert sterke prestaties op maar vereist training en het opslaan van extra parameters.
2. In-Context Learning (ICL): Het toevoegen van voorbeelden aan de prompt. Dit is training-vrij maar leidt tot lange prompts en hoge inferentie- overhead.

Recente studies hebben een derde paradigma geïntroduceerd: ICL-gedreven embedding-gebaseerde adaptatie. Hierbij worden taakspecifieke embeddings afgeleid van de tussenliggende activaties van het model (geëliciteerd door few-shot prompts) en tijdens de inferentie geïnjecteerd. Hoewel dit conceptueel aantrekkelijk is, hebben bestaande methoden niet consequent betere prestaties laten zien dan PEFT of standaard ICL. Een groot probleem is dat deze methoden vaak heuristieken gebruiken om te bepalen waar (in welke lagen of koppen) embeddings moeten worden geïnjecteerd, zonder rekening te houden met de specifieke functionaliteit van attention heads per taak.

Methodologie: SITE

De auteurs stellen SITE (Soft head-selection for ICL-derived Task Embeddings) voor, een methode die gebruikmaakt van een continue, op gradiënten gebaseerde optimalisatie om de meest relevante attention heads te identificeren voor injectie.

Het proces bestaat uit drie fasen:

1. Constructie van Task Embeddings:

   • Er worden $M$ few-shot ICL-prompts gebruikt (elk met $N$ voorbeelden).
   • Voor elke prompt wordt een forward pass uitgevoerd en de output van elke attention head in elke laag wordt geëxtraheerd.
   • De activaties van het laatste token worden gemiddeld over alle prompts om een taakembedding ($t^{(l,h)}$) te vormen voor elke head $h$ in laag $l$.

2. Optimalisatie van Soft Head-Selectie:

   • In plaats van heuristieken, formuleren de auteurs de selectie van heads als een continu optimalisatieprobleem.
   • Er wordt een leerbaar matrix $A$ geïntroduceerd met parameters $\alpha^{(l,h)} \in [0, 1]$. Deze parameter bepaalt de mate waarin de taakembedding wordt geïnjecteerd in de oorspronkelijke activatie van een specifieke head.
   • De injectie gebeurt via lineaire interpolatie:
     $$o^{(l,h)} \leftarrow (1 - \alpha^{(l,h)}) \cdot o^{(l,h)} + \alpha^{(l,h)} \cdot t^{(l,h)}$$
   • Het model (LLM) blijft bevroren; alleen de parameters in $A$ worden geoptimaliseerd via gradient descent (minimiseren van cross-entropy loss) op een zero-shot prompt, waarbij de embeddings worden geïnjecteerd volgens de huidige waarden van $A$.

3. Zero-Shot Inferentie:

   • Na optimalisatie worden de geleerde parameters $A$ en de taakembeddings gebruikt om het model te laten presteren op nieuwe queries zonder voorbeelden in de prompt.
   • De injectie gebeurt slechts één keer (bij het laatste token van de input prompt), wat de complexiteit verlaagt en de efficiëntie ten opzichte van KV-caching maximaliseert.

Belangrijkste Bijdragen

• SITE Framework: Een nieuwe methode die ICL-gedreven embeddings combineert met een gradient-based soft head-selection mechanisme. Dit elimineert de noodzaak voor handmatige zoekopdrachten naar injectiepunten.
• Superieure Prestaties: SITE presteert significant beter dan eerdere embedding-gebaseerde methoden en few-shot ICL, en benadert de prestaties van PEFT (zoals LoRA) met aanzienlijk minder trainbare parameters.
• Mechanistische Inzichten: Door middel van "activation patching" analyses tonen de auteurs aan dat attention heads taakspecifiek zijn.
  • Intra-task: De door SITE geselecteerde heads (hoge $\alpha$) bevatten cruciale taakinformatie; het patchen van deze heads verbetert prestaties drastisch.
  • Inter-task: Vergelijkbare taken delen belangrijke attention heads, terwijl verschillende taken dat niet doen. Dit bevestigt dat de functionaliteit van heads niet universeel is, maar afhankelijk van de taak.

Resultaten

De methode werd geëvalueerd op 12 verschillende LLMs (variërend van 4B tot 70B parameters, inclusief Llama-3.1, Mistral, Qwen en Gemma families) over 57 taken in de Function Vectors (FV) benchmark, evenals ANLI, MMLU-Pro en Big-Bench Hard.

• Prestaties: SITE overtreft few-shot ICL (10-shot) met gemiddeld 13.26% op de FV benchmark en presteert vergelijkbaar met of beter dan LoRA op de meeste taken, terwijl het slechts 1.02K trainbare parameters vereist (tegenover ~3.4M voor LoRA).
• Efficiëntie: Omdat de embeddings en selectieparameters slechts één keer per taak worden berekend (pre-computation), schaalt de inferentietijd bijna even goed als zero-shot inferentie en is aanzienlijk sneller dan 10-shot ICL of methoden die veelvuldige patching vereisen.
• Generalisatie: De methode werkt consistent goed over verschillende modelgroottes en families, wat aantoont dat het principe van soft head-selectie robuust is.

Betekenis en Impact

Dit paper is significant omdat het de kloof overbrugt tussen de efficiëntie van ICL en de prestaties van PEFT. Het introduceert een nieuw perspectief op mechanistische interpretatie: attention heads zijn niet statisch functioneel, maar hun rol is sterk afhankelijk van de specifieke taak.

Door te leren welke heads relevant zijn voor een taak in plaats van ze handmatig te selecteren, biedt SITE een schaalbare, training-vrije (voor het hoofdmodel) en parameter-efficiënte manier om LLMs aan te passen. Dit opent de deur voor snellere, goedkopere en meer controleerbare aanpassingen van grote modellen in real-world applicaties, zonder de noodzaak van zware fine-tuning of lange prompts. De bevindingen over taakafhankelijkheid van attention heads bieden ook waardevolle inzichten voor toekomstig onderzoek in mechanistische interpretatie en modelcontrole.