Unlocking Prompt Infilling Capability for Diffusion Language Models

1. Il Problema: Il Divario Addestramento-Inferenza

I modelli di linguaggio a diffusione mascherata (dLMs, diffusion language models) sono architetturalmente progettati per generare testo attraverso un processo di denoising bidirezionale. A differenza dei modelli autoregressivi (che generano da sinistra a destra), i dLMs possono, in teoria, riempire qualsiasi posizione mascherata all'interno di una sequenza, condizionandosi al contesto circostante (sia a sinistra che a destra).

Tuttavia, il paper identifica un limite critico nelle pratiche attuali di Fine-Tuning Supervisionato (SFT):

• Mascheramento solo della risposta (Response-Only Masking): I modelli esistenti (come LLaDA e Dream) vengono addestrati mascherando esclusivamente i token della risposta, lasciando il prompt completamente visibile e non corrotto.
• La conseguenza: Sebbene l'architettura supporti il denoising del prompt, il modello non riceve mai il segnale di addestramento necessario per imparare a ricostruire o "riempire" (infill) parti di un prompt.
• Risultato: Durante l'inferenza, quando si tenta di usare il modello per generare prompt basati su esempi (prompt infilling), i modelli pubblici falliscono miseramente, spesso riempiendo le parti mascherate con token di fine testo (EOS) o generando contenuti privi di senso. Questo non è un limite architetturale, ma un limite derivante dalle pratiche di addestramento.

2. Metodologia: Mascheramento della Sequenza Completa

Per sbloccare questa capacità, gli autori propongono una modifica semplice ma fondamentale al processo di SFT:

• Full-Sequence Masking (FS): Invece di mascherare solo la risposta, durante l'addestramento viene mascherata l'intera sequenza $X = (P, R)$, dove $P$ è il prompt e $R$ è la risposta.

  • Per ogni esempio di addestramento, viene campionata una ratio di mascheramento $t$ e un numero casuale di token viene mascherato dall'intera sequenza.
  • Questo espone il modello a pattern di corruzione diversificati: solo prompt mascherati, solo risposte mascherate, o sequenze intere mascherate.
  • La funzione di perdita (loss) viene calcolata su tutti i token mascherati, permettendo al modello di imparare a denoisare sia il prompt che la risposta.

• Raffinamento Opzionale (FS+RO): Dopo la fase di mascheramento completo, gli autori introducono un secondo stadio opzionale di fine-tuning con mascheramento solo della risposta. Questo serve a specializzare il modello per la generazione standard, mantenendo però le capacità di infilling apprese nella fase precedente.

• Procedura di Inferenza (Prompt Infilling):

  1. Viene fornito un template di prompt con alcune regioni mascherate e pochi esempi (few-shot) con le risposte di riferimento.
  2. Il dLM denoisa i token mascherati nel template per generare un prompt completo e adattato al task.
  3. I prompt candidati generati vengono validati su tutti gli esempi few-shot per selezionare quello migliore.
  4. Il prompt ottimizzato viene poi utilizzato per l'inferenza su tutti i dati di test.

3. Contributi Chiave

1. Identificazione del Collo di Bottiglia: Dimostrano che l'incapacità dei dLMs di eseguire il prompt infilling è dovuta esclusivamente alla pratica di mascheramento solo della risposta durante l'SFT, non a limiti architetturali.
2. Nuova Strategia di Addestramento: Introducono il Full-Sequence Masking come ingrediente fondamentale per abilitare il denoising del prompt.
3. Protocollo di Infilling: Propongono un metodo pratico per generare prompt ottimizzati partendo da esempi, che supera o eguaglia i template progettati manualmente.
4. Complementarità: Dimostrano che i prompt generati tramite infilling sono complementari ai metodi esistenti di ottimizzazione automatica dei prompt (come GEPA o COPRO), fornendo guadagni aggiuntivi quando combinati.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli come LLaDA-8B e Dream-7B, valutati su dataset come GSM8K (matematica), SummEval (valutazione di riassunti) e BigGen-Bench.

• Sblocco della Capacità:

  • I checkpoint pubblici (addestrati con mascheramento solo risposta) mostrano un crollo delle prestazioni quando si tenta l'infilling (es. su GSM8K, l'accuratezza scende dal 73.1% al 64.8% per LLaDA).
  • Con il Full-Sequence Masking, il modello impara a generare prompt efficaci. Su SummEval, la correlazione di Spearman con il giudizio umano passa da ~0.08 (checkpoint pubblico) a 0.535 (modello FS+RO).
  • Il modello riesce a inferire rubriche di punteggio complesse (inclusi punteggi decimali non interi come 1.2 o 4.8) direttamente dagli esempi, migliorando l'allineamento con i criteri umani.

• Trasferibilità (Prompt Transfer):

  • I prompt infilliti da un modello addestrato con FS possono essere trasferiti ad altri modelli (anche con configurazioni di addestramento diverse) con successo.
  • Su GSM8K, un prompt infillito permette di raggiungere il 76.4% di accuratezza con solo 315 token, superando l'ICL (In-Context Learning) standard a 16-shot (73.1%) che richiede oltre 3300 token.

• Combinazione con Ottimizzazione Automatica:

  • L'infilling è stato applicato a prompt già ottimizzati da metodi come COPRO e GEPA.
  • La combinazione ha portato a miglioramenti significativi, specialmente su COPRO, riducendo la varianza dei risultati e aumentando la correlazione con il giudizio umano (es. su Prometheus 8x7B, la correlazione Spearman è salita da 0.447 a 0.490).

• Analisi della Perplexità:

  • L'analisi della Diffusion Perplexity conferma che i modelli addestrati con FS si adattano meglio ai template strutturati (Judge template), mostrando una perplexità significativamente inferiore rispetto ai modelli addestrati con mascheramento solo risposta.

5. Significato e Implicazioni

Questo lavoro ribalta una percezione comune secondo cui le capacità avanzate dei modelli di linguaggio richiedano necessariamente innovazioni architetturali complesse.

• L'importanza dei Dati e dell'Addestramento: Dimostra che molte capacità "nascoste" (come l'infilling del prompt) sono già presenti nell'architettura bidirezionale dei dLMs, ma rimangono inaccessibili se il segnale di addestramento non le espone esplicitamente.
• Efficienza e Automazione: Il metodo proposto permette di generare prompt altamente ottimizzati e compatti partendo da pochi esempi, riducendo la necessità di ingegneria manuale dei prompt e migliorando l'efficienza computazionale (meno token necessari per ottenere prestazioni superiori).
• Flessibilità: La capacità di infillire prompt apre nuove possibilità per l'adattamento dinamico dei modelli a nuovi task e formati di valutazione senza dover riaddestrare l'intero modello da zero.

In sintesi, il paper stabilisce che il Full-Sequence Masking durante l'SFT è la chiave per trasformare i modelli di diffusione da semplici generatori di testo a potenti strumenti di ottimizzazione e adattamento dei prompt, sbloccando il pieno potenziale della loro natura bidirezionale.