Z-Cache: accelerare i transformer a diffusione tramite autoriflessione

Abstract degli autori

I transformer a diffusione sono diventati i modelli più potenti per la generazione visiva, ma richiedono ancora enormi risorse di calcolo. Per risolvere questo problema, è stato proposto il caching delle feature: le feature dei modelli a diffusione vengono memorizzate nei passaggi di calcolo precedenti e poi riutilizzate nei successivi passaggi di caching. Questo approccio consente un’accelerazione significativa, ma comporta anche un peggioramento della qualità della generazione. Per affrontare il problema, questo paper propone Z-cache, un metodo di caching delle feature che, grazie all’autoriflessione, consente di mantenere un’elevata qualità della generazione. In particolare, osserviamo che l’errore dovuto al caching delle feature tende a ridursi nettamente dopo ogni calcolo completo. Sulla base di questa osservazione, Z-Cache è progettato per prevedere prima le feature nei successivi passaggi di caching e poi eseguire un calcolo completo. In seguito, Z-Cache riprende i passaggi di caching e li riprevede sulla base dei passaggi di calcolo precedenti e di quello corrente, correggendo così le feature. Gli esperimenti dimostrano che, con Z-Cache, i transformer a diffusione raggiungono una qualità di generazione paragonabile a quella del modello originale, ma con una velocità di inferenza maggiore: per esempio, un’accelerazione di 6,22X su FLUX-dev per la generazione di immagini a partire da testo. Il nostro codice è disponibile su GitHub.