The Linear Geometry of Interpretable Tokens: Jailbreaking Attacks and Defenses for Unlearned Diffusion Models

Technische Samenvatting: De Lineaire Geometrie van Interpreteerbare Tokens in Unlearned Diffusion Models

Probleemstelling
Diffusionmodellen (DM's), in het bijzonder Text-to-Image (T2I) modellen, blinken uit in het genereren van hoogwaardige content, maar vormen veiligheidsrisico's omdat ze potentieel schadelijke concepten (bijv. naaktheid, auteursrechtelijk beschermde stijlen of specifieke objecten) kunnen onthouden en reproduceren. Machine Unlearning (MU)-methoden zijn ontwikkeld om deze concepten uit getrainde modellen te "wissen" via fine-tuning zonder vanaf nul te hoeven hertrainen. Echter, bestaande unlearning-methoden hebben moeite om schadelijke concepten volledig te verwijderen terwijl de generatiekwaliteit van veilige content behouden blijft. Bijgevolg blijven unlearned modellen kwetsbaar voor jailbreaking-aanvallen, waarbij tegenstanders prompts manipuleren om de gewiste concepten te regenereren.

Huidige jailbreaking-benaderingen (bijv. adversarial prompt optimalisatie of embedding-gebaseerde aanvallen) vertrouwen vaak op niet-interpreteerbare, optimalisatiegebaseerde perturbaties. Hoewel effectief, bieden deze methoden beperkt inzicht in hoe schadelijke concepten voortbestaan binnen de interne staat van het model. Dit gebrek aan interpreteerbaarheid belemmert de ontwikkeling van robuuste, gerichte defensiestrategieën. Bovendien zijn lichtgewicht, post-hoc defensiemechanismen die opereren op reeds unlearned modellen minder onderzocht vergeleken met adversarial re-training benaderingen.

Methodologie
De auteurs stellen voor dat gewiste concepten niet volledig verdwijnen, maar voortbestaan als een coherente, interpreteerbare lineaire subspace binnen de token embedding-ruimte van de CLIP text encoder die wordt gebruikt door Latent Diffusion Models (LDMs). Op basis van deze hypothese introduceren zij twee complementaire methoden: SubAttack en SubDefense.

1. SubAttack (Subspace Attack):

   • Doel: Om de residuele subspace van een schadelijk concept te identificeren en uit te buiten.
   • Mechanisme: De methode leert een set orthogonale attack token embeddings $\{v_{att, k}\}$ die de residuele subspace beslaan. Elke embedding wordt beperkt tot een niet-negatieve lineaire combinatie van bestaande vocabulaire token embeddings ($v_{att} = \sum \alpha_i v_i, \alpha_i \geq 0$).
   • Optimalisatie: Een MLP-netwerk wordt getraind om de coëfficiënten $\alpha_i$ te optimaliseren om de denoising loss te minimaliseren wanneer de gegenereerde afbeelding het doelgerichte schadelijke concept bevat.
   • Deflatie: Om meerdere orthogonale embeddings te leren, gebruiken de auteurs een deflatie-strategie. Na het leren van $v_{att, k}$, projecteren zij de gehele vocabulaire orthogonaal op $v_{att, k}$ voordat de volgende embedding wordt geleerd. Dit waarborgt diversiteit en legt de volledige subspace van het residuele concept vast.
   • Interpreteerbaarheid: Omdat de embeddings lineaire combinaties zijn van bestaande tokens, kunnen de top-gewogen tokens worden geëxtraheerd en geanalyseerd (bijv. via WordClouds) om de semantische componenten (bijv. "slave," "hips," "nip") te onthullen die het unlearned model nog steeds associeert met het gewiste concept.

2. SubDefense (Subspace Defense):

   • Doel: Om de geïdentificeerde residuele subspace van SubAttack te neutraliseren.
   • Mechanisme: Zodra de attack subspace $V_{att}$ (gespannen door de geleerde orthogonale embeddings) is geïdentificeerd, projecteert SubDefense de token embeddings van de vocabulaire van het unlearned model op de nullspace van $V_{att}$.
   • Update Regel: Voor elke token embedding $v_i$ in de vocabulaire wordt de embedding bijgewerkt als $v_{def, i} = v_i - \text{Proj}_{V_{att}}(v_i)$.
   • Plug-and-Play: Deze operatie is lichtgewicht en vereist geen hertraining van de gewichten van het diffusion model. Het kan worden toegepast op elk unlearned model, inclusief die welke een adversarial fine-tuning hebben ondergaan.

Belangrijkste Bijdragen

• Identificatie van Interpreteerbare Residuele Subspaces: Het artikel demonstreert dat schadelijke concepten in unlearned DM's voortbestaan als een lineaire subspace in de token embedding-ruimte, samengesteld uit menselijk interpreteerbare tekstuele elementen.
• SubAttack: Een nieuwe, interpreteerbare jailbreaking-aanval die orthogonale token embeddings leert om deze subspace te beslaan. Het presteert beter dan bestaande baselines in Attack Success Rate (ASR) en vertoont een hoge transfereerbaarheid over verschillende prompts, initiële ruis en unlearned modellen.
• SubDefense: Een robuuste, lichtgewicht defensiemechanisme dat de geïdentificeerde residuele subspace wegprojecteert. Het verbetert de robuustheid tegen diverse jailbreaking-aanvallen (inclusief adversarial prompts en embedding-aanvallen) terwijl het de kwaliteit van veilige generatie beter behoudt vergeleken met bestaande defensies zoals RECE.
• Inzicht in Unlearning Beperkingen: Het werk onthult dat unlearning vaak oppervlakkige trefwoorden onderdrukt, maar faalt in het elimineren van diepere, impliciete associaties (bijv. "dumpster" voor "garbage truck" of "slave" voor "nudity"), die worden gevangen door de lineaire subspace.

Experimentele Resultaten

• Attack Performance: SubAttack behaalde een significant hogere ASR dan baselines (UnlearnDiff en CCE) over meerdere concepten (naaktheid, Van Gogh-stijl, kerk) en unlearned modellen (ESD, FMN, UCE, SPM). Bijvoorbeeld, op het concept "nudity" over ESD, FMN, UCE en SPM, behaalde SubAttack ASR's van respectievelijk 97,56%, 100%, 81,67% en 74,89%, waarmee het CCE en UnlearnDiff overtrof.
• Transfereerbaarheid: Attack embeddings geleerd op één unlearned model (bijv. ESD) vertaalden effectief naar anderen (FMN, SPM, UCE), en zelfs terug naar het originele Stable Diffusion model, met ASR's die consequent boven de 80% lagen.
• Defense Performance: SubDefense verminderde de ASR aanzienlijk voor diverse aanvallen. Bijvoorbeeld, bij het verdedigen van UCE tegen SubAttack op het "Van Gogh"-concept, daalde de ASR van 98,33% naar 29,44%. Het verminderde ook de ASR voor UnlearnDiff en CCE aanvallen.
• Kwaliteitsbehoud: In tegen tegenstelling tot sommige defensies die de beeldkwaliteit verslechteren, behield SubDefense of verbeterde het de FID en CLIP scores op veilige prompts (MSCOCO-10k dataset) vergeleken met de baseline RECE-defensie.
• Complementariteit: Er werd aangetoond dat SubDefense effectief is, zelfs wanneer het wordt toegepast op top van adversarially gefinetunede modellen zoals STEREO, wat de ASR verder verlaagde van 21,67% naar 6,67%, wat aangeeft dat het een andere (lineaire residuele) kwetsbaarheid aanpakt dan adversarial training.

Betekenis en Claims
Het artikel claimt dat het voortbestaan van schadelijke concepten in unlearned modellen geen willekeurige ruis is, maar een gestructureerd, lineair fenomeen dat zowel geëxploiteerd als gemitigeerd kan worden. Door het probleem te kaderen vanuit de lens van lineaire geometrie in token embeddings, biedt de auteur een verenigd framework voor het begrijpen van kwetsbaarheden (SubAttack) en het construeren van defensies (SubDefense).

De betekenis ligt in de verschuiving van een opaak, optimalisatiegebaseerd paradigma naar een interpreteerbare, subspace-gebaseerde analyse. Deze benadering legt niet alleen diepere kwetsbaarheden bloot van huidige unlearning technieken, maar biedt ook een praktische, model-agnostische defensiestrategie die geen hertraining vereist. De auteurs benadrukken dat hoewel lineaire projectie geen niet-lineaire resten kan vatten (een beperking die zij erkennen), het aanpakken van de lineaire residuele subspace een cruciale eerste stap vormt naar robuustere en meer interpreteerbare veiligheidsmechanismen voor generatieve AI.