Learning What to Forget: Improving LLM Unlearning via Learned Token-Level Importance

Technische Samenvatting: Leren wat te vergeten: Verbeteren van LLM Unlearning via Geleerde Token-Level Belangrijkheid

Probleemstelling
Machine unlearning heeft als doel om gerichte kennis (bijv. persoonlijke gegevens, auteursrechtelijk beschermde tekst, schadelijke inhoud) te verwijderen uit getrainde Large Language Models (LLM's) zonder de training vanaf nul te hoeven herstarten. Een cruciale uitdaging in autoregressieve LLM's is dat niet alle tokens binnen een "forget sample" (de data die verwijderd moet worden) even relevant zijn voor de gerichte informatie. Veel tokens zijn structureel (bijv. functiewoorden, interpunctie, syntactische patronen) en zijn noodzakelijk voor algemene taalvaardigheden. Bestaande benaderingen voor token-gewogen unlearning proberen de forget-loss selectief toe te passen op specifieke tokens, maar vertrouwen op externe annotaties, hulpmodellen, waarschijnlijkheidsgebaseerde heuristieken of taakspecifieke regels om de relevantie van tokens in te schatten. Deze methoden slagen er vaak niet in om effectief onderscheid te maken tussen forget-specifieke inhoud en structurele tokens, wat leidt tot ofwel onvoldoende vergeten of onnodige degradatie van de algemene bruikbaarheid.

Methodologie: Alternating Token-Weighted Unlearning (ATWU)
De auteurs stellen voor dat token-level forget-specificiteit gekarakteriseerd kan worden door de interactie tussen de forget- en de retain-objectieven zelf, in plaats van door externe signalen. Ze formaliseren dit als een gezamenlijk optimalisatieprobleem over modelparameters ($\theta$) en token-gewichten ($z$).

1. Theoretische Formulering: Het artikel definieert een "retain conflict"-metriek, $\kappa(A)$, die de kosten meet van het onderdrukken van een set tokens $A$ terwijl de optimaliteit op de retain-set behouden blijft. Onder een "retain-conflict separatie"-aanname bewijzen de auteurs dat de oracle-set van forget-specifieke tokens kan worden teruggevonden als de minimizer van een gezamenlijk objectief dat de retain-loss en een gewogen forget-loss balanceert. Dit theoretische resultaat suggereert dat token-gewichten direct geleerd kunnen worden uit de unlearning-objectieven zonder externe supervisie.
2. Algoritme (ATWU): Om dit hanteerbaar te maken, introduceren de auteurs Alternating Token-Weighted Unlearning (ATWU).
   • Scorer: In plaats van onafhankelijke gewichten voor elke token te leren, gebruikt ATWU een lichtgewicht lineaire scorer, $g_w$, geparametriseerd door een vector $w$ in de hidden-state-ruimte van het model. De score voor een token is $\hat{z}_{x,t} = \sigma(\langle w, h_\theta(x_t) \rangle)$, waarbij $h_\theta$ de hidden representatie is.
   • Gezamenlijke Optimalisatie: Het framework optimaliseert een ontspannen, continue objectieven die een retain-loss, een verzadigde negatieve cross-entropy forget-loss (SATGA+) bevat, een entropie-regularisator om scores naar binaire waarden te duwen, en een budget-constraint om het aantal geselecteerde tokens te controleren.
   • Alternerende Updates: De modelparameters $\theta$ en de scorer-parameters $w$ worden alternerend bijgewerkt. Wanneer $\theta$ wordt bijgewerkt, worden de scores losgekoppeld (detached); wanneer $w$ wordt bijgewerkt, is $\theta$ bevroren. Dit stabiliseert de training door de feedbackloop tussen de veranderende modelgeometrie en de tokenselectie te ontkoppelen.

Belangrijkste Bijdragen

• Retain-Conflict Karakterisering: Het artikel biedt een theoretische karakterisering van token-level forget-specificiteit gebaseerd op de spanning tussen de forget- en de retain-objectieven, en bewijst dat een gezamenlijke optimalisatie de oracle forget-specifieke support kan terugvinden onder natuurlijke separatiecondities.
• Ongesuperviseerd Leren van Token-Gewichten: ATWU is een lichtgewicht framework dat token forget-specificiteit gezamenlijk met modelparameters leert met behulp van enkel de interne representaties van het model. Het vereist geen externe token-niveau supervisie, hulpmodellen of dure annotaties.
• Adaptieve Loss Integratie: De methode introduceert SATGA+, een score-gemoduleerde variant van verzadigde negatieve cross-entropy waarbij de geleerde token-score wordt geïnjecteerd in de saturatie-exponent. Dit maakt gladdere, onzekerheidsbewuste updates mogelijk tijdens de vroege fasen van de optimalisatie.

Resultaten
De auteurs evalueren ATWU op twee benchmarks: TOFU (synthetische fictieve biografieën van auteurs) en RWKU (werkelijke kennis over publieke figuren).

• Prestaties: ATWU bereikt state-of-the-art Unlearning Quality (UQ), een metriek die de balans vindt tussen forget quality (FQ) en retain degradation (RD). Op TOFU bereikt het een UQ van 91,7 (meer dan 5,8 pp boven de runner-up), en op RWKU, 58,7 (meer dan 6,8 pp boven de runner-up).
• Behoud van Bruikbaarheid: In tegenstelling tot agressieve unlearning-methoden die de generatiekwaliteit doen instorten, behoudt ATWU de algemene bruikbaarheid (gemeten via MMLU, repetitiviteit en pairwise win rates) dicht bij de originele checkpoint.
• Kwaliteit van Token Scores: De geleerde token scores stemmen significant beter overeen met de grondwaarheid van forget-specifieke spans (gemeten via AUROC) vergeleken met waarschijnlijkheidsgebaseerde heuristieken, hulpmodel-benaderingen en andere token-weging baselines.
• Veelzijdigheid: Het framework is loss-agnostisch; wanneer gecombineerd met andere forget losses (DPO, NPO, SIMNPO), verbetert ATWU consistent de forget–retain trade-off ten opzichte van hun ongegewogen tegenhangers.

Betekenis en Claims
Het artikel claimt dat de "retain conflict" een effectief, intrinsiek criterium biedt voor het identificeren van wat taalmodellen moeten vergeten. Door token-gewogen unlearning te verschuiven van een pipeline van externe heuristieken naar een end-to-end gezamenlijk optimalisatieprobleem, demonstreert ATU dat:

1. De unlearning-objectieven zelf voldoende latente signalen bevatten om forget-specifieke tokens te terugvinden zonder externe supervisie.
2. De bottleneck in selectieve unlearning niet louter het vinden van de juiste tokens is, maar de principiële co-design van de scoring-mechanisme en de forget-objectieven.
3. Schaalbare, privacy-bewuste en zeer gerichte unlearning haalbaar is zonder te vertrouwen op externe annotaties, die vaak privacyrisico's en schaalbaarheidsproblemen met zich meebrengen.

De auteurs vermelden beperkingen, waaronder de afhankelijkheid van een hyperparameter voor het token-selectiebudget en de aanname van een strikte scheiding tussen forget-specifieke en structurele tokens, wat mogelijk slechts benaderend standhoudt in natuurlijke taal. Daarnaast zijn de evaluaties momenteel beperkt tot specifieke benchmarks en enkele runs vanwege de computationele kosten.