Replay What Matters: Off-Policy Replay for Efficient LLM Reinforcement Unlearning

Technische Samenvatting: Replay What Matters: Off-Policy Replay voor Efficiënt LLM Reinforcement Unlearning

Probleemstelling

Het ontleren (unlearning) van Large Language Models (LLM's) heeft als doel schadelijke kennis (bijv. auteursrechtelijk beschermde inhoud, privé-informatie) te verwijderen uit getrainde modellen zonder de prohibitieve kosten van volledige hertraining. Recente reinforcement learning (RL)-gebaseerde benaderingen, zoals RULE, herformuleren ontleren als het aanleren van een weigeringsgedrag: het model leren om vragen gerelateerd aan een "forget set" te weigeren, terwijl de antwoorden op een "retain set" en een "boundary set" behouden blijven.

De auteurs identificeren echter een kritieke inefficiëntie in de Stage 2 on-policy optimalisatie (specifiek met behulp van GRPO) die wordt toegepast door methoden zoals RULE. Tijdens de training ontstaat er een uitgesproken hard–easy disparity (hard-gemak verschil):

1. Easy Cases (Gemakkelijke gevallen): Veel prompts convergeren snel naar hoog-belonende toestanden. Het voortdurend samplen van deze prompts in latere trainingsstadia levert weinig informatieve gradiënten op (advantage $\approx 0$), maar verbruikt aanzienlijke computationele middelen.
2. Hard Cases (Moeilijke gevallen): Prompts nabij de semantische grens tussen vergeten en behouden blijven vaak gedurende de hele training laag-belonend. In standaard on-policy settings worden de laag-belonende rollouts gegenereerd voor deze harde gevallen na eenmalig gebruik weggegooid, ondanks dat ze waardevolle informatie bevatten over de foutmodi van het model.

Deze dynamiek leidt tot suboptimale trainingsefficiëntie, waarbij rekenkracht wordt verspild aan geconvergeerde gemakkelijke gevallen, terwijl moeilijke gevallen ondergeoptimaliseerd blijven, wat potentieel de capaciteit van het model om algemene bruikbaarheid (Retain Quality) te behouden (terwijl specifieke doelen worden vergeten) verslechtert.

Methodologie: ReRULE

Om dit aan te pakken, stellen de auteurs ReRULE (Replay-enhanced Reinforcement UnLEarning) voor, een off-policy replay-verbeteringsmechanisme geïntegreerd in het RULE-framework.

Framework Overzicht

ReRULE behoudt de tweestapsstructuur van RULE:

• Stage 1 (Rejection Steering): Supervised Fine-Tuning (SFT) op een weigeringsdataset om het vermogen van het model om te weigeren een warm-start te geven.
• Stage 2 (Refusal Boundary Optimization): RL-gebaseerde optimalisatie om de grens tussen forget- en retain-queries te verfijnen.

De ReRULE Pipeline

De kerninnovatie vindt plaats in Stage of Stage 2, waar ReRULE een replay buffer en hybride training introduceert:

1. Hard-Case Identificatie en Opslag:
   Tijdens de vroege GRPO-training monitort het systeem de gemiddelde beloning van de rollout-groepen gegenereerd voor elke prompt. Als de gemiddelde beloning van een groep onder een vooraf gedefinieerde drempelwaarde $\tau$ valt, wordt deze geclassificeerd als een "hard case" en opgeslagen in een replay buffer $\mathcal{B}$. Deze buffer accumuleert laag-belonende trajecten die de huidige foutmodi van het model vertegenwoordigen.

2. Hybride Training (Off-Policy Updates):
   Zodra de buffer voldoende harde gevallen bevat ($|\mathcal{B}| \ge B_{min}$) en een warm-up periode ($t \ge T_{warm}$) is voltooid, activeert ReRULE een hybride trainingsregime. Bij elke stap voert het het volgende uit:

   • Standaard On-Policy Update: Een standaard GRPO-update met verse samples van de huidige policy $\pi_\theta$.
   • Off-Policy Replay Update: Een aanvullende gradiënt-update met behulp van samples getrokken uit de replay buffer $\mathcal{B}$.

3. Importance Sampling:
   Om de distributieverschuiving te corrigeren tussen de policy die de opgeslagen rollouts genereerde ($\pi_{\theta_{old}}$) en de huidige policy ($\pi_\theta$), gebruikt ReRULE importance sampling. De off-policy gradiënt wordt berekend als:
   $$\nabla_\theta J_{off}(\theta) = \mathbb{E}_{g \sim \mathcal{B}} \left[ \mathbb{E}_{y \sim \pi_{\theta_{old}}} \left[ \tilde{\rho}(y|x) \nabla_\theta \log \pi_\theta(y|x) \hat{A}(x, y) \right] \right]$$
   waarbij $\tilde{\rho}(y|x)$ een geclipt en genormaliseerd importance gewicht is. Dit stelt het model in staat om te hertrainen op moeilijke gevallen zonder dat daarvoor nieuwe, dure sampling vereist is.

Belangrijkste Bijdragen

1. Identificatie van Trainingsinefficiëntie: De paper toont empirisch aan dat on-policy GRPO in LLM-ontleren lijdt onder afnemende meeropbrengsten bij gemakkelijke gevallen en onderbenutting van moeilijke gevallen, wat leidt tot een "hard/easy disparity" die convergentie belemmert.
2. Voorstel van ReRULE: De auteurs introduceren een off-policy replay-mechanisme dat laag-belonende hard-case rollouts opslaat en hergebruikt via importance-sampled updates. Dit richt de computationele focus naar de randgevallen die verdere optimalisatie vereisen.
3. Theoretische Analyse: De auteurs leveren een theoretisch bewijs dat ReRULE een nauwere convergentiebound oplevert voor harde gevallen vergeleken met pure on-policy RULE. Door het effectieve aantal trainingstokens in de harde regio te verhogen ($N_{H}^{ReRULE} = N_{H}^{on} + \eta_\rho N_{H}^{off}$), vermindert de methode de exploratiefout in de harde regio.
4. Empirische Validatie: Uitgebreide experimenten over drie benchmarks (RWKU, MUSE-Books, TOFU) valideren de effectiviteit van de methode.

Experimentele Resultaten

De auteurs evalueerden ReRULE tegenover baselines inclusief GA, NPO, SimNPO en de originele RULE over drie datasets:

• MUSE-Books (Copyrighted Unlearning): Deze setting, die Harry Potter-content betreft, vertoont een hoog aandeel harde gevallen. ReRULE verbeterde de Retain Quality aanzienlijk van 46.3 (RULE) naar 56.2, terwijl het een top-tier Forget Quality behield. De trainingstijd nam slechts met 5–11% toe (van 6.8h naar 7.1h).
• RWKU (Real-World Knowledge): ReRULE behaalde de beste Forget Quality en behield top-2 prestaties op Forget Naturalness en Retain Quality, waarmee het RULE overtrof in efficiëntie op harde gevallen.
• TOFU (Entity Unlearning): In deze setting is de forget set kleiner en over het algemeen makkelijker (minder harde gevallen). ReRULE toonde beperkte verbetering ten opzichte van RULE, wat de auteurs toeschrijven aan het feit dat het replay-mechanisme het meest voordelig is wanneer er een uitgesproken hard/easy disparity aanwezig is.
• Ablatie-studies:
  • Het vervangen van de hard-case selectie door random storage zorgde ervoor dat de prestaties degradeerden naar het niveau van standaard RULE, wat bevestigt dat het targeten van harde gevallen cruciaal is.
  • Het variëren van de drempelwaarde $\tau$ toonde aan dat het te hoog instellen (inclusief gemakkelijke gevallen) de efficiëntie vermindert, terwijl de optimale instelling (0.4 voor MUSE) een balans vindt tussen selectiviteit en bufferomvang.

Betekenis en Claims

De paper claimt dat ReRULE een kosteneffectieve verbetering biedt voor RL-gebaseerd ontleren. Door "mislukte" rollouts van harde gevallen te recyclen, verbetert het de Retain Quality van het ontleerde model—een kritieke metriek die vaak wordt aangetast in ontleringstaken—zonder dat volledige hertraining of significante extra rekenkracht nodig is.

De auteurs benadrukken dat de methode conditioneel is: de voordelen zijn het meest uitgesproken in scenario's met een aanzienlijke dispariteit tussen gemakkelijke en moeilijke prompts (bijv. MUSE-Books). In eenvoudigere settings (bijv. TOFU) zijn de winsten marginaal, wat de hypothese ondersteunt dat het replay-mechanisme specifiek gericht is op de inefficiëntie van ondergeoptimaliseerde harde gevallen. Het werk positioneert ReRULE als een praktische stap naar efficiënter en robuuster machine ontleren, met name voor het behoud van de algemene bruikbaarheid van het model terwijl specifieke schadelijke kennis wordt verwijderd.