SpeeDiff: modello di diffusione latente end-to-end scalabile ancorato ai pixel

Abstract degli autori

Presentiamo Scalable Pixel-anchored End-to-end Diffusion (SpeeDiff), un metodo di diffusione latente che addestra congiuntamente il VAE e il modello di diffusione da zero. In linea di principio, l’addestramento congiunto permette al gradiente della funzione di perdita della diffusione di guidare direttamente l’encoder del VAE, favorendo la formazione di uno spazio latente adatto alla generazione e potenzialmente consentendo una convergenza più rapida rispetto al tradizionale approccio in due fasi, che impiega un VAE preaddestrato i cui parametri restano fissi. Un’implementazione end-to-end ingenua, tuttavia, peggiora drasticamente le prestazioni: la retropropagazione senza vincoli della perdita della diffusione provoca il collasso dello spazio latente. Il nostro principale contributo tecnico è una funzione di perdita semplice ma efficace, Tweedie Pixel Reconstruction (TPR), che fornisce un ulteriore segnale a livello di pixel decodificando, tramite la formula di Tweedie, una rappresentazione latente pulita prevista a partire da uno stato rumoroso intermedio, attenuando così il collasso. Il nostro metodo consente inoltre di aumentare congiuntamente la scala di un’architettura interamente basata su transformer e migliora l’allineamento delle rappresentazioni nell’ambito dell’approccio end-to-end. Il nostro modello SpeeDiff-XL raggiunge un addestramento oltre 140 volte più rapido rispetto a Vanilla SiT e 61 volte più rapido rispetto a REPA, ottenendo un FID di 1,50 senza guida nella generazione di immagini ImageNet a 256x256. Con un VAE più efficiente che applica una compressione 32x, il nostro modello raggiunge inoltre un FID di 1,53 senza guida nella generazione di immagini ImageNet a 512x512.