MEND: apprendimento per rinforzo per i modelli di flusso tramite matching prossimale delle velocità

Abstract degli autori

Il post-addestramento dei modelli di flusso basato sulle ricompense segue in genere una di due strade: ripesa i campioni generati dal modello stesso usando una penalità KL o un modello di riferimento fisso, spesso per migliaia di aggiornamenti; oppure propaga all’indietro il gradiente della ricompensa e sposta ogni campione senza verificare se il beneficio giustifichi l’entità dello spostamento. Presentiamo MEND, un metodo di apprendimento per rinforzo basato sul proximal velocity matching. MEND pone un tetto alle ricompense all’interno di ciascun gruppo di prompt, perciò i campioni che ottengono già un buon punteggio non vengono spostati. Al di sotto del tetto, propone spostamenti lungo il gradiente della ricompensa e ne accetta uno solo se l’aumento della ricompensa, calcolata con quel tetto, supera un costo quadratico dello spostamento. Il modello esegue quindi una regressione verso i target di velocità risultanti, senza un termine KL, un modello di riferimento fisso o pesi di vantaggio. In 100 aggiornamenti, MEND supera Flow-GRPO (circa 4k aggiornamenti) su cinque valutatori su sei, a parità di distanza dalle immagini del modello di base. Con un protocollo a parità di budget, supera ReFL e DiffusionNFT a ogni aggiornamento valutato per tutte e quattro le ricompense di addestramento, raggiungendo un PickScore di 24,03 contro, rispettivamente, 23,92 e 23,43. Anche un esperimento di 300 aggiornamenti con tre ricompense supera il modello DiffusionNFT con cinque ricompense su tutte e tre le ricompense con cui viene addestrato. MEND è generale e facile da adottare: si applica a qualsiasi backbone per modelli di flusso con una ricompensa differenziabile.