Joint Flashback Adaptation for Forgetting-Resistant Instruction Tuning

Probleemstelling: Catastrophical Forgetting bij LLM's

Grote Taalmodellen (LLM's) presteren uitstekend in diverse taken, maar ondervinden grote moeite met het incrementeel leren van nieuwe taken zonder het vergeten van eerder aangeleerde kennis. Dit fenomeen staat bekend als catastrophical forgetting.

• Huidige beperkingen: Bestaande oplossingen zoals Experience Replay (het hergebruiken van oude trainingsdata), Regularisatie (het beperken van parameterveranderingen) en Task Differentiatie (toewijzen van specifieke adapters per taak) stuiten op praktische obstakels.
  • Replay: Vaak onmogelijk door data-privacy, gebrek aan toegang tot oorspronkelijke trainingsdata (bij open-source modellen) en hoge rekentkosten.
  • Regularisatie: Kan leiden tot onder- of over-beperkingen en vereist soms nog steeds toegang tot oude data.
  • Task Differentiatie: Vereist kennis van het taaktype of een taak-ID, wat de flexibiliteit beperkt in real-world scenario's waar taken gemengd of onbekend zijn.

Het doel van dit onderzoek is een methode te ontwikkelen die een bestaand, goed afgestemd LLM aanpast aan nieuwe taken terwijl de prestaties op oude taken behouden blijven, zonder gebruik te maken van replay-data of taakdifferentiatie.

Methodologie: Joint Flashback Adaptation (JFA)

De auteurs stellen Joint Flashback Adaptation (JFA) voor, een raamwerk dat twee kerncomponenten combineert: het gebruik van beperkte "flashbacks" en gezamenlijk leren van latente taken.

1. Flashback Adaptation (Behoud van Oude Kennis)

In plaats van volledige datasets van oude taken te gebruiken, introduceert JFA flashbacks: een zeer beperkt aantal prompts (zonder labels) van oude taken.

• Implementatie: Voor elke oude taak worden willekeurig $N$ prompts geselecteerd (handmatig of gesynthetiseerd).
• Verliesfunctie: Het model wordt getraind om de output voor deze flashbacks zo dicht mogelijk bij de output van het originele (referentie) model te houden. Dit wordt gemeten via een divergentieverlies ($L_{DIV}$), specifiek de tweezijdige KL-divergentie tussen de token-distributies van het huidige model en het referentiemodel.
• Doel: Dit beperkt de afwijkingen in het modelgedrag voor oude taken, waardoor kennisbehoud wordt geforceerd zonder de volledige dataset te hoeven opslaan.

2. Joint Task Learning (Gezamenlijk Leren & Interpolatie)

Om de data-schaarste van de flashbacks te overbruggen en kennisdeling te bevorderen, introduceert JFA latente taken.

• Latente Taken: Deze fungeren als interpolaties tussen flashbacks en nieuwe taken. Ze worden weergegeven als paren van een sleutel (vector) en een bijbehorende gewichtsvergroting ($\Delta\theta$).
• Retrieval: Voor elke input (nieuwe taak of flashback) wordt de input geëncodeerd (bijv. met RoBERTa). Vervolgens worden de $k$ meest vergelijkbare latente taken opgehaald via cosine-sequentie met de opgeslagen sleutels.
• Reparameterisatie: De uiteindelijke modelparameters voor een specifieke sample worden berekend door de basisparameters te combineren met een gewogen som van de gewichtsvergrotingen van de opgehaalde latente taken. Dit stelt het model in staat om kennis te delen tussen nieuwe taken en flashbacks via deze latente tussenruimte.
• Gradient Projection: Om conflicten te voorkomen tussen het leren van nieuwe taken en het verminderen van vergeten, wordt PCGrad (Projecting Gradients) gebruikt. Dit projecteert de gradiënten van de verschillende doelen (nieuwe taak vs. flashback-behoud) zodat ze niet in tegenstrijdige richtingen werken.

De totale trainingsdoelstelling is een combinatie van het toezicht op de nieuwe taken ($L_{SFT}$) en het behoud van de flashbacks ($L_{DIV}$), gebalanceerd door een hyperparameter $\alpha$.

Belangrijkste Bijdragen

1. Nieuwe Methode (JFA): Een task-agnostische aanpak voor continu leren die geen toegang vereist tot replay-data of taak-identificatie.
2. Flashback-mechanisme: Het bewijzen dat een zeer beperkt aantal prompts (zonder labels) voldoende is om vergeten te voorkomen, wat privacy en opslagkosten drastisch verlaagt.
3. Gezamenlijk Leren: Het introduceren van latente taken om de data-schaarste van flashbacks te compenseren en kennisdeling tussen oude en nieuwe taken te faciliteren.
4. Gradient Projection: Het effectief toepassen van PCGrad om de stabiliteit-plasticiteitsdilemma op te lossen.

Resultaten

De auteurs hebben JFA uitgebreid getest op state-of-the-art modellen (Vicuna-13B en Llama-3.1-8B) over meer dan 1000 instructie-volgende taken, rekenkundige redenering en algemeen redeneren. De datasets omvatten Super Natural Instructions (SNI), GSM-8K, SVAMP, ARC-Challenge, BIG-Bench-Hard (BBH) en MMLU-Pro.

• Prestaties op Nieuwe Taken: JFA behaalde de beste prestaties op de nieuwe taak (SNI) vergeleken met SFT (Supervised Fine-Tuning), PACE en andere continu-leer methoden, wat wijst op superieure generalisatie.
• Verminderen van Vergeten: Op oude taken (zoals GSM-8K en BBH) presteerde JFA consistent het beste of als tweede beste, aanzienlijk beter dan standaard SFT en vaak beter dan methoden die wel replay-data gebruiken (zoals Replay en FROMP).
• Efficiëntie: De methode vereist slechts een handvol prompts (bijv. 30 prompts in totaal verdeeld over oude taken) in plaats van duizenden trainingsvoorbeelden.
• Ablatie Studies:
  • Het toevoegen van Joint Task Learning verbeterde zowel de nieuwe als de oude taken.
  • Het gebruik van Flashbacks (zelfs met 10-20 prompts) was cruciaal; zonder flashbacks daalden de prestaties op oude taken drastisch.
  • De hyperparameter $\alpha$ bleek effectief in het balanceren tussen het leren van nieuwe taken en het behouden van oude kennis.

Significantie en Conclusie

Deze paper introduceert een praktische en schaalbare oplossing voor het probleem van catastrophical forgetting in LLM's. De belangrijkste doorbraak is dat data-privacy en opslagbeperkingen niet langer een barrière hoeven te zijn voor continu leren. Door slechts een klein aantal "flashback" prompts te gebruiken en deze te verrijken met latente taken, kan een model nieuwe vaardigheden leren zonder de oude te verliezen.

Dit maakt JFA bij uitstek geschikt voor real-world toepassingen waar:

1. Oude trainingsdata niet beschikbaar is (door privacy of licenties).
2. Taken dynamisch veranderen en niet vooraf gekend zijn (task-agnostisch).
3. Rekenkracht beperkt is, waardoor het opslaan van grote replay-datasets onpraktisch is.

De auteurs erkennen wel beperkingen, zoals de afhankelijkheid van de kwaliteit van het referentiemodel en het gebrek aan fijnmazige controle over welke specifieke vaardigheden behouden moeten worden, maar JFA biedt een sterke basis voor toekomstig onderzoek in continu leren voor generatieve AI.