SSDD: Single-Step Diffusion Decoder for Efficient Image Tokenization

Tokenizers are a key component of state-of-the-art generativeimage models, extracting the most important features from the signalwhile reducing data dimension and redundancy. Most current imagetokenizers are based on KL-regularized variational autoencoders (KL-VAE), trained with reconstruction, perceptual and adversarial losses.Diffusion decoders have been proposed as a more principled alternativeto model the distribution over images conditioned on the latent. However,matching the performance of KL-VAE still required adversarial losses,as well as a higher decoding time due to iterative sampling. To addressthese limitations, we introduce a new pixel diffusion decoder architecturefor improved scaling and training stability, benefiting from transformercomponents and GAN-free training. We use distillation to replicate theperformance of the diffusion decoder in an efficient single-step decoder.This makes SSDD the first GAN-free single-step diffusion decoder toreach state-of-the-art continuous tokenization, with higher reconstructionquality and faster sampling than KL-VAE. In particular, SSDD improvesreconstruction FID from 0.87 to 0.46 with 1.4× higher throughput andpreserves generation quality of DiTs with 3.8× faster sampling. As such,SSDD can be used as a drop-in replacement for KL-VAE, and for buildinghigher-quality and faster generative models.