Pivot-SD: autodistillazione efficiente per modelli linguistici a diffusione con mascheramento

Abstract degli autori

I modelli linguistici a diffusione con mascheramento (dLM) offrono una promettente alternativa parallela ai modelli autoregressivi per il ragionamento complesso. Presentano però un problema specifico di attribuzione del merito: poche decisioni prese durante il denoising riducono drasticamente l’incertezza sulle restanti posizioni mascherate e determinano gran parte della risposta. La maggior parte delle procedure di post-training per i dLM non usa questo segnale per decidere su quali token addestrare il modello: in genere si addestra il modello sul testo finale o si assegnano ricompense a interi passaggi di denoising, anziché selezionare le singole decisioni che plasmano la risposta. Presentiamo Pivot-SD, un metodo efficiente di autodistillazione offline che supervisiona soltanto queste decisioni ad alto impatto (pivot). Pivot-SD seleziona i pivot mediante una metrica del guadagno di informazione che misura la riduzione dell’incertezza sulle restanti posizioni mascherate. I pivot delle traiettorie riuscite vengono addestrati con la cross-entropy, quelli delle traiettorie fallite con unlikelihood mirata, lasciando intatto il resto della traiettoria fallita. Usando soltanto 200 domande e quattro rollout per ciascuna, Pivot-SD migliora le prestazioni di LLaDA-8B-Instruct rispetto alle baseline di SFT sull’intera sequenza e di RL per modelli a diffusione con budget equivalente, nei benchmark di matematica e programmazione.