Replay What Matters: Off-Policy Replay for Efficient LLM Reinforcement Unlearning

Riepilogo Tecnico: Replay What Matters: Off-Policy Replay per un Efficiente Unlearning di LLM tramite Reinforcement Learning

Problematica

L'unlearning di Large Language Models (LLM) mira a rimuovere la conoscenza pericolosa (ad es., contenuti protetti da copyright, informazioni private) dai modelli pre-addestrati senza i costi proibitivi di un intero riaddestramento. Recenti approcci basati sul reinforcement learning (RL), come RULE, riformulano l'unlearning come l'apprendimento di un comportamento di rifiuto: insegnare al modello a rifiutare query relative a un "forget set" preservando al contempo le risposte su un "retain set" e un "boundary set".

Tuttavia, gli autori identificano un'inefficienza critica nello Stage 2 dell'ottimizzazione on-policy (specificamente utilizzando GRPO) impiegata da metodi come RULE. Durante l'addestramento, emerge una pronunciata disparità hard–easy:

1. Casi Easy (Facili): Molti prompt convergono rapidamente verso stati ad alto reward. Continuare a campionare questi prompt nelle fasi successive dell'addestramento produce un segnale di gradiente poco informativo (vantaggio $\approx 0$) ma consuma ingenti risorse computazionali.
2. Casi Hard (Difficili): I prompt vicini al confine semantico tra dimenticare e mantenere spesso rimangono a basso reward durante tutto l'addestramento. Nelle impostazioni on-policy standard, le rollout a basso reward generate per questi casi hard vengono scartate dopo un singolo utilizzo, nonostante contengano informazioni preziose sulle modalità di fallimento del modello.

Questa dinamica porta a una sottottimale efficienza dell'addestramento, dove la computazione viene sprecata su casi easy già convergenti, mentre i casi hard rimangono sotto-ottimizzati, potenzialmente degradando la capacità del modello di preservare l'utilità generale (Retain Quality) pur dimenticando obiettivi specifici.

Metodologia: ReRULE

Per affrontare questo problema, gli autori propongono ReRULE (Replay-enhanced Reinforcement UnLEarning), un meccanismo di potenziamento off-policy replay integrato nel framework RULE.

Panoramica del Framework

ReRULE mantiene la struttura a due stadi di RULE:

• Stage 1 (Rejection Steering): Supervised Fine-Tuning (SFT) su un dataset di rifiuto per dare un warm-start alla capacità del modello di rifiutare.
• Stage 2 (Refusal Boundary Optimization): Ottimizzazione basata su RL per raffinare il confine tra query di dimenticanza e di mantenimento.

La Pipeline ReRULE

L'innovazione principale avviene nello Stage 2, dove ReRULE introduce un replay buffer e un addestramento ibrido:

1. Identificazione e Archiviazione dei Casi Hard:
   Durante l'addestramento GRPO iniziale, il sistema monitora il reward medio dei gruppi di rollout generati per ogni prompt. Se il reward medio di un gruppo scende al di sotto di una soglia prestabilita $\tau$, esso viene classificato come un "caso hard" e memorizzato in un replay buffer $\mathcal{B}$. Questo buffer accumula traiettorie a basso reward che rappresentano le attuali modalità di fallimento del modello.

2. Addestramento Ibrido (Aggiornamenti Off-Policy):
   Una volta che il buffer contiene un numero sufficiente di casi hard ($|\mathcal{B}| \ge B_{min}$) e un periodo di warm-up ($t \ge T_{warm}$) è completato, ReRULE attiva un regime di addestramento ibrido. Ad ogni step, esso esegue:

   • Aggiornamento On-Policy Standard: Un aggiornamento GRPO standard utilizzando campioni freschi dalla policy corrente $\pi_\theta$.
   • Aggiornamento di Replay Off-Policy: Un ulteriore aggiornamento del gradiente utilizzando campioni tratti dal replay buffer $\mathcal{B}$.

3. Importance Sampling:
   Per correggere lo shift della distribuzione tra la policy che ha generato le rollout archiviate ($\pi_{\theta_{old}}$) e la policy corrente ($\pi_\theta$), ReRULE utilizza l'importance sampling. Il gradiente off-policy è calcolato come:
   $$\nabla_\theta J_{off}(\theta) = \mathbb{E}_{g \sim \mathcal{B}} \left[ \mathbb{E}_{y \sim \pi_{\theta_{old}}} \left[ \tilde{\rho}(y|x) \nabla_\theta \log \pi_\theta(y|x) \hat{A}(x, y) \right] \right]$$
   dove $\tilde{\rho}(y|x)$ è un peso di importanza tagliato (clipped) e normalizzato. Ciò consente al modello di ri-addestrarsi sui casi difficili senza richiedere nuovi costosi campionamenti.

Contributi Chiave

1. Identificazione dell'Inefficienza dell'Addestramento: Il paper dimostra empiricamente che l'on-policy GRPO nell'unlearning di LLM soffre di rendimenti decrescenti sui casi easy e di un sotto-utilizzo dei casi hard, portando a una "disparità hard/easy" che ostacola la convergenza.
2. Proposta di ReRULE: Gli autori introducono un meccanismo di replay off-policy che memorizza le rollout dei casi hard a basso reward e le riutilizza tramite aggiornamenti con importance sampling. Questo reindirizza il focus computazionale verso i casi di confine che richiedono un'ulteriore ottimizzazione.
3. Analisi Teorica: Gli autori forniscono una prova teorica che mostra come ReRULE produca un limite di convergenza più stretto per i casi hard rispetto al puro RULE on-policy. Aumentando il numero effettivo di token di addestramento nella regione hard ($N_{H}^{ReRULE} = N_{H}^{on} + \eta_\rho N_{H}^{off}$), il metodo riduce l'errore di esplorazione nella regione hard.
4. Validazione Empirica: Esperimenti estesi su tre benchmark (RWKU, MUSE-Books, TOFU) validano l'efficacia del metodo.

Risultati Sperimentali

Gli autori hanno valutato ReRULE contro i baseline inclusi GA, NPO, SimNPO e l'originale RULE su tre dataset:

• MUSE-Books (Unlearning di Contenuti Protetti da Copyright): Questa impostazione, che coinvolge contenuti di Harry Potter, presenta un'alta proporzione di casi hard. ReRULE ha migliorato significativamente la Retain Quality dal 46.3 (RULE) al 56.2, mantenendo al contempo una Forget Quality di alto livello. Il tempo di addestramento è aumentato solo del 5–11% (da 6.8h a 7.1h).
• RWKU (Conoscenza del Mondo Reale): ReRULE ha ottenuto la migliore Forget Quality e ha mantenuto una performance top-2 in Forget Naturalness e Retain Quality, superando RULE in efficienza sui casi hard.
• TOFU (Unlearning di Entità): In questa impostazione, il forget set è più piccolo e generalmente più facile (meno casi hard). ReRULE ha mostrato miglioramenti limitati rispetto a RULE, il che gli autori citano come prova che il meccanismo di replay è più vantaggioso quando esiste una pronunciata disparità hard/easy.
• Studi di Ablazione:
  • Sostituire la selezione dei casi hard con lo stoccaggio casuale (random storage) ha causato un degrado delle prestazioni ai livelli del RULE standard, confermando che il targeting dei casi hard è cruciale.
  • Variando la soglia $\tau$, si è osservato che impostarla troppo alta (includendo casi più facili) riduce l'efficienza, mentre l'impostazione ottimale (0.4 per MUSE) bilancia selettività e dimensione del buffer.

Significato e Rivendicazioni

Il paper sostiene che ReRULE offra un potenziamento cost-effective all'unlearning basato su RL. Riciclando le rollout "fallite" dei casi hard, migliora la Retain Quality del modello sottoposto a unlearning — un parametro critico spesso compromesso nei compiti di unlearning — senza richiedere un intero riaddestramento o un ulteriore calcolo significativo.

Gli autori sottolineano che il metodo è condizionale: i suoi benefici sono più pronunciati in scenari con una significativa disparità tra prompt easy e hard (ad es., MUSE-Books). In impostazioni più semplici (ad es., TOFU), i guadagni sono marginali, a supporto dell'ipotesi che il meccanismo di replay colpisca specificamente l'inefficienza dei casi hard sotto-ottimizzati. Il lavoro posiziona ReRULE come un passo pratico verso un machine unlearning più efficiente e robusto, in particolare per preservare l'utilità generale del modello pur rimuovendo specifiche conoscenze pericolose.