Ottenere accelerazioni senza perdite negli LLM grazie alla diffusione discreta

Abstract degli autori

I modelli linguistici di grandi dimensioni (LLM) devono gran parte del loro successo alla previsione del token successivo (NTP), ma la loro struttura autoregressiva (AR) richiede una generazione dei token lenta e sequenziale. Per superare questo collo di bottiglia, introduciamo gli LLM potenziati con la diffusione, una nuova classe di modelli che definisce la distribuzione di un modello AR, ma usa la diffusione per campionare in parallelo più token da tale distribuzione. Suddividiamo i parametri di questi modelli in due insiemi: i pesi AR, addestrati con il consueto obiettivo NTP, e i pesi di diffusione, leggeri e addestrati a generare più token contemporaneamente. I pesi di diffusione vengono appresi attraverso una semplice fase di Diffusion Distillation, che aggiunge un sovraccarico trascurabile alle pipeline esistenti per l’addestramento degli LLM. Introduciamo inoltre $Ψ$-Spec, una famiglia di metodi di campionamento che consente accelerazioni senza perdite e una maggiore scalabilità in fase di inferenza a lunghezza di contesto fissa. A differenza della decodifica speculativa, il nostro metodo non richiede un modello separato per generare bozze. A differenza degli LLM a diffusione (d-LLM), accelera la generazione senza sacrificare la qualità del modello AR sottostante. I modelli risultanti, chiamati Uno, possono essere addestrati da zero oppure ottenuti potenziando LLM AR esistenti con pesi aperti. Uno raggiunge un throughput superiore a quello dei principali metodi di decodifica speculativa per ogni dimensione del batch valutata e offre accelerazioni fino a $3\times$ rispetto al modello AR di base, anche alla massima dimensione del batch supportata dal dispositivo. In particolare, il nostro modello Uno 8B supera il principale d-LLM aperto, DiffusionGemma 26B, e il modello proprietario Mercury 2 in tutti i benchmark valutati relativi all’uso agentico degli strumenti, alla programmazione e al ragionamento su contesti lunghi. Rendiamo disponibili il codice e i checkpoint all’indirizzo: https://s-sahoo.github.io/uno/