Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

Deze paper introduceert UOT-Unlearn, een nieuw framework dat onbalans optimale transport gebruikt om klassen effectief te laten vergeten in één-staps generatieve modellen terwijl de algehele generatiekwaliteit behouden blijft.

De kern

De Probleemstelling: De "Snelle Kunstenaar" die niet stopt

Stel je voor dat je een kunstenaar hebt die foto's kan maken.

• De oude kunstenaars (Diffusiemodellen): Zij werken heel langzaam. Ze beginnen met een vlekje ruis en "ontdoen" zich stap voor stap van de ruis, alsof ze een beeld uit een blok marmer hakken. Het duurt lang, maar het resultaat is prachtig. Als je wilt dat ze stoppen met het maken van bepaalde onderwerpen (bijvoorbeeld geen foto's meer van auto's), kun je ze een paar keer "terugfluiten" tijdens het hakken.
• De nieuwe kunstenaars (One-Step Modellen): Deze zijn revolutionair. Ze maken een foto in één flits. Geen stap-voor-stap hakwerk, maar direct van ruis naar plaatje. Het is razendsnel.

Het probleem: Wat doe je als deze supersnelle kunstenaar iets maakt dat je niet wilt (bijvoorbeeld copyright-materiaal of onfatsoenlijke beelden)?
De oude methoden om hem te "leren vergeten" werken niet. Die methoden zijn gebouwd voor het stap-voor-stap proces. Als je ze op de snelle kunstenaar toepast, is het alsof je probeert een auto te remmen terwijl hij al op de snelweg staat en je geen rempedaal hebt. Het resultaat is vaak dat de kunstenaar helemaal stopt met werken, of dat hij rare, vervormde plaatjes maakt.

De Oplossing: UOT-Unlearn (De Slimme Verhuizer)

De auteurs van dit paper hebben een nieuwe methode bedacht, genaamd UOT-Unlearn. Ze gebruiken een wiskundig concept dat "Unbalanced Optimal Transport" heet. Laten we dit uitleggen met een analogie.

Stel je voor dat de kunstenaar een groot magazijn met ballonnen heeft.

• De rode ballonnen zijn de dingen die je wilt vergeten (bijvoorbeeld auto's).
• De blauwe en gele ballonnen zijn de dingen die je wilt behouden (bijvoorbeeld dieren en landschappen).

Hoe werkt het nu?
In de oude methoden probeerde je de rode ballonnen gewoon weg te gooien. Het probleem was dat het magazijn dan een gat kreeg. De kunstenaar vulde dat gat dan op met rotzooi (ruis) of maakte de blauwe ballonnen lelijk.

De nieuwe methode (UOT-Unlearn) doet iets slims:

1. De "Boete" (De Kosten): De kunstenaar krijgt een zware boete als hij een rode ballon probeert te maken.
2. De "Slimme Verhuizing": In plaats van de rode ballonnen weg te gooien en een gat te laten, zegt de methode: "Oké, je mag die rode ballonnen niet maken, maar je mag ze wel veranderen in blauwe of gele ballonnen."
3. De "Balans": De methode zorgt ervoor dat de rode ballonnen soepel worden omgevormd naar de andere kleuren, zonder dat het magazijn instort of dat de andere ballonnen lelijk worden.

Dit heet "Unbalanced" (ongebalanceerd) omdat we de totale hoeveelheid ballonnen niet strikt vasthouden; we laten toe dat de verdeling verandert, zolang de kwaliteit maar hoog blijft.

Waarom is dit zo goed?

1. Geen oude foto's nodig: Meestal moet je om iets te vergeten, de kunstenaar weer laten oefenen met de goede foto's (de "retaint" data). Deze methode heeft dat niet. Hij gebruikt alleen de foto's die de kunstenaar zelf maakt. Het is alsof de kunstenaar zichzelf corrigeert zonder dat je hem een boekje hoeft te geven.
2. Geen kwaliteitsverlies: Bij andere methoden werden de goede foto's vaak wazig of vreemd. Hier blijven de blauwe en gele ballonnen (de goede foto's) er nog steeds perfect uitzien.
3. Snelheid: Omdat het werkt met de snelle "één-flits" modellen, is het vergeten ook razendsnel.

Samenvattend in één zin:

De auteurs hebben een slimme manier bedacht om supersnelle AI-kunstenaars te "resetten" zodat ze stoppen met het maken van ongewenste onderwerpen, door die onderwerpen niet weg te gooien, maar ze op een natuurlijke manier om te vormen naar iets anders, zonder de kwaliteit van de rest van de kunst te verpesten.

Het is alsof je een snelkookpan hebt die te heet wordt: in plaats van de pan kapot te slaan (de AI volledig hertrainen), regel je de hitte zo dat het eten niet verbrandt, maar wel gaar blijft.

Technische samenvatting

Probleemstelling

Recente doorbraken in generatieve modellen, zoals Flow Map-modellen (bijv. Consistency Trajectory Models en MeanFlow), hebben de efficiëntie van beeldgeneratie aanzienlijk verbeterd door een directe mapping van ruis naar data in één enkele forward-pass mogelijk te maken. Dit lost het probleem van de trage inferentie op bij traditionele diffusion-modellen, die honderden iteratieve stappen vereisen.

Echter, deze krachtige en snelle modellen brengen nieuwe risico's met zich mee, zoals de generatie van ongepaste inhoud (NSFW) of auteursrechtelijk beschermd materiaal. Machine Unlearning (het verwijderen van specifieke concepten uit een getraind model zonder het opnieuw te trainen) is essentieel om deze risico's te mitigeren.

• De Kruis: Bestaande unlearning-methoden zijn ontworpen voor multi-step diffusion-modellen en vertrouwen op het aanpassen van denoising-stappen op specifieke tijdstippen.
• De Uitdaging: Deze methoden zijn fundamenteel incompatibel met one-step generatieve modellen, omdat deze modellen geen tussenstappen hebben tijdens de generatie. Er is geen bestaand kader om concepten veilig te verwijderen uit deze snelle, single-pass architecturen zonder de algehele kwaliteit van de generatie te vernietigen.

Methodologie: UOT-Unlearn

De auteurs stellen UOT-Unlearn voor, het eerste plug-and-play framework voor class unlearning specifiek ontworpen voor one-step generatieve modellen. De kern van de methode is het gebruik van Unbalanced Optimal Transport (UOT).

1. Formulering als Transportprobleem:
   In plaats van het model opnieuw te trainen, wordt het unlearning-proces gezien als het transporteren van de verdeling van het vooraf getrainde model ($p_{pre}$) naar een nieuwe verdeling waarbij het "vergeten" concept is verwijderd.

   • Standaard OT: Vereist strikte randvoorwaarden (marginalen) waarbij de totale massa exact moet worden behouden. Dit is te rigide voor unlearning, waarbij massa juist moet worden verwijderd of herschikt.
   • UOT: Verslapt deze strikte randvoorwaarden door een f-divergentie-straf toe te voegen. Dit creëert een afweging tussen de transportkosten en de afwijking van de oorspronkelijke verdeling.

2. De Unlearning-Doelstelling:
   De methode definieert een unlearning-kostfunctie ($c_{ul}$) die twee rollen vervult:

   • Forget Cost (Actieve Verwijdering): Voor samples die vallen binnen het "vergeten" gebied (bepaald door een cosine-afstand tot een centroid van het vergeten concept), wordt een zware straf opgelegd. Dit dwingt het model om waarschijnlijkheidsmassa uit dit gebied te verplaatsen.
   • Retain Cost (Behoud van Kwaliteit): Voor samples buiten het vergeten gebied wordt een $L_2$-straf opgelegd om de output dicht bij de oorspronkelijke output van het vooraf getrainde model te houden. Dit behoudt de structuur en kwaliteit van de overige klassen.

3. Data-vrije Optimalisatie:
   Een belangrijk kenmerk is dat UOT-Unlearn geen toegang vereist tot echte "retain"-data (de data die bewaard moet blijven) tijdens het optimalisatieproces.

   • Het gebruikt alleen gegenereerde samples van het vooraf getrainde model en een vooraf berekende centroid van het vergeten concept.
   • De doelstelling wordt geoptimaliseerd via een semi-dual formulering waarbij een transportmap ($\Delta T_\theta$) wordt geleerd die de vooraf getrainde generator ($G_{pre}$) fine-tunt naar de nieuwe generator ($G_\theta$).

Belangrijkste Bijdragen

• Eerste Framework voor One-Step Modellen: Het introduceert UOT-Unlearn als het eerste unlearning-systeem dat specifiek is ontworpen voor single-pass generatieve modellen, waardoor de kloof tussen snelle generatie en veiligheidscontrole wordt overbrugd.
• UOT-gebaseerde Formulering: Het stelt een nieuwe doelstelling voor die de waarschijnlijkheidsmassa van het vergeten concept soepel herschikt naar de overige klassen via een f-divergentie-straf, in plaats van dat het model instabiel wordt of ruisachtige samples produceert.
• Data-Efficiëntie: Het framework vereist geen real retain-data tijdens het unlearning-proces, wat privacy en rekenefficiëntie ten goede komt.
• Plug-and-Play: De methode werkt als een fine-tuning-laag en is compatibel met verschillende bestaande one-step architecturen zonder wijzigingen in de modelstructuur.

Experimentele Resultaten

De auteurs evalueren hun methode op CIFAR-10 en ImageNet-256 met modellen zoals CTM en MeanFlow, en vergelijken deze met gevestigde baselines (Gradient Ascent, Selective Amnesia, SalUn, VDU).

• Unlearning Succes (PUL): UOT-Unlearn bereikt consistent de hoogste Percentage of Unlearning (PUL) scores. Het verwijdert het doelconcept effectief (bijv. >90% vermindering in generatie van de doelklasse).
• Behoud van Kwaliteit (u-FID): In tegenstelling tot baselines die vaak leiden tot een sterke degradatie van de beeldkwaliteit (hoge u-FID scores) of instabiliteit, behoudt UOT-Unlearn een lage unlearned FID. Dit betekent dat de overige klassen van hoge kwaliteit blijven.
• Visuele Analyse: Op 2D synthetische data en CIFAR-10 toont de methode aan dat de waarschijnlijkheidsmassa van het vergeten concept soepel wordt verplaatst naar de geldige gebieden van de retain-klassen. Baselines zoals VDU leiden vaak tot vervorming van de totale verdeling of het genereren van samples in ongeldige gebieden.
• ImageNet-256: Zelfs bij hoge resolutie en complexe semantische verschuivingen (bijv. het verwijderen van "goudvissen" terwijl andere waterdieren behouden blijven), behoudt de methode de structurele integriteit van de retain-data, terwijl baselines zoals Gradient Ascent leiden tot ernstige corruptie.

Betekenis en Conclusie

Dit werk is significant omdat het een fundamentele beperking van de huidige generatieve AI-oplossingen oplost: de onmogelijkheid om snelle, one-step modellen veilig te maken via unlearning.

Door unlearning te formuleren als een beperkt optimalisatieprobleem voor probabilistisch transport, biedt de auteurs een theoretisch onderbouwde manier om de afweging tussen "sterk verwijderen" en "kwaliteit behouden" te balanceren. De methode demonstreert dat het mogelijk is om schadelijke of ongewenste concepten uit zeer snelle generatieve modellen te verwijderen zonder de noodzaak van dure hertraining of toegang tot gevoelige trainingsdata, wat een cruciale stap is voor de veilige en ethische implementatie van next-generation generatieve AI.