DA-VAE: Plug-in Latent Compression for Diffusion via Detail Alignment

摘要

Reducing the token count is crucial for efficient training and inference of latent diffusion models, especially at high resolution. A common approach is to build high-compression image tokenizers that store more information by allocating more channels per token. However, when trained solely with reconstruction objectives, high-dimensional latent spaces often fail to maintain meaningful structure, which complicates diffusion training. Existing methods introduce additional objectives, such as semantic alignment or selective dropout, to enforce structure, but typically require costly retraining of the diffusion model. In contrast, pretrained diffusion models already exhibit a structured, lower-dimensional latent space, suggesting a simpler strategy: expand the latent dimensionality while preserving this structure. To this end, we propose Detail-Aligned VAE (DA-VAE), which increases the compression ratio of a pretrained VAE with only lightweight adaptation of the diffusion backbone. DA-VAE imposes an explicit latent layout: the first C channels are inherited from the pretrained VAE at a base resolution, while additional D channels encode high-resolution details. We introduce a simple yet effective detail-alignment mechanism that explicitly regularizes the expanded latent space to preserve the structural properties of the original space defined by the first C channels. Finally, we present a warm-start fine-tuning strategy that enables 1024 x1024 image generation with Stable Diffusion 3.5 using only 32 x32 tokens, 4xfewer than the original model, within a compute budget of 5 H100-days. It further enables 2048 x2048 generation with SD3.5, achieving a 6xspeedup while preserving image quality. We also validate the method and its design choices quantitatively on ImageNet.