DC-SAE: autoencoder semantico a compressione profonda per una convergenza più rapida della diffusione

Abstract degli autori

I tokenizer ad alta compressione sono essenziali per aumentare la scala dei modelli generativi di immagini nello spazio latente. Tuttavia, una compressione spinta comporta un compromesso fondamentale tra fedeltà della ricostruzione ed efficienza della generazione: un encoder di immagini ad alta compressione aumenta sempre la difficoltà dell’addestramento dei modelli di diffusione, rallentandone la convergenza. I recenti autoencoder basati su rappresentazioni accelerano l’addestramento dei modelli di diffusione sostituendo gli encoder VAE con encoder semantici preaddestrati, che migliorano la capacità espressiva delle caratteristiche latenti. In genere, però, si limitano a una compressione moderata e perdono i dettagli a livello di pixel necessari per una ricostruzione fedele. Per ottenere sia un’elevata compressione sia un addestramento rapido dei modelli di diffusione, proponiamo DC-SAE, un Decoupled Compact Semantic Autoencoder progettato per generare immagini ad alta compressione accelerando la convergenza dei modelli di diffusione. DC-SAE comprende due componenti principali: (1) un’architettura progettata a livello macroscopico che sfrutta encoder semantici per consentire rapporti di compressione più elevati e (2) un encoder a livello di pixel che conserva i dettagli di basso livello, garantendo una ricostruzione delle immagini ad alta fedeltà. Dimostriamo empiricamente che DC-SAE ottiene risultati solidi nei compiti di generazione di immagini, combinando rappresentazioni latenti compatte e un addestramento efficiente. Nello specifico, sul dataset ImageNet alla risoluzione di $512 \times 512$, DC-SAE raggiunge una compressione spaziale di $32\times$, con un PSNR di 29,79 e un gFID di 3,37: risultati superiori a quelli di DC-AE, il precedente tokenizer ad alta compressione di riferimento, rispettivamente del 13,5% per il PSNR e del 54,9% per il gFID, mantenendo un throughput comparabile e una convergenza più rapida durante l’addestramento dei modelli di diffusione. Oltre alla generazione condizionata dalla classe, un DiT con $1.6$B parametri che utilizza DC-SAE ottiene 0,84 su GenEval e 86,007 su DPG-Bench nella generazione di immagini da testo alla risoluzione di $1024\times1024$.