Epsilon-VAE: Denoising as Visual Decoding

TL;DR

Epsilon-VAE rethinks visual tokenization by substituting the traditional deterministic decoder with a conditional diffusion-based denoiser that iteratively refines encoded latents. Through a carefully designed loss suite and a rectified-flow diffusion parameterization, it achieves superior reconstruction fidelity and enhances downstream generation, with notable speedups at higher compression. The approach demonstrates robust resolution generalization and strong gains in class-conditional generation when paired with latent diffusion models. This work offers a pathway to integrate iterative generation into autoencoding, improving both compression efficiency and perceptual quality.

Abstract

In generative modeling, tokenization simplifies complex data into compact, structured representations, creating a more efficient, learnable space. For high-dimensional visual data, it reduces redundancy and emphasizes key features for high-quality generation. Current visual tokenization methods rely on a traditional autoencoder framework, where the encoder compresses data into latent representations, and the decoder reconstructs the original input. In this work, we offer a new perspective by proposing denoising as decoding, shifting from single-step reconstruction to iterative refinement. Specifically, we replace the decoder with a diffusion process that iteratively refines noise to recover the original image, guided by the latents provided by the encoder. We evaluate our approach by assessing both reconstruction (rFID) and generation quality (FID), comparing it to state-of-the-art autoencoding approaches. By adopting iterative reconstruction through diffusion, our autoencoder, namely Epsilon-VAE, achieves high reconstruction quality, which in turn enhances downstream generation quality by 22% at the same compression rates or provides 2.3x inference speedup through increasing compression rates. We hope this work offers new insights into integrating iterative generation and autoencoding for improved compression and generation.

Figures (7)

• Figure 1: An overview of $\epsilon$-VAE. We frame visual decoding as an iterative denoising problem by replacing the autoencoder decoder with a diffusion model, optimized using a combination of score, perception, and trajectory matching losses. During inference, images are reconstructed (or generated) from encoded (or sampled) latents through an iterative denoising process. The number of sampling steps $N$ can be flexibly adjusted within small NFE regimes (from 1 to 3). We empirically confirm that $\epsilon$-VAE significantly outperforms the standard VAE schema, even with just a few steps.
• Figure 2: Architecture and compression analysis. The $\epsilon$-VAE decoder uses either a UNet-based ADM or Transformer-based DiT (left). $\epsilon$-VAE and VAE are evaluated under different compression rates by varying latent channel dimensions (middle) or encoder downsampling factors (right). We follow the $\epsilon$-VAE-lite configuration in these experiments.
• Figure 3: Impact of our major diffusion decoder designs. Improved training objectives, particularly perceptual matching loss and adversarial denoising trajectory matching loss, significantly contribute to better rFID scores and NFE (left). Effective noise scheduling by modulating the scaling factor $\gamma$ further enhances rFID, with an optimum value of 0.6 in our experiments (middle). Lastly, adjusting time step spacing during inference ensures stable sampling in low NFE regimes (right). We report results under the $\epsilon$-VAE-lite configuration.
• Figure 4: Image reconstruction results under the SD-VAE configuration rombach2022high at the resolution of $512 \times 512$. We find that $\epsilon$-VAE produces more accurate visual details than SD-VAE in the highlighted regions with text or human face. Best viewed when zoomed-in and in color.
• Figure 5: Reconstruction results with varying downsampling ratios.$\epsilon$-VAE maintains both high fidelity and perceptual quality, even under extreme downsampling conditions, whereas VAE fails to preserve semantic integrity. Best viewed when zoomed-in and in color.