Reinforcement Learning for Neural Model Editing

Technische Samenvatting: Reinforcement Learning voor het Bewerken van Neurale Modellen

Probleemstelling
Het aanpassen van vooraf getrainde neurale netwerken om specifieke doelstellingen te bereiken—zoals conceptverwijdering, bias-mitigatie of het opnemen van nieuwe kennis—vereist traditioneel gespecialiseerde, handmatig ontworpen algoritmen. Deze methoden zijn vaak tijdrovend om te ontwerpen en missen generaliseerbaarheid over verschillende taken heen. Dit artikel onderzoekt of het proces van het bewerken van neurale modellen geautomatiseerd kan worden door het te formuleren als een Reinforcement Learning (RL) probleem. De kernonderzoeksvraag is of een RL-agent een beleid (policy) kan leren om de gewichten van een neuraal netwerk te wijzigen met behulp van uitsluitend beloningsfeedback, zonder toegang tot gradiënten of handmatig geconstrueerde datasets voor de specifieke bewerkingsopdracht.

Methodologie
De auteurs stellen een exploratief kader voor waarbij een agent interageert met een neuraal netwerk door iteratief de gewichten bij te werken op basis van een beloningssignaal. Het kader bestaat uit twee afzonderlijke omgevingen:

1. MaskWorld: De agent wijzigt de gewichten door ze multiplicatief te schalen. De actie $a_t$ schaalt de gewichten van een specifieke laag $l$ zodanig dat $W^{(l)}_t = W^{(l)} \odot a_t$, waarbij $a_t$ een vector van constanten tussen 0.0 en 1.0 is.
2. ShiftWorld: De agent wijzigt de gewichten door additieve updates toe te passen. De actie $a_t$ verschuift de gewichten zodanig dat $W^{(l)}_t = W^{(l)} + a_t$, waarbij $a_t$ constanten bevat tussen -0.05 en 0.05.

Om de computationele onhandelbaarheid van het voorspellen van hoogdimensionale actievectoren (bijv. $m \times n$ voor een gewichtsmatrix) aan te pakken, maken de auteurs gebruik van een door Low-Rank Adaptation (LoRA) geïnspireerd mechanisme voor de taak van machine unlearning. In plaats van de volledige gewicht-update-matrix te voorspellen, voorspelt de agent twee laag-rang matrices $A$ en $B$, wat de dimensionaliteit van de actieruimte vermindert en de optimalisatie van de policy stabiliseert.

Beloningsfunctie
De beloningsfunctie $R$ is een gewogen combinatie van twee componenten:
$$R = \alpha U(w) + \beta O(w)$$

• Utility Term ($U(w)$): Moedigt het behoud van de algemene modelprestaties aan (bijv. nauwkeurigheid op een "retain"-set of een algemene dataset).
• Editing Objective Term ($O(w)$): Kwantificeert het succes van de specifieke bewerkingsdoelstelling (bijv. het minimaliseren van de nauwkeurigheid op een "forget"-set of het reduceren van bias-metrieken).
• Coëfficiënten ($\alpha, \beta$): Regelen de afweging tussen het behoud van utility en het bereiken van de bewerkingsdoelstelling. Opvallend genoeg wordt $\beta$ in de context van unlearning-taken op een negatieve waarde ingesteld om prestaties op de forget-set te bestraffen.

De agent wordt getraind met Proximal Policy Optimization (PPO) met een MLP policy-netwerk. Cruciaal is dat de agent geen gradiënten ontvangt; hij leert uitsluitend door middel van trial-and-error op basis van de prestaties van het aangepaste model op gesamplede batches.

Belangrijkste Bijdragen en Experimentele Resultaten
Het kader werd geëvalueerd op twee verschillende taken: machine unlearning in beeldclassificatie en bias-mitigatie in tekstclassificatie.

• Machine Unlearning (MNIST): Het doel was om de cijferklasse 7 te "vergeten" (unlearn) terwijl de prestaties op klassen 0–6 en 8–9 behouden blijven.

  • Resultaten: De geleerde policies verminderden de nauwkeurigheid op de forget-set (klasse 7) tot bijna 0% (0.0000 voor MaskWorld, 0.0006 voor ShiftWorld).
  • Utility Behoud: In tegen tegen de verwachtingen van prestatievermindering, vertoonden de aangepaste modellen een lichte verbetering in nauwkeurigheid op de retain-set (van 92.84% naar ~93.5%) vergeleken met het originele model.
  • Observatie: De agents bereikten deze resultaten terwijl ze slechts 15% van de trainingsdata tijdens de optimalisatie hadden ingezien, wat duidt op een sterke generalisatie naar de testset.

• Bias Mitigation (Jigsaw Toxic Comments): Het doel was om spurieuze correlaties te verminderen (bijv. de associatie van het woord "black" met toxiciteit) terwijl de algemene classificatie-utility behouden blijft.

  • Resultaten: De aangepaste modellen verbeterden de nauwkeurigheid op een synthetische onbevooroordeelde testset met ongeveer 6–7% (bijv. ShiftWorld verbeterde van 55.36% naar 62.52%).
  • Utility Behoud: De nauwkeurigheid op de algemene dataset bleef stabiel en daalde met slechts ~1% voor MaskWorld, wat aangeeft dat de agent erin slaagde de balans tussen bias-reductie en utility-behoud te vinden.

Significantie en Claims
Het artikel beweert niet beter te presteren dan gespecialiseerde state-of-the-art algoritmen voor deze specifieke taken. In plaats daarvan ligt de primaire significantie in het demonstreren van een paradigmaverschuiving: het bewerken van neurale modellen kan worden beschouwd als een leerbaar RL-probleem in plaats van een handmatig ontworpen proces.

De auteurs beweren dat hun bevindingen suggereren dat RL een algemeen en uitbreidbaar kader biedt voor het leren van bewerkingspolicies vanuit beloningsfeedback. Deze aanpak maakt de ontdekking van gewichtswijzigingen mogelijk die gerichte doelstellingen bereiken (zoals unlearning of debiasing) zonder dat er voor elke nieuwe doelstelling een taakspecifiek algoritmisch ontwerp nodig is. Het werk dient als een proof-of-concept om toekomstig onderzoek naar geleerde benaderingen voor modelbewerking op gebieden zoals eerlijkheid (fairness), robuustheid en privacy te stimuleren.

Beperkingen
De auteurs erkennen dat het huidige kader te maken heeft met de schaalbaarheidsuitdagingen die inherent zijn aan RL. Hoewel het door LoRA geïnspireerde mechanisme helpt bij het omgaan met hoogdimensionale lagen, blijft het volledig optimaliseren van alle lagen van een groot model onhaalbaar vanwege de instabiliteit veroorzaakt door de enorme actieruimte. Toekomstig werk suggereert het verkennen van multi-agent settings en speltheoretische benaderingen om deze beperkingen aan te pakken.