IA2: Alignment with ICL Activations Improves Supervised Fine-Tuning

1. Het Probleem

Grote Taalmodellen (LLMs) worden vaak aangepast voor specifieke taken via Supervised Fine-Tuning (SFT). SFT traint de modelgewichten om gewenste antwoorden te genereren op basis van gelabelde datasets. Hoewel SFT efficiënt is in inferentie (geen extra context nodig), heeft het twee belangrijke beperkingen, vooral in data-scarce (few-shot) scenario's:

• Slechte generalisatie en overfitting: SFT-modellen neigen naar "shortcut learning" (het leren van oppervlakkige patronen) en overfitting op de trainingsdata, wat leidt tot slechtere prestaties op nieuwe, onbekende data.
• Moeilijk kalibratie: SFT-modellen zijn vaak slecht gekalibreerd; ze geven hoge zekerheid aan foutieve antwoorden.

In tegenstelling hieraan biedt In-Context Learning (ICL) (het geven van voorbeelden in de prompt tijdens inferentie) vaak betere generalisatie en kalibratie, maar ten koste van aanzienlijke inferentiële rekenkracht en contextruimte.

De kernvraag van dit paper is: Kunnen de rijke, informatieve interne berekeningen (activaties) van ICL worden gebruikt om de kwaliteit van SFT te verbeteren?

2. Methodologie: IA2 (ICL Activation Alignment)

De auteurs stellen vast dat ICL en SFT, hoewel ze soms vergelijkbare output geven, fundamenteel verschillende interne activatiepatronen genereren. ICL activeert complexe circuits om generaliseerbare patronen uit de context te halen, terwijl SFT vaak directer naar de output streeft.

Om dit gat te dichten, introduceren ze IA2 (ICL Activation Alignment), een zelfdistillatie-techniek die als een "priming"-stap dient vóór standaard SFT. De methode bestaat uit twee fasen:

1. IA2 Priming (Fase 1):

   • Het doel is om de interne activaties van het model te laten overeenkomen met die van ICL, zonder noodzakelijkerwijs de outputtekst te kopiëren.
   • Het proces:
     • Het model voert ICL uit op een dataset (input $X$ + demonstraties $I$) om een respons $\hat{Y}$ en bijbehorende activatietensor $A_{ICL}$ te genereren.
     • Vervolgens wordt het model geprompt met alleen de query $X$ (zonder demonstraties), maar wordt het getraind om de activaties die ontstaan bij deze query ($A_{SFT}$) te laten overeenkomen met de activaties van de ICL-respons ($A_{ICL}$).
   • De verliesfunctie is de Mean Squared Error (MSE) tussen de activatietensors:
     $$L_{IA2} = \sum ||A_{SFT} - A_{ICL}||^2$$
   • Dit dwingt het model om de "functionele logica" van ICL in zijn gewichten te internaliseren.

2. Standaard SFT (Fase 2):

   • Na de IA2-priming wordt het model verder getraind met standaard cross-entropy loss op de grondwahrheid (ground truth) antwoorden ($Y$).
   • Dit zorgt voor de uiteindelijke uitlijning met de gewenste output, maar het model start nu vanuit een "ICL-gealigneerde" gewichtsubruimte.

3. Belangrijkste Bijdragen

• Ontdekking van Functionele Discrepantie: De auteurs tonen empirisch aan dat ICL en SFT verschillende activatiepatronen produceren, zelfs op dezelfde data. SFT-modellen zijn functioneel niet uitgelijnd met ICL, wat verklaart waarom ze minder goed generaliseren en slechter gekalibreerd zijn.
• IA2 Methode: Ze introduceren een nieuwe trainingstroom die ICL-activaties gebruikt als een trainingssignaal om SFT-modellen te "primen". Dit creëert een brug tussen de inferentie-gebaseerde aanpassing van ICL en de gewichtsgebaseerde aanpassing van SFT.
• Subruimte Overlap Analyse: Door analyse van de LoRA-gewichten tonen ze aan dat IA2-modellen een aanzienlijke overlap hebben met ICL-gedrag, terwijl standaard SFT-modellen orthogonaal (onafhankelijk) zijn. Dit bewijst dat IA2 toegang geeft tot een deel van de gewichtsspace dat voor standaard SFT ontoegankelijk is.

4. Experimentele Resultaten

De auteurs hebben hun methode getest op 12 populaire benchmarks (inclusief sentimentanalyse, nieuwsclassificatie, wiskunde en meerkeuzevragen) met twee model families: Qwen3-4B en Llama-3.2 (1B/3B). Ze trainden in totaal meer dan 13.000 modellen.

• Prestaties (Accuracy): De combinatie IA2 → SFT overtreft consistent standaard SFT (alleen) in nauwkeurigheid, vaak zelfs in data-scarce settings (bijv. N=4 voorbeelden). In veel gevallen komt de prestatie dicht in de buurt van of zelfs boven die van ICL.
• Kalibratie (ECE): Een van de sterkste resultaten is de verbetering in Expected Calibration Error (ECE). IA2 → SFT-modellen zijn aanzienlijk beter gekalibreerd dan standaard SFT-modellen, wat betekent dat hun betrouwbaarheidsschattingen nauwkeuriger zijn.
• Multi-token vs. Single-token: De methode werkt goed voor zowel korte antwoorden (single-token) als langere generaties (multi-token, zoals wiskundige redeneringen), hoewel de verbetering in multi-token settings soms iets minder extreem is vanwege de complexiteit van het comprimeren van lange contexten in kleine LoRA-gewichten.
• Vergelijking met andere methoden: IA2 presteert beter dan kennisdistillatie op zachte labels (soft labels) en is consistent effectief in zowel single- als multi-token taken.

5. Betekenis en Conclusie

Dit paper biedt een conceptueel inzicht in hoe modellen zich aanpassen: ICL en SFT zijn niet functioneel equivalent, maar ICL bevat waardevolle informatie over generalisatie die in SFT verloren gaat.

• Praktische Impact: IA2 biedt een eenvoudige, maar krachtige techniek om SFT-modellen te verbeteren zonder de inferentiekosten van ICL te hoeven betalen. Het maakt modellen robuuster tegen overfitting en beter gekalibreerd.
• Wetenschappelijke Impact: Het paper ondermijnt de aanname dat ICL en SFT slechts verschillende manieren zijn om tot hetzelfde doel te komen. Het toont aan dat het internaliseren van de interne processen (activaties) van ICL cruciaal is voor het verbeteren van de aanpassingskwaliteit van SFT.

Kortom, door de "gedachten" van het model tijdens ICL te kopiëren naar de gewichten van een SFT-model, kunnen we de beste eigenschappen van beide werelden combineren: de efficiëntie van SFT met de generalisatiekracht van ICL.