MMD nello spazio delle rappresentazioni per i modelli linguistici a diffusione

Abstract degli autori

Presentiamo un metodo di post-training per i modelli linguistici a diffusione (DLM) che riduce al minimo la Maximum Mean Discrepancy (MMD) tra le distribuzioni generate e quelle di riferimento nello spazio delle caratteristiche di un DLM preaddestrato i cui parametri sono congelati. Per stimare la MMD, conserviamo le caratteristiche contestuali nelle posizioni dei singoli token, ottenendo più osservazioni per sequenza con un solo passaggio dell’estrattore. Ottimizziamo questo obiettivo usando gradienti di policy per i modelli discreti e la differenziazione diretta attraverso le variabili latenti generate per quelli continui. In entrambi i casi, calcolare la funzione di perdita direttamente da queste caratteristiche consente un post-training efficiente senza traiettorie complete di campionamento né modelli ausiliari addestrati congiuntamente. Gli esperimenti mostrano una perplexity generativa inferiore a parità di entropia comparabile su OpenWebText e migliori compromessi tra accuratezza e costo computazionale su GSM8K. Nei modelli 16B DMax-LLaDA2.0 con diffusione ibrida mascherata-uniforme, aumentiamo il parallelismo nella decodifica mantenendo un’accuratezza simile o superiore nei benchmark di matematica e codice.