Unlearning What Matters: Token-Level Attribution for Precise Language Model Unlearning

1. Enunciato del Problema

I Large Language Models (LLM) memorizzano spesso dati sensibili, protetti da copyright o dannosi, rendendo necessario l'unlearning automatico per rimuovere conoscenze specifiche senza riaddestrare da zero.

• Limitazione Attuale: I metodi di unlearning esistenti operano a livello di sequenza, applicando aggiornamenti del gradiente uniformi su tutti i token in una sequenza.
• Il Problema Centrale: La conoscenza non è codificata in modo uniforme; solo un sottoinsieme sparso di token (ad esempio, entità specifiche, date, fatti) trasporta le informazioni mirate, mentre altri servono ruoli strutturali o contestuali.
• Conseguenze degli Aggiornamenti Uniformi:
  1. Rumore del Gradiente: Aggiornare token non critici introduce rumore, degradando il rapporto segnale-rumore (SNR) del processo di unlearning.
  2. Degradazione dell'Utilità: Aggiornamenti indiscriminati danneggiano le prestazioni del modello su conoscenze non correlate (danni collaterali).
  3. Dimenticanza Subottimale: Il modello non riesce a rimuovere efficacemente la conoscenza mirata perché il segnale del gradiente è diluito da token irrilevanti.

2. Metodologia: TokenUnlearn

Gli autori propongono TokenUnlearn, un framework che sposta l'unlearning dal livello di sequenza al livello di token identificando e prendendo di mira selettivamente i token "critici per la conoscenza".

A. Meccanismo di Attribuzione dei Token

Per identificare quali token codificano la conoscenza da dimenticare, il framework calcola un punteggio di importanza composito ($\phi_i$) per ogni token utilizzando due segnali:

1. Attribuzione Consapevole della Conoscenza (Mascheramento):
   • Meccanismo: Il modello genera previsioni sull'input originale ($z_{orig}$) e su un input in cui le porzioni rilevanti per la conoscenza (in particolare i sostantivi nel prompt) sono mascherate ($z_{mask}$).
   • Metrica: Viene calcolato lo spostamento assoluto della log-probabilità ($\Delta^{unlearn}_i$). Un grande spostamento indica che la previsione del token dipende fortemente dalla conoscenza mirata.
2. Segnale Complementare Consapevole dell'Entropia:
   • Meccanismo: Calcola l'entropia predittiva ($H_i$) del token.
   • Razionale: Un'alta entropia può indicare punti decisionali in cui il modello è incerto su continuazioni dipendenti dalla conoscenza (ad esempio, scegliere tra più fatti plausibili).
   • Combinazione: Il punteggio finale è una somma ponderata: $\phi_i = \alpha \cdot \bar{\Delta}^{unlearn}_i + (1-\alpha) \cdot \bar{H}_i$.

B. Strategie di Ottimizzazione Selettiva dei Token

Una volta generati i punteggi di importanza, il framework applica obiettivi di unlearning solo ai token ad alta importanza utilizzando due strategie:

1. Selezione Rigida:
   • Seleziona la frazione top-$r$ dei token (ad esempio, top 20%) basandosi su $\phi_i$.
   • Applica la perdita di unlearning esclusivamente a questi token; gli altri token ricevono zero aggiornamenti del gradiente.
2. Ponderazione Morbida:
   • Modula il contributo del gradiente di ogni token in base a un punteggio di importanza normalizzato (utilizzando un softmax con temperatura $\tau$).
   • Consente un'interpolazione fluida tra aggiornamenti uniformi e aggiornamenti completamente selettivi.

C. Integrazione con Metodi Esistenti

Il framework è agnostico rispetto al modello e si integra con obiettivi di unlearning basati su gradienti standard, inclusi:

• Gradient Ascent (GA)
• Weighted Gradient Ascent (WGA)
• Negative Preference Optimization (NPO)
• Representation Misdirection Unlearning (RMU)
• Ritenzione: Tutti i metodi sono combinati con regolarizzazione della divergenza KL sull'"insieme di ritenzione" per preservare le capacità generali.

3. Analisi Teorica

Il paper fornisce una giustificazione teorica del perché la selezione a livello di token migliora l'unlearning:

• Rapporto Segnale-Rumore (SNR): Gli autori modellano l'unlearning come movimento lungo un sottospazio a bassa dimensionalità ($U$) dei parametri.
• Riduzione del Rumore: Gli aggiornamenti uniformi accumulano rumore da tutti i $T$ token. La selezione a livello di token concentra il peso del gradiente sul sottoinsieme $K$ di token critici per la conoscenza.
• Risultato: L'SNR dello stimatore del gradiente migliora proporzionalmente all'inverso del rapporto di selezione ($T/|K|$). Sopprimendo gli aggiornamenti sui token non critici, il metodo riduce il rumore del gradiente e preserva più efficacemente le conoscenze non correlate.

4. Risultati Sperimentali

Gli autori hanno valutato TokenUnlearn su due benchmark attraverso tre architetture di modelli (Phi-1.5, Llama-2-7B, Qwen-3-8B).

Benchmark

1. TOFU (Task of Fictitious Unlearning): Unlearning di conoscenza fine-granularità (rimozione di fatti su autori fittizi).
2. WMDP (Weapons of Mass Destruction Proxy): Allineamento alla sicurezza (rimozione di conoscenze pericolose su bio/cybersecurity).

Risultati Chiave

• Migliore Dimenticanza: Le varianti a livello di token (T-* e S-*) hanno costantemente superato le baseline a livello di sequenza nell'efficacia della dimenticanza (minore accuratezza sugli insiemi di dimenticanza).
  • Esempio: Su TOFU con Qwen-3-8B, T-WGA ha migliorato significativamente l'efficacia della dimenticanza rispetto al WGA standard.
• Migliore Preservazione dell'Utilità: I metodi a livello di token hanno mantenuto un'utilità del modello più elevata (punteggi di ritenzione) su dati non mirati.
  • Esempio: Su Qwen-3-8B, il punteggio di ritenzione è migliorato fino al 19,0% rispetto alle baseline a livello di sequenza, ottenendo al contempo una migliore dimenticanza.
• Generalizzabilità: L'approccio ha funzionato coerentemente attraverso diverse dimensioni e architetture di modelli.
• Studi di Ablazione:
  • Rapporto di Selezione ($r$): Un rapporto del 20% (top 20% dei token) ha fornito il compromesso ottimale tra dimenticanza e ritenzione.
  • Peso dell'Entropia: Combinare segnali di entropia (circa 30% di peso) ha fornito benefici complementari, catturando punti decisionali incerti che il mascheramento da solo potrebbe perdere.

5. Contributi Chiave

1. Framework di Attribuzione a Livello di Token: Ha introdotto un metodo principiato per identificare i token critici per la conoscenza utilizzando mascheramento controfattuale e segnali di entropia senza richiedere modelli ausiliari.
2. Strategie di Unlearning Selettivo: Ha proposto strategie di "Selezione Rigida" e "Ponderazione Morbida" che estendono gli algoritmi di unlearning esistenti per operare a livello di token.
3. Insight Teorico: Ha dimostrato che gli aggiornamenti selettivi per token migliorano l'SNR del gradiente, spiegando perché gli aggiornamenti uniformi non riescono a ottimizzare i duali obiettivi di dimenticanza e ritenzione.
4. Validazione Empirica: Ha dimostrato miglioramenti coerenti su tre modelli e due benchmark distinti (conoscenza e sicurezza), ottenendo fino al 32,6% di miglioramento nella dimenticanza e al 19,0% di miglioramento nella preservazione dell'utilità.

6. Significato e Impatto

• Privacy e Regolamentazione: Abilita una conformità più precisa al "Diritto all'Oblio" (GDPR/CCPA) rimuovendo tracce di dati specifiche senza degradare l'utilità generale del modello.
• Sicurezza: Offre un meccanismo più efficace per rimuovere capacità pericolose (ad esempio, minacce alla biosicurezza) dai modelli pre-addestrati.
• Efficienza: Mette in discussione l'assunzione secondo cui l'unlearning richiede aggiornamenti di sequenza uniformi, suggerendo che la conoscenza negli LLM è sparsa e può essere presa di mira con precisione.
• Limitazioni: Il metodo introduce un sovraccarico computazionale dovuto al passaggio in avanti aggiuntivo per il mascheramento. È un metodo approssimato e non garantisce la rimozione al 100% delle tracce di conoscenza (possono esistere tracce residue).

In conclusione, TokenUnlearn rappresenta un cambiamento di paradigma nell'unlearning automatico, passando da aggiornamenti di sequenza a grana grossa ad attribuzione di token a grana fine, risolvendo così la tensione fondamentale tra dimenticanza efficace e preservazione dell'utilità.