Unlearning for One-Step Generative Models via Unbalanced Optimal Transport

Die Arbeit stellt UOT-Unlearn vor, einen neuartigen Plug-and-Play-Ansatz auf Basis des unausgeglichenen optimalen Transports, der das maschinelle Vergessen in effizienten einstufigen Generativmodellen ermöglicht, indem er das Entfernen unerwünschter Klassen mit dem Erhalt der generativen Qualität durch eine glatte Umverteilung der Wahrscheinlichkeitsmasse verbindet.

Das Wesentliche

Das große Problem: Der "Ein-Schritt"-Künstler und das vergessene Geheimnis

Stell dir vor, du hast einen genialen Künstler (einen KI-Modell), der Bilder malt.

• Die alten Künstler (Diffusionsmodelle): Diese brauchen viel Zeit. Sie nehmen einen Haufen grauen Nebel und entfernen Schritt für Schritt den Nebel, bis ein Bild entsteht. Das dauert lange (vielleicht 50 oder 100 Schritte), aber das Ergebnis ist toll.
• Der neue Künstler (Ein-Schritt-Modelle): Dieser ist ein Genie. Er nimmt den grauen Nebel und malt das fertige Bild in einem einzigen, blitzschnellen Hieb. Das ist super schnell, aber es gibt ein Problem: Was passiert, wenn dieser Künstler Bilder von Dingen malt, die wir nicht mehr sehen wollen? Zum Beispiel urheberrechtlich geschützte Charaktere oder unangemessene Inhalte?

Normalerweise müsste man den Künstler komplett neu ausbilden, um ihn zu "vergessen" zu lassen. Das ist aber so, als würde man einen ganzen Kunstunterricht von vorne beginnen, nur weil er ein bestimmtes Motiv nicht mehr malen soll. Das kostet zu viel Zeit und Geld.

Bisher gab es Methoden, um Künstlern das Vergessen beizubringen, aber diese funktionierten nur für die langsamen, schrittweisen Künstler. Wenn man sie auf den schnellen Ein-Schritt-Künstler anwendete, ging das Bild kaputt oder der Künstler wurde verrückt.

Die Lösung: UOT-Unlearn (Der geschickte Umverteiler)

Die Forscher aus Korea haben eine neue Methode namens UOT-Unlearn entwickelt. Sie nutzen ein mathematisches Konzept namens "Unbalanced Optimal Transport" (Ungleichgewichtiger Optimaler Transport).

Klingt kompliziert? Stell dir das so vor:

Die Analogie: Das überfüllte Festmahl

Stell dir vor, der KI-Künstler ist ein Koch, der ein riesiges Buffet mit vielen Gerichten (Klassen) anbietet: Pizza, Sushi, Burger, Salat und... Giftige Beeren (die Klasse, die wir vergessen wollen).

1. Das Problem: Der Koch serviert immer noch die giftigen Beeren. Wir wollen, dass er aufhört, sie zu servieren.
2. Die alte, falsche Methode: Man sagt dem Koch: "Mach die Beeren einfach weg!" Aber weil der Koch so schnell ist (Ein-Schritt), wenn er die Beeren einfach weglässt, entsteht eine Lücke auf dem Buffet. Die Gäste bekommen dann plötzlich nur noch leere Teller oder seltsame, undefinierbare Matsch-Sachen. Das Buffet sieht kaputt aus.
3. Die neue Methode (UOT-Unlearn):
   • Wir sagen dem Koch nicht einfach "Weg damit". Wir sagen: "Wenn du eine giftige Beere auf den Teller legst, bekommst du einen riesigen, schmerzhaften Stachel in den Finger." (Das ist die Strafkosten).
   • Aber wir sagen ihm auch: "Du darfst die Beeren nicht einfach in den Müll werfen. Du musst sie in andere, leckere Gerichte umwandeln." (Das ist die Verteilung).
   • Weil der Koch die giftigen Beeren nicht mehr servieren will (wegen des Stachels), fängt er an, die Zutaten der Beeren geschickt in die Pizza oder den Salat zu mischen.

Das Ergebnis?

• Es gibt keine giftigen Beeren mehr auf dem Buffet (die KI lernt das Verbotene).
• Aber das Buffet sieht immer noch toll aus, und die Pizza schmeckt sogar noch besser, weil sie jetzt ein bisschen mehr "Beeren-Zutat" hat, die aber harmlos ist (die KI behält ihre Qualität).

Was macht die Methode so besonders?

1. Kein neues Lernen nötig: Der Koch muss nicht das ganze Kochbuch neu lernen. Er passt nur seine Handbewegung für das eine verbotene Gericht an.
2. Keine echten Daten nötig: Normalerweise braucht man Beispiele von den "guten" Gerichten, um zu lernen, wie man sie verbessert. Diese Methode kommt aber nur mit den Bildern, die der Koch selbst schon gemalt hat, aus. Er braucht keine echten Fotos von echten Pizzen mehr. Das ist wie ein Koch, der nur mit seinen eigenen Erinnerungen experimentiert.
3. Kein Chaos: Andere Methoden haben oft dazu geführt, dass die KI anfing, wirre, unscharfe Bilder zu machen, wenn sie etwas vergaß. Diese Methode sorgt dafür, dass die verbotenen Bilder einfach in harmlose, schöne Bilder verwandelt werden.

Das Ergebnis im echten Leben

Die Forscher haben das an echten KI-Modellen getestet (z. B. auf Bildern von Autos, Fröschen oder Schiffen).

• Ergebnis: Die KI hat die verbotenen Bilder fast zu 100% vergessen.
• Qualität: Die anderen Bilder (die sie behalten durfte) sahen genauso gut aus wie vorher.
• Vergleich: Andere Methoden haben entweder die Bilder nicht richtig vergessen oder die Qualität der restlichen Bilder total ruiniert.

Zusammenfassung in einem Satz

Stell dir vor, du willst einem schnellen Künstler beibringen, ein bestimmtes Motiv zu vergessen, ohne ihn neu ausbilden zu müssen. Die neue Methode sagt ihm: "Mal das nicht mehr, aber verwandle die Idee stattdessen in etwas anderes Schönes, damit dein gesamtes Kunstwerk nicht kaputtgeht."

Das ist UOT-Unlearn: Ein cleverer Trick, um KI-Künstler sicher zu machen, ohne ihre Geschwindigkeit oder ihre Kunstfertigkeit zu opfern.

Technische Zusammenfassung

1. Problemstellung

Generative Modelle, insbesondere Diffusionsmodelle, haben zwar hohe Bildqualitäten erreicht, leiden jedoch unter langsamen Inferenzzeiten aufgrund ihrer iterativen Denoisingschritte. Um dies zu lösen, wurden One-Step-Generative Modelle (z. B. Consistency Models, Flow Maps) entwickelt, die die Abbildung von Rauschen zu Daten in einem einzigen Vorwärtsschritt durchführen.

Mit der zunehmenden Leistungsfähigkeit und Geschwindigkeit dieser Modelle steigt jedoch das Risiko der Generierung unerwünschter Inhalte (z. B. NSFW-Inhalte oder urheberrechtlich geschütztes Material). Machine Unlearning (das gezielte Vergessen bestimmter Datenklassen) ist eine notwendige Sicherheitsmaßnahme.

• Das Kernproblem: Bestehende Unlearning-Methoden für Diffusionsmodelle basieren auf der Manipulation von Rauschvorhersagen oder Gradienten in mehreren iterativen Denoisingschritten. Diese Ansätze sind für One-Step-Modelle nicht direkt anwendbar, da diese keine Zwischenzustände besitzen, in die eingegriffen werden könnte.
• Die Herausforderung: Es fehlt ein Framework, das das Vergessen einer Klasse in einem einzigen Vorwärtsschritt ermöglicht, ohne die Qualität der verbleibenden Klassen zu zerstören oder auf teure Retraining-Prozesse zurückzugreifen.

2. Methodik: UOT-Unlearn

Die Autoren schlagen UOT-Unlearn vor, das erste Plug-and-Play-Framework für das Klassen-Unlearning in One-Step-Modellen, basierend auf Unbalanced Optimal Transport (UOT).

Kernkonzept: Unbalanced Optimal Transport (UOT)

Im Gegensatz zum klassischen Optimal Transport (OT), der eine strikte Übereinstimmung der Randverteilungen (Marginals) erzwingt, erlaubt UOT eine Relaxierung dieser Randbedingungen.

• Ziel: Es wird ein Ziel-Funktionswert minimiert, der einen Trade-off zwischen den Transportkosten und der Abweichung der Verteilungen (gemessen durch $f$-Divergenzen) darstellt.
• Anwendung auf Unlearning: Dies ermöglicht es, die Wahrscheinlichkeitsmasse der zu vergessenden Klasse („Forget Class") nicht einfach zu löschen (was zu Artefakten führen würde), sondern sie glatt auf die verbleibenden Klassen („Retain Classes") zu verteilen.

Der Algorithmus

1. Formulierung als Transportproblem:
   • Die Quelle ($\mu$) ist die vortrainierte Verteilung des Modells.
   • Das Ziel ($\nu$) ist die ursprüngliche Datenverteilung (ohne die zu vergessende Klasse).
   • Ein Unlearning-Kostenfunktion ($c_{ul}$) wird eingeführt, die Proben innerhalb des „Vergessens-Bereichs" ($R_f$) stark bestraft.
2. Kostenfunktion ($c_{ul}$):
   • Forget Cost: Für generierte Bilder, die im Merkmalsraum nahe dem „Vergessens-Anker" ($\mu_f$) liegen, wird eine Strafe angewendet, die die Features über einen bestimmten Rand ($m$) hinaus vom Anker wegbewegt.
   • Retain Cost: Für Bilder außerhalb dieses Bereichs wird eine $L_2$-Distanz zwischen dem Output des vortrainierten Modells ($G_{pre}$) und des feinabgestimmten Modells ($G_\theta$) bestraft, um die Fidelity (Treue) der verbleibenden Klassen zu bewahren.
3. Optimierung:
   • Das Problem wird über ein semi-duales UOT-Formular gelöst, das ein neuronales Netzwerk für die Transportkarte ($\Delta T_\theta$) und ein duales Potential ($v_\phi$) verwendet.
   • Dateneffizienz: Ein entscheidender Vorteil ist, dass das Verfahren keine echten Retain-Daten benötigt. Es nutzt nur synthetische Proben des vortrainierten Modells und einen einmalig berechneten Ankervektor für die zu vergessende Klasse.

3. Hauptbeiträge

• Erstes Framework für One-Step-Modelle: UOT-Unlearn ist die erste Methode, die Unlearning speziell für die Architektur von One-Step-Generatoren (Flow Maps) entwickelt hat.
• Neue Zielfunktion: Formulierung des Unlearning-Problems als UOT-Problem, das die Wahrscheinlichkeitsmasse der Zielklasse durch eine $f$-Divergenz-Strafe sanft auf andere Klassen umverteilt, anstatt sie in Rauschen kollabieren zu lassen.
• Datenschutz und Effizienz: Das Verfahren benötigt während der Optimierungsphase keine realen Daten (Zero-Data-Requirement), was Datenschutzbedenken minimiert und den Prozess effizienter macht.
• Plug-and-Play: Die Methode ist architekturunabhängig und kann in bestehende One-Step-Modelle integriert werden, ohne deren Struktur zu ändern.

4. Experimentelle Ergebnisse

Die Methode wurde auf CIFAR-10 und ImageNet-256 mit Architekturen wie Consistency Trajectory Models (CTM) und MeanFlow evaluiert.

• Metriken:
  • PUL (Percentage of Unlearning): Misst den Erfolg beim Entfernen der Zielklasse.
  • u-FID (Unlearned FID): Misst die Qualität der verbleibenden Klassen.
• Ergebnisse:
  • UOT-Unlearn übertrifft etablierte Baselines (Gradient Ascent, Selective Amnesia, SalUn, VDU) signifikant.
  • Es erreicht hohe PUL-Werte (z. B. >90% bei CIFAR-10), während der u-FID-Wert (Qualitätsverlust) minimal bleibt.
  • Im Gegensatz zu Baselines, die bei starkem Unlearning oft die gesamte Verteilung verzerren oder Artefakte erzeugen, behält UOT-Unlearn die strukturelle Integrität der verbleibenden Daten bei.
  • Auf ImageNet-256 zeigte sich, dass UOT-Unlearn komplexe semantische Verschiebungen (z. B. bei aquatischen Klassen) handhabt, ohne die Bildqualität drastisch zu senken (u-FID von ~20 vs. ~80 bei Baselines).

5. Bedeutung und Ausblick

Die Arbeit schließt eine kritische Lücke im Bereich des Machine Unlearning für die nächste Generation schneller Generativer Modelle.

• Sicherheit: Sie ermöglicht es, One-Step-Modelle sicher zu machen, ohne deren Hauptvorteil (Geschwindigkeit) zu opfern.
• Theoretischer Fortschritt: Die Anwendung von Unbalanced Optimal Transport auf das Unlearning-Problem bietet einen neuen, prinzipiellen Ansatz, um den Zielkonflikt zwischen „Löschen" und „Erhalten" mathematisch zu lösen.
• Praxis: Da keine echten Retain-Daten benötigt werden, ist die Methode besonders für Szenarien geeignet, in denen Datenzugriff eingeschränkt ist oder Datenschutz eine Rolle spielt.

Zusammenfassend demonstriert das Paper, dass Unlearning in One-Step-Modellen nicht durch Anpassung alter iterativer Methoden, sondern durch eine Neuformulierung als Verteilungstransportproblem effizient und qualitativ hochwertig gelöst werden kann.