DualOptim+: Bridging Shared and Decoupled Optimizer States for Better Machine Unlearning in Large Language Models

Technische Samenvatting van DualOptim+

Probleemstelling
Machine Unlearning (MU) in Large Language Models (LLMs) heeft tot doel de invloed van specifieke trainingsdata (de 'forget set') te wissen, terwijl de algemene bruikbaarheid van het model op de overige data (de 'retain set') behouden blijft. Bestaande aanpakken staan voor een aanzienlijke uitdaging in het balanceren van deze tegenstrijdige doelstellingen.

• Gecombineerde Bijwerking: De meeste methoden minimaliseren de som van de 'forget' en 'retain' verliezen gelijktijdig. Dit leidt vaak tot een aanzienlijke verslechtering van de modelbruikbaarheid door overmatig vergeten.
• Afwisselende Bijwerking: Het afwisselen tussen 'forget' en 'retain' doelstellingen verbetert de prestaties, maar lijdt aan instabiliteit en gevoeligheid voor hyperparameters.
• Gekoppelde Optimizers (DualOptim): Een eerdere aanpak, DualOptim, koppelde de optimizer-toestanden (bijvoorbeeld momentum-termen) voor elke doelstelling los om gradiëntverstrengeling te voorkomen. Hoewel effectief in computer vision, levert DualOptim slechts marginale verbeteringen op bij LLM's. De auteurs observeren dat bij LLM-unlearning de gradiënten voor 'forget' en 'retain' niet altijd in direct conflict staan; ze delen vaak gemeenschappelijke representaties, vooral in latere trainingsfasen. Volledig gekoppelde toestanden maken geen gebruik van deze gedeelde kenmerken, terwijl volledig gedeelde toestanden geen omgaan met conflicten.

Methodologie: DualOptim+
Het artikel stelt DualOptim+ voor, een plug-and-play optimalisatiekader ontworpen om de kloof te overbruggen tussen volledig gedeelde en volledig gekoppelde optimizer-toestanden. De kerninnovatie is de decompositie van de optimizer-toestand in twee componenten:

1. Gedeelde Basistoestand ($B$): Deze toestand vangt de gemeenschappelijke representaties op die worden gedeeld door zowel de 'forget' als de 'retain' doelstellingen. Deze wordt gezamenlijk bijgewerkt met behulp van gradiënten van beide doelstellingen ($\nabla_\theta L_f$ en $\nabla_\theta L_r$).
2. Gekoppelde Delta-toestanden ($\Delta_f, \Delta_r$): Deze toestanden behouden de doelstellingsspecifieke residuen. Ze worden bijgewerkt op basis van het verschil tussen de gradiënt van de specifieke doelstelling en de bias-correcteerde basistoestand.
   • $\Delta_f \leftarrow \beta \Delta_f + (1-\beta)(\nabla_\theta L_f - \hat{B})$
   • $\Delta_r \leftarrow \beta \Delta_r + (1-\beta)(\nabla_\theta L_r - \hat{B})$

Parameterbijwerking:
De uiteindelijke momentum-term die wordt gebruikt voor parameterbijwerkingen, wordt gereconstrueerd door de bias-correcteerde basistoestand te combineren met de respectievelijke bias-correcteerde delta-toestand ($\hat{B} + \hat{\Delta}_f$ of $\hat{B} + \hat{\Delta}_r$).

Belangrijk Theoretisch Inzicht:
Theoretische en numerieke analyses tonen aan dat DualOptim+ fungeert als een adaptief tussenniveau.

• Wanneer gradiënten positief gecorreleerd zijn (gedeelde kenmerken domineren), convergeren de delta-toestanden naar nul, en gedraagt de methode zich als een Afwisselend schema (met gebruik van de gedeelde basis).
• Wanneer gradiënten negatief gecorreleerd zijn (conflicterende doelstellingen), convergeert de basistoestand naar nul, en gedraagt de methode zich als een DualOptim schema (met gebruik van gekoppelde toestanden).
  Deze adaptiviteit stelt de optimizer in staat zich dynamisch aan te passen op basis van de directionele conflicten tussen 'forget' en 'retain' gradiënten.

DualOptim+ 8bit
Om het geheugenvolume te adressen dat gepaard gaat met het opslaan van extra optimizer-toestanden (3x het geheugen van een standaard Adam-optimizer), stellen de auteurs DualOptim+ 8bit voor. Deze variant kwantiseert de optimizer-toestanden naar 8-bit precisie, wat het geheugengebruik aanzienlijk verlaagt (tot ongeveer 3/4 van een standaard Adam) zonder de prestaties te compromitteren.

Belangrijkste Bijdragen

1. Nieuw Kader: De voorstelling van DualOptim+, die een gedeelde basistoestand en gekoppelde delta-toestanden introduceert om effectief gedeelde en gekoppelde optimizer-toestanden te overbruggen.
2. Plug-and-Play Ontwerp: Een kader dat compatibel is met elke optimizer die toestanden opslaat (bijvoorbeeld AdamW, Muon) en toepasbaar is op multi-objectief optimalisatie.
3. Empirische Validatie: Uitgebreide experimenten over fictieve datasets (TOFU), real-world scenario's (Llama-3 unlearning), veiligheidsuitlijning en multi-task learning taken.
4. Efficiëntie: De introductie van de 8-bit gekwantiseerde variant die het geheugenvolume beperkt.

Experimentele Resultaten
Het artikel evalueert DualOptim+ tegenover Joint, Alternate en DualOptim-baselines over diverse LLM's (Phi-1.5, Llama-2, Llama-3) en taken:

• Fictieve Scenario's (TOFU): DualOptim+ behaalt consistent de beste Algemene Prestatie (OVR), gedefinieerd als een afweging tussen Vergeet-Efficiëntie (FE) en Modelbruikbaarheid (MU). Het presteert beter dan DualOptim, vooral in gericht unlearning, door beter gebruik te maken van gedeelde kenmerken.
• Real-World Scenario's: Bij het onthouden van echte individuen uit Llama-3 behoudt DualOptim+ een superieure bruikbaarheid op downstream-taken (ARC-c, MMLU, enz.) in vergelijking met baselines, terwijl DualOptim soms onderpresteert ten opzichte van de eenvoudigere Alternate-aanpak.
• Veiligheidsuitlijning: In een taak om onveilige kennis te vergeten terwijl instructievolgende capaciteiten behouden blijven, bereikt DualOptim+ de hoogste veiligheids-bruikbaarheid-afweging en de laagste over-weigeringsgraad op de XSTest-benchmark.
• Multi-Task Learning: In tegenstelling tot unlearning, waar gradiëntconflicten hoog zijn, levert DualOptim+ bij multi-task learning, waar conflicten minder uitgesproken zijn, nog steeds aanzienlijke prestatieverbeteringen op, terwijl DualOptim negatieve winsten toont.
• Geheugenefficiëntie: De 8-bit varianten bereiken een bijna gelijkwaardig geheugengebruik met standaard baselines, terwijl ze de prestatievoordelen van de full-precision methode behouden.

Betekenis en Claims
De auteurs claimen dat DualOptim+ een stabielere en effectievere afweging biedt tussen kennisverwijdering en behoud van bruikbaarheid dan bestaande methoden. Ze positioneren de methode niet alleen als een unlearning-tool, maar als een generaliseerbaar optimalisatiekader dat toepasbaar is op bredere scenario's met meerdere conflicterende leerdoelstellingen, zoals LLM-uitlijning en multi-objectief learning. Het artikel benadrukt dat de adaptieve aard van het overbruggen van gedeelde en gekoppelde toestanden cruciaal is voor het hanteren van de dynamische gradiëntrelaties die worden aangetroffen bij LLM-unlearning.