StepX-OPD: ripensare la distillazione on-policy con più insegnanti nei modelli di diffusione

Abstract degli autori

Un modello di diffusione text-to-image pronto per la distribuzione deve possedere contemporaneamente diverse capacità, come seguire le istruzioni, rendere il testo e rispecchiare le preferenze umane, e deve generare campioni in pochi passaggi. Le capacità e la generazione in pochi passaggi vengono acquisite in fasi separate: un modello a passaggi completi consolida prima le capacità tramite distillazione on-policy con più insegnanti, poi funge da insegnante fisso per la distillazione dei passaggi. Tutti i modelli a passaggi completi che valutiamo ottengono punteggi inferiori dopo la compressione, e la seconda fase non può iniziare prima che la prima sia terminata. Presentiamo StepX-OPD, che unifica le due fasi: un unico modello studente acquisisce tutte le capacità mentre impara a generare campioni in pochi passaggi, così gli specialisti congelati lo supervisionano sugli stati visitati dal suo stesso rollout, anziché su una traiettoria più lunga che non percorre mai. Poiché una scomposizione della KL del rollout assegna il peso maggiore alla discrepanza di velocità in prossimità dell’estremo pulito, regrediamo anche la velocità dello studente sulla parte del rollout a basso rumore. Separatamente, riscontriamo che l’integrazione delle capacità induce i modelli a inserire nelle immagini le parole dei prompt anche quando questi non richiedono testo, e proponiamo una variante che riduce questo fenomeno. Valutiamo il metodo su SD3.5-Medium e FLUX.1-dev in base alla capacità di seguire le istruzioni, rendere il testo, alla qualità delle immagini e alle preferenze umane. Su entrambe le architetture, StepX-OPD raggiunge il punteggio medio normalizzato più alto tra i metodi a 4 passaggi messi a confronto e supera il modello Flow-OPD a 40 passaggi; su SD3.5-Medium, inoltre, si addestra con circa il 37% di ore-GPU in meno rispetto alla distillazione sequenziale DiffusionOPD seguita dalla distillazione dei passaggi.