Technical Summary: SSDD: Single-Step Diffusion Decoder for Efficient Image Tokenization

Problem Statement

State-of-the-art generative image models rely on tokenizers to compress high-dimensional visual signals into compact latent representations. While current standard tokenizers, such as KL-regularized Variational Autoencoders (KL-VAE), are effective, they face a fundamental trade-off between distortion (pixel-level similarity) and distribution shift (perceptual realism and diversity). Deterministic decoders minimize distortion but often fail to model the complex distribution of realistic images, leading to blurry outputs. Conversely, generative decoders can model this distribution but typically require iterative sampling (slow inference) or rely on adversarial losses (GANs), which are difficult to scale and stabilize.

Existing diffusion-based decoders offer a principled approach to modeling the conditional distribution of images given latents but suffer from two main limitations:

1. Slow Inference: They require multiple denoising steps, creating a bottleneck for downstream generative models.
2. Training Instability: Achieving state-of-the-art results often necessitates GAN losses, which are hard to scale and can lead to diverging training runs.

Methodology

The authors introduce SSDD (Single-Step Diffusion Decoder), a novel architecture designed to overcome these limitations by combining a scalable hybrid architecture, GAN-free training, and a single-step distillation strategy.

1. Scalable Hybrid Architecture

Instead of relying solely on convolutional U-Nets (common in pixel-space diffusion) or pure Transformers (common in latent modeling), SSDD utilizes a U-ViT based architecture.

• Structure: It employs a U-Net backbone with four levels of convolutional ResNet blocks but replaces the middle attention block with a Transformer operating on tokens representing $8 \times 8$ pixel patches.
• Conditioning: The model conditions on the latent representation $z$ by upsampling the latent grid and concatenating it with the noisy input. It further enhances conditioning via Adaptive Group Normalization (AdaGN) and Adaptive Layer Normalization (AdaLN) within both ResNet and Transformer blocks.
• Scaling: The architecture scales efficiently from 13.4M (SSDD-S) to 153.8M (SSDD-XL) parameters by increasing channel widths and transformer depth.

2. GAN-Free Training Scheme

SSDD eliminates the need for adversarial training by combining three specific loss functions:

• Flow Matching Loss: The primary objective uses optimal transport flow matching to predict the velocity field ($v = x - \epsilon$) between noise and data, trained with an $L_2$ loss.
• Perceptual Loss (LPIPS): Added to guide generation toward perceptually correct areas, aligning high-level features.
• REPA Loss: A feature regularization loss using DINOv2-B features to stabilize the transformer representations and accelerate training.
• Multi-Stage Training: The model is trained in stages: first on $128 \times 128$ crops with random resizing augmentation, followed by fine-tuning at target resolutions ($256 \times 256$).

3. Single-Step Distillation

To address the inference speed bottleneck, SSDD employs a distillation strategy:

• Teacher-Student Setup: A multi-step teacher decoder (trained with 8 steps) generates high-quality, diverse reconstructions. A student decoder is then fine-tuned to reproduce these outputs in a single step.
• Loss Preservation: Unlike standard distillation that might rely solely on $L_2$ regression, the SSDD distillation preserves the full training objective. The student minimizes flow-matching and LPIPS losses computed against the teacher's multi-step outputs.
• Result: This allows the model to inherit the perceptual and generative properties of the multi-step teacher while operating with the throughput of a single-step model.

4. Shared Encoders

The authors demonstrate that SSDD decoders can be trained with a shared encoder (specifically the standard KL-regularized encoder from Rombach et al., 2022). This allows SSDD to serve as a drop-in replacement for existing VAE decoders without requiring the retraining of the encoder or the creation of a new latent space for each decoder size.

Key Contributions

1. First GAN-Free Single-Step Diffusion Decoder: SSDD is the first architecture to achieve state-of-the-art continuous tokenization without adversarial losses, ensuring training stability across model scales.
2. Hybrid Architecture: The integration of convolutional inductive biases with transformer generative capabilities in a pixel-space decoder yields superior scaling properties compared to pure U-Nets.
3. Efficient Distillation: The proposed distillation method successfully compresses multi-step diffusion behavior into a single step with minimal quality degradation, significantly improving throughput.
4. Shared Latent Space Compatibility: The ability to use standard, pre-trained encoders makes SSDD immediately applicable to existing generative pipelines.

Experimental Results

Evaluated on ImageNet-1k (256$\times$256) and COCO, SSDD demonstrates significant improvements over KL-VAE and other diffusion decoders:

• Reconstruction Quality (f8c4 setting):
  • SSDD-S (13.4M params) improves reconstruction FID (rFID) from 0.87 (KL-VAE) to 0.46.
  • SSDD-M (48M params, similar size to KL-VAE) achieves an rFID of 0.39.
  • SSDD outperforms all baselines in perceptual metrics (LPIPS, DreamSim) while maintaining competitive low-level metrics (PSNR, SSIM), though it prioritizes perceptual quality over pixel-perfect distortion.
• Throughput:
  • SSDD-S achieves 1.4$\times$ higher throughput than KL-VAE while offering better quality.
  • Compared to multi-step diffusion decoders (e.g., $\epsilon$-VAE), SSDD offers up to 10$\times$ faster decoding at similar rFID levels.
• Downstream Generation:
  • When used to train Diffusion Transformers (DiT), SSDD improves generation FID (gFID) compared to KL-VAE.
  • Using a higher compression ratio (f16c4), SSDD enables 3.8$\times$ faster sampling for the DiT model while preserving generation quality.
• Generalization: SSDD models generalize well to COCO and higher resolutions (up to 1024$\times$1024) without retraining from scratch, and can decode latents from frozen encoders like DiTo and DC-AE.

Significance

The paper claims that SSDD establishes a new standard for image tokenization by resolving the tension between speed, stability, and generative fidelity. By removing the reliance on GANs and iterative sampling, SSDD provides a scalable, drop-in replacement for KL-VAE that enables the training of higher-quality and faster generative models. The work demonstrates that diffusion-based tokenization can be both efficient and stable, paving the way for future large-scale generative modeling without the computational overhead of multi-step inference or the instability of adversarial training.