Semplificare e migliorare la distillazione on-policy per la diffusione in pochi passaggi

Abstract degli autori

La distillazione in pochi passaggi riduce il costo del campionamento dei modelli di diffusione e di flusso, ma gli approcci esistenti richiedono spesso modelli ausiliari per la stima dello score o architetture specializzate. Mostriamo che una semplice regressione on-policy della velocità offre già una solida base per la generazione a 4-NFE. Per rappresentare le variazioni della velocità del modello insegnante all’interno di ogni transizione a grana grossa, introduciamo Pivot, un proxy residuo usato solo durante l’addestramento che separa le correzioni locali dalla previsione dello studente della velocità media sull’intervallo. Inoltre, allontaniamo lo studente dalle previsioni meno valide del modello insegnante, offrendo un meccanismo leggero per evitare modi indesiderati. Il framework risultante, Simplified Matching Distillation (SMD), si basa interamente sulla regressione, non richiede né la stima online dello score né l’addestramento avversario e lascia invariato il generatore impiegato. Valutiamo SMD nella generazione e modifica di immagini, nella generazione 3D, nella sintesi video e nella generazione congiunta di audio e video, scalando fino a 33 miliardi di parametri. A 4-NFE, SMD supera solidi metodi di riferimento per la generazione in pochi passaggi e, su diversi benchmark, supera i modelli insegnanti a più passaggi. Supporta inoltre la distillazione tra modelli, con più modelli insegnanti e per video causali.