Knowledge Editing in Masked Diffusion Language Models

Sintesi Tecnica: Editing della Conoscenza nei Modelli di Linguaggio a Diffusione Mascherata

Problematica
L'editing della conoscenza mira ad aggiornare o correggere conoscenze fattuali specifiche all'interno dei Large Language Models (LLM) senza riaddestrare l'intero modello. Il paradigma dominante, "locate-then-edit" (es. ROME, MEMIT), è stato sviluppato esclusivamente per i Modelli Autoregressivi (ARM). Questi metodi si basano sul causal tracing per identificare che il richiamo fattuale è localizzato nelle MLP (Multi-Layer Perceptrons) degli strati iniziali-intermedi al livello dell'ultimo token del soggetto. Successivamente, modificano i pesi in questa posizione per associare un nuovo oggetto al soggetto.

Tuttavia, i Modelli di Diffusione Mascherata (MDM) rappresentano un paradigma distinto. A differenza degli ARM, che predicono i token sequenzialmente da sinistra a destra, gli MDM modellano il testo in modo bidirezionale e generano testo attraverso un processo iterativo di denoising, recuperando i token mascherati da uno stato completamente mascherato. Rimane una questione aperta se le assunzioni alla base del "locate-then-edit" — specificamente la localizzazione dei fatti e l'efficacia dell'editing in tale posizione — si applichino agli MDM, dato il loro processo di generazione e i loro obiettivi di addestramento fondamentalmente diversi.

Metodologia
Gli autori trasferiscono il framework "locate-then-edit" (specificamente l'algoritmo MEMIT) agli MDM senza alterarne l'algoritmo centrale, applicando solo minimi aggiustamenti per adattarlo al meccanismo di predizione mascherata degli MDM.

1. Configurazione Sperimentale: Lo studio confronta due coppie accoppiate di modelli:

   • ARM: LLaMA-3-8B e Qwen2.5-7B.
   • MDM: LLaDA-8B (addestrato da zero) e Dream-v0-7B (inizializzato dai pesi di Qwen2.5).
   • Dataset: COUNTERFACT (benchmark standard, principalmente con target a singolo token) e KAMEL (adattato per l'analisi di target a multi-token).

2. Causal Tracing: Gli autori eseguono il causal tracing per determinare dove i fatti sono causalmente localizzati negli MDM. Corrompono gli embedding dei token del soggetto e misurano il recupero delle predizioni corrette ripristinando gli stati nascosti (hidden states) in specifici livelli e posizioni dei token.

3. Valutazione dell'Editing: Misurano le prestazioni di editing (efficacia, generalizzazione, specificità) attraverso diversi posizionamenti dei token e intervalli di strati. Fondamentalmente, valutano le prestazioni attraverso target di lunghezza variabile (singolo token vs multi-token) e diverse impostazioni di generazione per identificare i breakdown specifici del processo di generazione degli MDM.

4. Correzione Proposta: Dopo aver identificato un degrado delle prestazioni negli MDM per i target multi-token, gli autori propongono una correzione. Lo standard edit ottimizza l'aggiornamento del peso ($\delta$) per un singolo stato di condizionamento (lo stato del target completamente mascherato). La correzione proposta amplia l'ottimizzazione includendo stati intermedi parzialmente de-mascherati che il modello attraversa durante il denoising iterativo.

Risultati Chiave

1. Localizzazione Condivisa: Il nucleo dell'assunzione del "locate-then-edit" si trasferisce tra i paradigmi. Nonostante l'addestramento bidirezionale, gli MDM localizzano il richiamo fattuale nello stesso sito degli ARM: la MLP dei primi-medi strati all'ultimo token del soggetto. Il causal tracing conferma che il "sito precoce" (richiamo del fatto) e il "sito tardivo" (readout della predizione) esistono in entrambi i tipi di modello.

2. Esiti Divergenti (Il Breakdown del Multi-Token): Sebbene la posizione sia condivisa, l'esito non lo è.

   • Target a Singolo Token: L'editing performa in modo comparabile tra ARM e MDM.
   • Target Multi-Token: Le prestazioni di editing negli MDM degradano sistematicamente all'aumentare della lunghezza del target. Ad esempio, su target a 4 token, l'efficacia di LLaDA scende da 0.92 (1-token) a ~0.27, mentre LLaMA rimane elevata (0.89).
   • Causa Radice: Il fallimento deriva da un mismatch tra lo stato di ottimizzazione e il processo di generazione. MEMIT ottimizza l'edit per lo stato completamente mascherato. Negli ARM, la generazione procede da sinistra a destra, corrispondendo a questo stato. Negli MDM, la generazione procede attraverso una catena di stati intermedi dove alcuni token sono de-mascherati e altri rimangono mascherati. L'edit, ottimizzato solo per lo stato completamente mascherato, non tiene conto del contesto variabile creato da questi token intermedi de-mascherati.

3. Restaurazione tramite Ottimizzazione dello Stato Intermedio: Modificando l'obiettivo di ottimizzazione per adattare il residuo $\delta$ attraverso tutti gli stati intermedi parzialmente de-mascherati (ciclando attraverso il numero di token rivelati), gli autori ripristinano sostanzialmente le prestazioni di editing multi-token. Per LLaDA, l'efficacia su target a 4 token sale da 0.27 a 0.73.

Significatività e Rivendicazioni
Il paper sostiene di fornire la prima analisi del "locate-then-edit" per la conoscenza negli MDM. I suoi principali contributi sono:

• Validazione della Localizzazione: Dimostrare che la localizzazione causale dei fatti nella MLP del token del soggetto è una proprietà indipendente dal paradigma, valida sia per i modelli linguistici autoregressivi che per quelli a diffusione.
• Identificazione di un Fallimento Indotto dalla Generazione: Rivelare che i benchmark standard (dominati da target a singolo token) mascherano un fallimento sistematico del "locate-then-edit" negli MDM per i fatti multi-token. Questo fallimento è attribuito al mismatch tra l'ottimizzazione a singolo stato e il denoising iterativo.
• Una Correzione Semplice: Proporre un metodo per ottimizzare gli edit per gli stati intermedi, che colma efficacemente il divario di prestazioni per l'editing multi-token negli MDM senza cambiare la posizione dell'edit o la struttura dell'algoritmo sottostante.

Gli autori concludono che il porting di metodi di editing tra architetture richiede di far corrispondere non solo il sito dell'edit, ma anche il processo di generazione attraverso il quale il fatto editato viene realizzato. Notano che, sebbene la loro correzione affronti la causa principale del fallimento, rimane un divario di prestazioni per i target più lunghi, suggerendo che il lavoro futuro potrebbe dover tenere conto dell'ordine specifico con cui gli MDM de-mascherano i token.