DeCoVec: Building Decoding Space based Task Vector for Large Language Models via In-Context Learning

Probleemstelling

Task vectors (taakvectoren) zijn een veelbelovende techniek om Large Language Models (LLM's) te sturen naar specifieke gedragingen of vaardigheden. Bestaande methoden werken echter voornamelijk in de modelruimte (gewichtsruiimte) of de activatieruimte (interne verborgen staten).

• Gewichtsruiimte: Vereist vaak volledige fine-tuning per taak om de vector te extraheren, wat rekenkundig duur en niet schaalbaar is.
• Activatieruimte: Vereist vaak ingrijpende manipulatie van interne staten of complexe optimalisatie, wat de flexibiliteit beperkt en de toepasbaarheid in de praktijk vermindert.

De auteurs stellen dat er behoefte is aan een methode die trainingsvrij (training-free) en niet-invasief is, die geen gewichtsupdates vereist en geen extra input-tokenkosten met zich meebrengt, terwijl het toch effectief de prestaties van LLM's verbetert.

Methodologie: DeCoVec

De auteurs introduceren DeCoVec (Decoding Space based Task Vector), een raamwerk dat taakvectoren direct construeert in de decodingsruimte (output logit-distributies) door gebruik te maken van In-Context Learning (ICL).

Het proces verloopt in drie hoofdstappen:

1. Context Constructie:

   • Er worden twee soorten contexten voor dezelfde taak $T$ gegenereerd:
     • Zero-Shot Context ($C_{zs}$): Alleen de instructie en de query, zonder voorbeelden. Dit vertegenwoordigt de intrinsieke instructie-volgende capaciteit van het model.
     • Few-Shot ICL Context ($C_{icl}$): De instructie, de query, en een set van $k$ voorbeelden (demonstraties) die via een samplingstrategie (bijv. KATE of BM25) zijn geselecteerd.

2. Constructie van de Task Vector:

   • In plaats van gewichten of interne activaties te analyseren, berekent DeCoVec de logit-distributies ($z$) voor beide contexten tijdens het generatieproces.
   • De taakvector $v_T$ wordt gedefinieerd als het verschil tussen de logits van de few-shot context en de zero-shot context:
     $$v_T = z_{icl} - z_{zs}$$
   • Deze vector encodeert de "essentie" van de taak (de distilleerde taakspecifieke signalen) die door de ICL-context wordt geactiveerd.

3. Sturing tijdens Decoding (Steering):

   • Tijdens het genereren van een nieuw token voor een testinput, wordt de standaard few-shot context gebruikt om de basis-logits ($z_{de}$) te berekenen.
   • De berekende task vector $v_T$ wordt vervolgens in de decoding-proces "geinjecteerd" via een schalingsfactor $\lambda$:
     $$\tilde{z}_t = z_{de} + \lambda \cdot v_T$$
   • De gegenereerde token wordt bepaald op basis van deze aangepaste logits $\tilde{z}_t$. Dit stuurmechanisme is online en token-voor-token, zonder de modelarchitectuur te wijzigen.

Belangrijkste Bijdragen

• Nieuwe Ruimte voor Task Vectors: De eerste methode die task vectors direct in de decodingsruimte (logits) construeert in plaats van in gewichten of interne activaties.
• Trainingsvrij en Niet-Invasief: De methode vereist geen backpropagation, geen fine-tuning, en geen toegang tot de interne parameters van het model. Het werkt puur via de output-distributie.
• Plug-and-Play: Het is een lichtgewicht oplossing die losgekoppeld is van de interne modelarchitectuur en werkt met bestaande, bevroren LLM's.
• Robuustheid: De methode is robuust tegen variaties in de volgorde van demonstraties en werkt effectief zonder extra input-tokenkosten (in feite wordt de outputlengte vaak verkort).

Resultaten

De auteurs hebben DeCoVec geëvalueerd op zeven verschillende open-source LLM's (variërend van 0.5B tot 9B parameters, waaronder Qwen2, Yi, Llama-2/3 en Gemma-2) op drie benchmarks:

1. TruthfulQA: Een dataset voor het testen van waarheidsgetrouwheid.
2. Math-500: Een subset van wiskundige redeneertaken.
3. AQUA-RAT: Een dataset voor algebraïsche woordproblemen.

Kernbevindingen:

• Prestatieverbetering: DeCoVec overtreft consistent standaard few-shot baselines. De gemiddelde verbetering in nauwkeur varieerde van +0.59% tot +5.50%. De grootste winst (+5.50%) werd behaald op Qwen2-1.5B.
• Schaalbaarheid: De methode werkt even goed op kleine modellen (0.5B) als op grotere modellen (9B), wat suggereert dat het taakrepresentaties effectief extrahert ongeacht het aantal parameters.
• Foutanalyse: Op wiskundige taken (Math-500) reduceerde DeCoVec significant logische fouten en degeneratie (herhaling/onzin), hoewel er een lichte toename was in rekenfouten. Dit bevestigt dat de vector de hoog-level logische structuur van de taak verbetert.
• Efficiëntie: De methode leidt tot kortere outputtekens (reductie van 10-40 tokens), wat de inferentiekosten verlaagt ondanks de extra berekening van de vector.
• Robuustheid: De prestaties blijven stabiel ongeacht de volgorde van de voorbeelden (sequentieel, omgekeerd, willekeurig), wat een zwak punt is bij traditionele ICL.

Significantie

DeCoVec biedt een fundamenteel nieuw perspectief op het sturen van LLM's. In plaats van de "zwarte doos" van de modelgewichten of interne lagen aan te raken, gebruikt het de decodingsruimte als een semantisch rijk en direct interface.

• Het lost het probleem op van hoge rekenkosten en complexiteit bij het aanpassen van LLM's voor specifieke taken.
• Het demonstreert dat de verschillen in output-distributies tussen zero-shot en few-shot contexten voldoende informatie bevatten om een krachtige taakvector te vormen.
• Het biedt een praktische, schaalbare oplossing voor het verbeteren van de betrouwbaarheid en logica van LLM's zonder de noodzaak van extra training of zware infrastructuur.

Samenvattend positioneert dit paper DeCoVec als een efficiënte, niet-invasieve alternatief voor traditionele fine-tuning en gewichtsbased task vectors, met name waardevol in scenario's waar resources beperkt zijn of waar snelle aanpassing aan nieuwe taken vereist is.