SAEmnesia: Erasing Concepts in Diffusion Models with Supervised Sparse Autoencoders

Sintesi Tecnica: SAEmnesia: Cancellazione di Concetti nei Modelli di Diffusione tramite Autoencoder Sparsi Supervisionati

Problematica
La cancellazione dei concetti (concept unlearning) nei modelli di diffusione text-to-image è attualmente ostacolata dal "feature splitting", un fenomeno in cui un singolo concetto è distribuito su molti neuroni latenti anziché essere localizzato in un'unica unità. Questa rappresentazione distribuita, nota come polisemanticità, rende difficile l'identificazione e la rimozione di specifici concetti. Gli approcci esistenti richiedono spesso ricerche esaustive tra numerosi accoppiamenti latenti per trovare il sottoinsieme corretto da modificare, rendendo il processo computazionalmente costoso. Inoltre, i latenti sovrapposti sfumano i confini dei concetti, aumentando il rischio di effetti collaterali indesiderati su concetti correlati durante l'intervento.

Metodologia
Gli autori introducono SAEmnesia, un framework che utilizza un autoencoder sparso (SAE) supervisionato per imporre mappature uno-a-uno tra concetti e neuroni latenti. A differenza degli approcci SAE non supervisionati che scoprono le feature post-hoc, SAEmnesia integra la etichettatura supervisionata direttamente nel processo di addestramento utilizzando i prompt testuali usati per generare le immagini.

I componenti tecnici principali includono:

1. Centralizzazione delle Feature: Etichettando sistematicamente i concetti durante l'addestramento, il metodo ottiene una "centralizzazione delle feature", legando ogni concetto a un singolo neurone interpretabile. Questo impedisce la frammentazione dei concetti su più neuroni.
2. Funzione di Perdita Supervisionata: L'obiettivo dell'addestramento combina la standard perdita di ricostruzione non supervisionata con una componente supervisionata ($L_{supSAE}$). Questa perdita supervisionata consiste in:
   • Perdita di Assegnazione del Concetto (CA Loss): Incoraggia latenti assegnati specifici ad attivarsi fortemente quando i corrispondenti concetti sono presenti nell'input.
   • Perdita di Decorrelazione (DC Loss): Penalizza le correlazioni tra i pattern di attivazione dei latenti assegnati a diversi macro-gruppi di concetti (ad esempio, oggetti vs stili) per promuovere il disincanto (disentanglement).
   • Regolarizzazione della Sparsità: Un termine $L_1$ per garantire che la maggior parte dei latenti rimanga inattiva per ogni dato input.
3. Inferenza e Cancellazione: Per cancellare un concetto target, il metodo richiede la guida (steering) di un singolo latente. L'attivazione di questo specifico latente viene moltiplicata per un valore negativo (normalizzato rispetto all'attivazione media) per sopprimere l'influenza del concetto. Questo processo è reversibile; scollegando l'SAE si ripristina esattamente il modello originale.

Contributi Chiave
Il documento delinea tre contributi primari:

1. Framework SAE Supervisionato: L'introduzione di un SAE supervisionato che impone esplicitamente mappature uno-a-uno tra concetto e latente, eliminando il feature splitting e consentendo un controllo trasparente sulle rappresentazioni concettuali.
2. Efficienza nella Cancellazione: Dimostrare che questa struttura permette la cancellazione di un concetto guidando un singolo latente, riducendo sostanzialmente la necessità di ricerca degli iperparametri durante l'inferenza.
3. Miglioramenti delle Prestazioni: Ottenere risultati superiori sul benchmark UnlearnCanvas rispetto agli approcci meccanicistici allo stato dell'arte, un miglioramento della scalabilità nella cancellazione sequenziale, una maggiore robustezza agli attacchi avversari e una soppressione efficace dei contenuti NSFW.

Risultati Sperimentali
Valutato sul benchmark UnlearnCanvas utilizzando Stable Diffusion v1.5, SAEmnesia dimostra miglioramenti significativi rispetto al metodo basato su SAE allo stato dell'arte attuale (SAeUron) e altri approcci di unlearning:

• Prestazioni del Benchmark: SAEmnesia raggiunge un punteggio medio del 91,51% nella cancellazione di oggetti, un miglioramento assoluto del 9,22% rispetto a SAeUron (82,29%). Mostra guadagni costanti in Unlearning Accuracy (UA), In-domain Retain Accuracy (IRA) e Cross-domain Retain Accuracy (CRA).
• Efficienza Computazionale: Riducendo lo spazio di ricerca dall'esplorazione di molteplici combinazioni di latenti alla regolazione di un singolo moltiplicatore per concetto, il metodo riduce i requisiti di ricerca degli iperparametri del 96,67% (da 210 valutazioni a 7).
• Unlearning Sequenziale: Nei compiti che coinvolgono la rimozione sequenziale di nove oggetti, SAEmnesia migliora l'accuratezza della cancellazione del 28,4% rispetto al baseline, dimostrando una scalabilità superiore.
• Robustezza: Sotto attacchi avversari white-box (UnlearnDiffAtk), SAEmnesia mantiene una precisione di cancellazione del 57,50% rispetto al 34,20% di SAeUron. Mostra inoltre resilienza agli attacchi black-box (Ring-A-Bell).
• Soppressione NSFW: Sul benchmark I2P per la rimozione della nudità, il metodo raggiunge risultati allo stato dell'arte (9 rilevamenti contro 18 per SAeUron) preservando al contempo la qualità della generazione.

Significato e Rivendicazioni
Il documento posiziona SAEmnesia come un passo avanti verso una "cancellazione precisa e controllata dei concetti" nei modelli generativi. Trasformando le rappresentazioni distribuite in feature monosemantiche concentrate, il metodo affronta le sfide fondamentali di interpretabilità che hanno ostacolato le tecniche di unlearning precedenti. Gli autori affermano che questo approccio consente un controllo fine-grained sul comportamento generativo senza modificare i pesi del modello base, offrendo un meccanismo completamente reversibile e trasparente per rimuovere contenuti dannosi, inappropriati o protetti da copyright. Il lavoro sottolinea che, sebbene serva come strumento di ricerca per un deployment più sicuro, deve essere applicato insieme a rigorosi meccanismi di valutazione e governance per garantirne l'affidabilità.