Unlearning What Matters: Token-Level Attribution for Precise Language Model Unlearning

1. Probleemstelling

Grote Taalmodellen (LLM's) memoriseren vaak gevoelige, auteursrechtelijk beschermde of schadelijke gegevens, wat machine unlearning noodzakelijk maakt om specifieke kennis te verwijderen zonder het model vanaf nul opnieuw te trainen.

• Huidige Beperking: Bestaande unlearning-methoden werken op sequenti niveau, waarbij uniforme gradiëntupdates worden toegepast over alle tokens in een sequentie.
• De Kernkwestie: Kennis wordt niet uniform gecodeerd; slechts een schaarse subset van tokens (bijvoorbeeld specifieke entiteiten, datums, feiten) draagt de gerichte informatie, terwijl andere tokens structurele of contextuele rollen vervullen.
• Gevolgen van Uniforme Updates:
  1. Gradiëntruis: Het updaten van niet-kritieke tokens introduceert ruis, wat het signaal-ruisverhouding (SNR) van het unlearning-proces verslechtert.
  2. Verslechtering van Nut: Indiscriminate updates schaden de prestaties van het model op niet-gerelateerde kennis (bijwerkingen).
  3. Suboptimale Vergeting: Het model faalt om de gerichte kennis effectief te verwijderen omdat het gradiëntsignaal wordt verdund door irrelevante tokens.

2. Methodologie: TokenUnlearn

De auteurs stellen TokenUnlearn voor, een raamwerk dat unlearning verschuift van sequenti niveau naar token niveau door "kennis-kritieke" tokens te identificeren en selectief te targeten.

A. Token Attributiemechanisme

Om te identificeren welke tokens de te vergeten kennis coderen, berekent het raamwerk een samengestelde belangrijkheidsscore ($\phi_i$) voor elke token met behulp van twee signalen:

1. Kennisbewuste Attributie (Maskering):
   • Mechanisme: Het model genereert voorspellingen op de originele invoer ($z_{orig}$) en op een invoer waarbij kennis-relevante delen (specifiek zelfstandige naamwoorden in de prompt) gemaskeerd zijn ($z_{mask}$).
   • Meting: De absolute log-probabiliteitsverschuiving ($\Delta^{unlearn}_i$) wordt berekend. Een grote verschuiving geeft aan dat de voorspelling van de token zwaar afhankelijk is van de gerichte kennis.
2. Entropie-bewust Complementair Signaal:
   • Mechanisme: Berekent de voorspellende entropie ($H_i$) van de token.
   • Redenering: Hoge entropie kan beslispunten aangeven waar het model onzeker is over kennis-afhankelijke voortzettingen (bijvoorbeeld kiezen tussen meerdere plausibele feiten).
   • Combinatie: De uiteindelijke score is een gewogen som: $\phi_i = \alpha \cdot \bar{\Delta}^{unlearn}_i + (1-\alpha) \cdot \bar{H}_i$.

B. Token-selectieve Optimalisatiestrategieën

Zodra belangrijkheidsscores zijn gegenereerd, past het raamwerk unlearning-doelstellingen alleen toe op tokens met hoge belangrijkheid, gebruikmakend van twee strategieën:

1. Harde Selectie:
   • Selecteert de top-$r$ fractie van tokens (bijvoorbeeld de top 20%) gebaseerd op $\phi_i$.
   • Past unlearning-verlies uitsluitend toe op deze tokens; andere tokens ontvangen nul gradiëntupdates.
2. Zachte Weging:
   • Moduleert de gradiëntbijdrage van elke token op basis van een genormaliseerde belangrijkheidsscore (met behulp van een softmax met temperatuur $\tau$).
   • Staat een gladde interpolatie toe tussen uniforme updates en volledig selectieve updates.

C. Integratie met Bestaande Methoden

Het raamwerk is model-onafhankelijk en integreert met standaard gradiënt-gebaseerde unlearning-doelstellingen, waaronder:

• Gradient Ascent (GA)
• Weighted Gradient Ascent (WGA)
• Negative Preference Optimization (NPO)
• Representation Misdirection Unlearning (RMU)
• Behoud: Alle methoden worden gecombineerd met KL-divergentie regularisatie op de "retain set" om algemene capaciteiten te behouden.

3. Theoretische Analyse

Het paper biedt een theoretische rechtvaardiging waarom token-selectie unlearning verbetert:

• Signaal-ruisverhouding (SNR): De auteurs modelleren unlearning als beweging langs een laag-dimensionale deelruimte ($U$) van parameters.
• Ruisreductie: Uniforme updates accumuleren ruis van alle $T$ tokens. Token-selectie concentreert de gradiëntgewicht op de subset $K$ van kennis-kritieke tokens.
• Resultaat: De SNR van de gradiëntschatter verbetert evenredig met het omgekeerde van de selectieverhouding ($T/|K|$). Door updates op niet-kritieke tokens te onderdrukken, reduceert de methode gradiëntruis en behoudt ze ongerelateerde kennis effectiever.

4. Experimentele Resultaten

De auteurs evalueerden TokenUnlearn op twee benchmarks over drie modelarchitecturen (Phi-1.5, Llama-2-7B, Qwen-3-8B).

Benchmarks

1. TOFU (Task of Fictitious Unlearning): Gedetailleerde kennisunlearning (verwijderen van feiten over fictieve auteurs).
2. WMDP (Weapons of Mass Destruction Proxy): Veiligheidsuitlijning (verwijderen van gevaarlijke bio/cybersecuritykennis).

Belangrijkste Bevindingen

• Superieure Vergeting: Token-level varianten (T-* en S-*) presteerden consistent beter dan sequentie-level baselines in vergetingseffectiviteit (lagere nauwkeurigheid op forget sets).
  • Voorbeeld: Op TOFU met Qwen-3-8B verbeterde T-WGA de vergetingseffectiviteit aanzienlijk in vergelijking met standaard WGA.
• Betere Behoud van Nut: Token-level methoden behielden een hogere modelnut (retain scores) op niet-gerichte data.
  • Voorbeeld: Op Qwen-3-8B verbeterde de retain score met maximaal 19,0% in vergelijking met sequentie-level baselines, terwijl tegelijkertijd een betere vergeting werd bereikt.
• Generaliseerbaarheid: De aanpak werkte consistent over verschillende modelgroottes en architecturen.
• Ablatiestudies:
  • Selectieverhouding ($r$): Een verhouding van 20% (top 20% van tokens) bood de optimale afweging tussen vergeten en behoud.
  • Entropiegewicht: Het combineren van entropiesignalen (ongeveer 30% gewicht) bood complementaire voordelen, waarbij onzekere beslispunten werden gevangen die alleen maskering zou kunnen missen.

5. Belangrijkste Bijdragen

1. Token-Level Attributieraamwerk: Geïntroduceerd een principiële methode om kennis-kritieke tokens te identificeren met behulp van counterfactuele maskering en entropiesignalen zonder aanvullende modellen.
2. Selectieve Unlearning-strategieën: Voorstellen van "Harde Selectie" en "Zachte Weging" strategieën die bestaande unlearning-algoritmen uitbreiden om op token niveau te werken.
3. Theoretisch Inzicht: Bewezen dat token-selectieve updates de gradiënt SNR verbeteren, wat verklaart waarom uniforme updates falen om de dubbele doelstellingen van vergeten en behoud te optimaliseren.
4. Empirische Validatie: Gedemonstreerde consistente verbeteringen over drie modellen en twee verschillende benchmarks (kennis en veiligheid), met tot 32,6% verbetering in vergeten en 19,0% verbetering in behoud van nut.

6. Betekenis en Impact

• Privacy & Regulering: Maakt nauwkeurigere naleving van het "Recht om Vergeten" (GDPR/CCPA) mogelijk door specifieke datasporen te verwijderen zonder de algemene nut van het model te verslechteren.
• Veiligheid: Biedt een effectievere mechanisme voor het verwijderen van gevaarlijke capaciteiten (bijvoorbeeld biosecurity-bedreigingen) uit voorgetrainde modellen.
• Efficiëntie: Daagt de aanname uit dat unlearning uniforme sequentieupdates vereist, en suggereert dat kennis in LLM's schaars is en nauwkeurig kan worden aangepakt.
• Beperkingen: De methode introduceert rekenkundige overhead door de extra forward pass voor maskering. Het is een benaderende methode en garandeert geen 100% verwijdering van kennis-sporen (residusporen kunnen bestaan).

Concluderend vertegenwoordigt TokenUnlearn een paradigma-verschuiving in machine unlearning, die overgaat van grofkorrelige sequentie-updates naar fijnkorrelige token-attributie, en zo de fundamentele spanning tussen effectief vergeten en behoud van nut oplost.