$\bm{\epsilon}$-VAE: Denoising as Visual Decoding

Generative modeling aims to capture the underlying distribution of training data, enabling realistic sample generation during inference. A key preprocessing step is tokenization, which converts raw data into discrete tokens or continuous latent representations. These compact representations allow models to efficiently learn complex patterns, enhancing the quality of generated outputs.

Two dominant paradigms in modern generative modeling are autoregression (Radford et al., 2018) and diffusion (Ho et al., 2020). Autoregressive models generate samples incrementally, conditioning each new element on previously generated ones, while diffusion models progressively refine noisy inputs into clear representations, without relying on element-wise factorization. Tokenization is an essential in both: discrete tokens direct step-by-step conditional generation in autoregressive models, while continuous latents streamline the denoising process in diffusion models. Empirical results across language (Achiam et al., 2023; Anil et al., 2023; Dubey et al., 2024) and vision (Baldridge et al., 2024; Esser et al., 2024; Brooks et al., 2024) tasks show that tokenization—whether discrete or continuous—improves generative performance. We focus on tokenization for latent diffusion models, which excel at generating high-dimensional visual data.

Given its central role in both paradigms, understanding how tokenization works is essential. In language processing, tokenization is relatively straightforward, involving segmenting text into discrete units such as words, subwords, or characters (Sennrich et al., 2015; Kudo & Richardson, 2018; Kudo, 2018). However, tokenization in visual domains poses greater challenges due to the continuous, high-dimensional, and redundant nature. Instead of direct segmentation, compact representations are typically learned using an autoencoding (Hinton & Salakhutdinov, 2006). Despite rapid advancements in visual generation techniques, the design of tokenizers has received relatively little attention. This is evident in the minimal evolution of tokenizers used in state-of-the-art models, which have remained largely unchanged since their initial introduction (Van Den Oord et al., 2017).

In this paper, we address this gap by revisiting the widely adopted visual autoencoding formulation (Esser et al., 2021), aiming to achieve higher compression rates and improved reconstruction quality, thereby enhancing generation quality of downstream generative models. Our key idea is to rethink the traditional autoencoding pipeline, which typically involves an encoder that compresses the input into a latent representation, followed by a decoder that reconstructs the original data in a single step. In our approach, we replace the deterministic decoder with a diffusion process. Here, the encoder still compresses the input into a latent representation, but instead of a one-step reconstruction, the diffusion model iteratively denoises the data to recover the original. This reframing turns the reconstruction phase into a step-by-step refinement, where the diffusion model, guided by the latent representation, progressively restores the original data.

To implement our approach effectively, several key design factors must be carefully considered. First, the architectural design must ensure effective conditioning of the diffusion decoder on the latent representations provided by the encoder. Second, the objectives for training the diffusion decoder should also explore potential synergies with traditional autoencoding losses, such as LPIPS (Zhang et al., 2018) and GAN (Esser et al., 2021). Finally, diffusion-specific design choices are crucial, including: (1) the model parameterization, which defines the prediction target for the diffusion decoder; (2) the noise schedule, which dictates the optimization trajectory; and (3) the distribution of time steps during training and testing, which balances noise levels during learning and generation. Our study systematically explores all these components under controlled experiments.

In summary, our contributions are: (1) introducing a novel approach that reframes the autoencoding pipeline using diffusion decoder; (2) presenting key design choices for optimal performance; and (3) conducting extensive controlled experiments demonstrating that our method achieves high-quality reconstruction and generation results compared to leading visual auto-encoding paradigms.

We start by briefly reviewing the basic concepts required to understand the proposed method. A more detailed summary of related work is deferred to Appendix A.

Visual tokenization. To achieve efficient and scalable high-resolution image synthesis, common generative models, including autoregressive models (Razavi et al., 2019; Esser et al., 2021; Chang et al., 2022) and diffusion models (Rombach et al., 2022), are typically trained in a low-resolution latent space by first downsampling the input image using a tokenizer. The tokenizer is generally implemented as a convolutional autoencoder consisting of an encoder, ${\mathcal{E}}$, and a decoder, ${\mathcal{G}}$. Specifically, the encoder, ${\mathcal{E}}$, compresses an input image ${\bm{x}}\in\mathbb{R}^{H\times W\times 3}$ into a set of latent codes (i.e., tokens), ${\mathcal{E}}({\bm{x}})={\bm{z}}\in\mathbb{R}^{H/f\times W/f\times n_{z}}$, where $f$ is the downsampling factor and $n_{z}$ is the latent channel dimensions. The decoder, ${\mathcal{G}}$, then reconstructs the input from ${\bm{z}}$, such that ${\mathcal{G}}({\bm{z}})={\bm{x}}$.

Training an autoencoder primarily involves several losses: reconstruction loss ${\mathcal{L}}_{\text{rec}}$, perceptual loss (LPIPS) ${\mathcal{L}}_{\text{LPIPS}}$, and adversarial loss ${\mathcal{L}}_{\text{adv}}$. The reconstruction loss minimizes pixel differences (i.e., typically measured by the $\ell_{1}$ or $\ell_{2}$ distance) between ${\bm{x}}$ and ${\mathcal{G}}({\bm{z}})$. The LPIPS loss (Zhang et al., 2018) enforces high-level structural similarities between inputs and reconstructions by minimizing differences in their intermediate features extracted from a pre-trained VGG network (Simonyan & Zisserman, 2015). The adversarial loss (Esser et al., 2021) introduces a discriminator, ${\mathcal{D}}$, which encourages more photorealistic outputs by distinguishing between real images, ${\mathcal{D}}({\bm{x}})$, and reconstructions, ${\mathcal{D}}({\mathcal{G}}({\bm{z}}))$. The final training objective is a weighted combination of these losses:

$${\mathcal{L}}_{\text{VAE}}={\mathcal{L}}_{\text{rec}}+\lambda_{\text{LPIPS}}\cdot{\mathcal{L}}_{\text{LPIPS}}+\lambda_{\text{adv}}\cdot{\mathcal{L}}_{\text{adv}},$$

(1)

where the $\lambda$ values are weighting coefficients. In this paper, we consider the autoencoder optimized by Eq. 1 as our main competing baseline (Esser et al., 2021), as it has become a standard tokenizer training scheme widely adopted in state-of-the-art image and video generative models (Chang et al., 2022; Rombach et al., 2022; Yu et al., 2022; 2023; Kondratyuk et al., 2024; Esser et al., 2024).

Diffusion. Given a data distribution $p_{{\bm{x}}}$ and a noise distribution $p_{\bm{\epsilon}}$, a diffusion process progressively corrupts clean data ${\bm{x}}_{0}\sim p_{{\bm{x}}}$ by adding noise $\bm{\epsilon}\sim p_{\bm{\epsilon}}$ and then reverses this corruption to recover the original data (Song & Ermon, 2019; Ho et al., 2020), represented as:

$${\bm{x}}_{t}=\alpha_{t}\cdot{\bm{x}}_{0}+\sigma_{t}\cdot\bm{\epsilon},$$

(2)

where $t\in[0,\text{T}]$ and $\bm{\epsilon}$ is drawn from a standard Gaussian distribution, $p_{\bm{\epsilon}}=\mathcal{N}(0,I)$. The functions $\alpha_{t}$ and $\sigma_{t}$ govern the trajectory between clean data and noise, affecting both training and sampling. The basic parameterization in Ho et al. (2020) defines $\sigma_{t}=\sqrt{1-a_{t}^{2}}$ with $a_{t}=\left(\prod_{s=0}^{t}(1-\beta_{s})\right)^{\frac{1}{2}}$ for discrete timesteps. The diffusion coefficients $\beta_{t}$ are lineary interpolated values bewteen $\beta_{0}$ and $\beta_{T-1}$ as $\beta_{t}=\beta_{0}+\frac{t}{T-1}(\beta_{T-1}-\beta_{0})$, with start and end values are set empirically.

The forward and reverse diffusion processes are described by the following factorizations:

$$q({\bm{x}}_{\Delta t:\text{T}}|{\bm{x}}_{0})=\prod_{i=1}^{\text{T}}q({\bm{x}}_{i\cdot\Delta t}|{\bm{x}}_{(i-1)\cdot\Delta t})\;\;\textrm{and}\;\;p({\bm{x}}_{0:\text{T}})=p({\bm{x}}_{\text{T}})\prod_{i=1}^{\text{T}}p({\bm{x}}_{(i-1)\cdot\Delta t}|{\bm{x}}_{i\cdot\Delta t}),$$

(3)

where the forward process $q({\bm{x}}_{\Delta t:\text{T}}|{\bm{x}}_{0})$ transitions clean data ${\bm{x}}_{0}$ to noise ${\bm{x}}_{\text{T}}=\bm{\epsilon}$, while the reverse process $p({\bm{x}}_{0:\text{T}})$ recovers clean data from noise. $\Delta t$ denotes the time step interval or step size.

During training, the model learns the score function $\nabla\log p_{t}({\bm{x}})\propto-\frac{\epsilon}{\sigma_{t}}$, which represents gradient pointing toward the data distribution along the noise-to-data trajectory. In practice, the model $s_{\Theta}({\bm{x}}_{t},t)$ is optimized by minimizing the score-matching objective:

$${\mathcal{L}}_{\text{score}}=\min_{\Theta}\mathbb{E}_{t\sim\pi(t),\epsilon\sim\mathcal{N}(0,I)}\left[w_{t}\|\sigma_{t}s_{\Theta}({\bm{x}}_{t},t)+\bm{\epsilon}\|^{2}\right],$$

(4)

where $\pi(t)$ defines the time-step sampling distribution and $w_{t}$ is a time-dependent weight. These elements together influence which time steps or noise levels are prioritized during training.

Conceptually, the diffusion model learns the tangent of the trajectory at each point along the path. During sampling, it progressively recovers clean data from noise based on its predictions.

Rectified flow. Rectified flow provides a specific parametrization of $\alpha_{t}$ and $\sigma_{t}$ such that the trajectory between data and noise follows a “straight” path (Liu et al., 2023; Albergo & Vanden-Eijnden, 2023; Lipman et al., 2022). This trajectory is represented as:

$${\bm{x}}_{t}=(1-t)\cdot{\bm{x}}_{0}+t\cdot\bm{\epsilon},$$

(5)

where $t\in[0,1]$. In this formulation, the gradient along the trajectory, $\bm{\epsilon}-{\bm{x}}_{0}$, is deterministic, often referred to as the velocity. The model $v_{\Theta}({\bm{x}}_{t},t)$ is parameterized to predict velocity by minimizing:

$$\min_{\Theta}\mathbb{E}_{t\sim\pi(t),\bm{\epsilon}\sim\mathcal{N}(0,I)}\left[\|v_{\Theta}({\bm{x}}_{t},t)-(\bm{\epsilon}-{\bm{x}})\|^{2}\right].$$

(6)

We note that this objective is equivalent to a score matching form (Eq. 4), with the weight $w_{t}=(\frac{1}{1-t})^{2}$. This equivalence highlights that alternative model parameterizations reduce to a standard denoising objective, where the primary difference lies in the time-dependent weighting functions and the corresponding optimization trajectory (Kingma & Gao, 2024).

During sampling, the model follows a simple probability flow ODE:

$${\textnormal{d}}{\bm{x}}_{t}=v_{\Theta}({\bm{x}}_{t},t)\cdot{\textnormal{d}}t.$$

(7)

Although a perfect straight path could theoretically be solved in a single step, the independent coupling between data and noise often results in curved trajectories, necessitating multiple steps to generate high-quality samples (Liu et al., 2023; Lee et al., 2024). In practice, we iteratively apply the standard Euler solver (Euler, 1845) to sample data from noise.

We introduce $\bm{\epsilon}$-VAE, with an overview provided in Figure 1. The core idea is to replace single-step, deterministic decoding with an iterative, stochastic denoising process. By reframing autoencoding as a conditional denoising problem, we anticipate two key improvements: (1) more effective generation of latent representations, allowing the downstream latent diffusion model to learn more efficiently, and (2) enhanced decoding quality due to the iterative and stochastic nature of the diffusion process.

We systematically explore the design space of model architecture, objectives, and diffusion training configurations, including noise and time scheduling. While this work primarily focuses on generating continuous latents for latent diffusion models, the concept of iterative decoding could also be extended to discrete tokens, which we leave for future exploration.

We evaluate the effectiveness of $\bm{\epsilon}$-VAE on image reconstruction and generation tasks using the ImageNet (Deng et al., 2009). The VAE formulation by Esser et al. (2021) serves as a strong baseline due to its widespread use in modern image generative models (Rombach et al., 2022; Peebles & Xie, 2023; Esser et al., 2024). We perform controlled experiments to compare reconstruction and generation quality by varying model scale, latent dimension, downsampling rates, and input resolution.

Model configurations. We use the encoder and discriminator architectures from VQGAN (Esser et al., 2021) and keep consistent across all models. The decoder design follows BigGAN (Brock et al., 2019) for VAE and from ADM (Dhariwal & Nichol, 2021) for $\bm{\epsilon}$-VAE. Additionally, we experiment with the DiT architecture (Peebles & Xie, 2023) for $\bm{\epsilon}$-VAE. To evaluate model scaling, we test five decoder variants: base (B), medium (M), large (L), extra-large (XL), and huge (H), by adjusting width and depth accordingly. Further model specifications are provided in Appendix B.

Training. The autoencoder loss follows Eq. 1, with weights set to $\lambda_{\text{LPIPS}}=0.5$ and $\lambda_{\text{adv}}=0.5$. We use the Adam optimizer (Kingma & Ba, 2015) with $\beta_{1}=0$ and $\beta_{2}=0.999$, applying a linear learning rate warmup over the first 5,000 steps, followed by a constant rate of 0.0001 for a total of one million steps. The batch size is 256, with data augmentations including random cropping and horizontal flipping. An exponential moving average of model weights is maintained with a decay rate of 0.999. Both VAE and $\bm{\epsilon}$-VAE are trained to reconstruct $128\times 128$ images. The default setup uses a downsampling factor of 16 and latent dimensions with 8 channels. We also test downsampling rates of 4, 8, and 32, as well as latent dimensions with 4, 16, and 32 channels. All models are implemented in JAX/Flax (Bradbury et al., 2018; Heek et al., 2024) and trained on TPU-v5lite pods.

Evaluation. We evaluate the autoencoder on both reconstruction and generation quality using Fréchet Inception Distance (FID) (Heusel et al., 2017) as the primary metric, computed on 10,000 validation images. For reconstruction quality (rFID), FID is computed at both training and higher resolutions to assess generalization across resolutions. For generation quality (FID), we generate latents from the trained autoencoders and use them to train the DiT-XL/2 latent generative model (Peebles & Xie, 2023). This latent model remains fixed across all generation experiments, meaning improved autoencoder latents directly enhance generation quality. We also report Inception Score (IS) (Salimans et al., 2016) and Precision/Recall (Kynkäänniemi et al., 2019) as secondary metrics.

Distribution-aware compression. Traditional image compression methods optimize the rate-distortion trade-off (Shannon et al., 1959), prioritizing compactness over input fidelity. Building on this, we also aim to capture the broader input distribution during compression, generating compact representations suitable for latent generative models. This approach introduces an additional dimension to the trade-off, perception or distribution fidelity (Blau & Michaeli, 2018), which aligns more closely with the rate-distortion-perception framework (Blau & Michaeli, 2019).

Iterative and stochastic decoding. A key question within the rate-distortion-perception trade-off is whether the iterative, stochastic nature of diffusion decoding offers advantages over traditional single-step, deterministic methods (Kingma, 2013). The strengths of diffusion (Ho et al., 2020) lie in its iterative process, which progressively refines the latent space for more accurate reconstructions, while stochasticity allows for capturing complex variations within the distribution. Although iterative methods may appear less efficient, our formulation is optimized to achieve optimal results in just three steps and also supports single-step decoding, ensuring decoding efficiency remains practical (see Figure 3 (left)). While stochasticity might suggest the risk of “hallucination” in reconstructions, the outputs remain faithful to the underlying distribution by design, producing perceptually plausible results. This advantage is particularly evident under extreme compression scenarios (see Figure 4), with the degree of stochasticity adapting based on compression levels (see Figure 5).

Scalability. As discussed in Section 4.1, our diffusion-based decoding method maintains the resolution generalizability typically found in standard autoencoders. This feature is highly practical: the autoencoder is trained on lower-resolution images, while the subsequent latent generative model is trained on latents derived from higher-resolution inputs. However, we acknowledge that memory overhead and throughput become concerns with our UNet-based diffusion decoder, especially for high-resolution inputs. This challenge becomes more pronounced as models, datasets, or resolutions scale up. A promising future direction is patch-based diffusion (Ding et al., 2024; Wang et al., 2024b), which partitions the input into smaller, independently processed patches. This approach has the potential to reduce memory usage and enable faster parallel decoding.

We present $\bm{\epsilon}$-VAE, an effective visual tokenization framework that introduces a diffusion decoder into standard autoencoders, turning single-step decoding into a multi-step probabilistic process. By exploring key design choices in modeling, objectives, and diffusion training, we demonstrate significant performance improvements. Our approach outperforms traditional visual autoencoders in both reconstruction and generation quality, particularly in high-compression scenarios. We hope our concept of iterative generation during decoding inspires further advancements in visual autoencoding.