$\epsilon$-VAE: Denoising as Visual Decoding

Thank you for your valuable comments. We believe there are misunderstandings of our method in Q1, Q4, and Q5. We would like to address all your concerns below. Please let us know if you have any further questions.

Q1: Why is FID used to validate generation quality? I see you used IS in the experiment section instead.

FID is a standard metric to measure image generation quality which is widely used in the literature. We reported both FID and IS (together with Precision and Recall) in Table 2 of the main paper for evaluating image generation quality.

Q2: The novelty is limited. The core innovation is only improving the VAE decoder, which offers minimal technical contribution.

Please refer to Q1 in the common response for justifications on our novelty and contributions.

Q3: The Introduction mentions various tokenizers, but the actual comparison seems to focus solely on VAE. What about discrete tokenizers like VQVAE and VQGAN?

We provide a comparison with VQ models in the table below. Since the main focus of our paper is on improving VAEs for latent diffusion models, which require continuous latents, implementing VQ versions of our model is left for future work. We believe the techniques proposed in our paper can be easily transferred to VQ models, as our method does not rely on specific assumptions about the encoder outputs.

Q4: You mention that the proposed model’s inference time is better than VAE, but also admit that "it requires more compute costs than VAE due to its U-Net design." In Section 5’s Discussion, you briefly touch on potential future optimizations. Does this imply that this limitation in your model is unsolvable?

Thank you for your question. To clarify, we have not claimed in the paper that the proposed model’s inference time is better than that of VAEs. As noted in L357, we acknowledge that our model has lower throughput compared to traditional VAEs due to the UNet-based diffusion decoder. However, we do not consider this limitation "unsolvable." In Section 5, we outline potential future optimizations, such as adopting lightweight architectures or implementing a patch-based decoder, which we believe can significantly reduce the compute cost while retaining the model's performance. These approaches are promising directions for addressing this limitation, and we are optimistic about their feasibility.

Q5: You discuss VAE optimization in the context of Diffusion Models, but the inference time you’re referring to compares VAE’s reconstruction speed. In practice, the major time consumption in generation models isn’t in the VAE part. The faster inference time you mention is insignificant because the time it saves is negligible within 50 DDIM steps.

The latent diffusion model (LDM) remains consistent across all experiments, resulting in the constant inference time for the LDM. Consequently, the decoding time of the VAEs become a key factor influencing the overall inference speed. This is why our analysis in the main paper focuses on comparing the decoding speeds of different VAEs.

On the other hand, our e-VAE supports a high compression rate, reducing the token length to one-fourth (x1/4) while preserving as much information as the original token length. This reduction can significantly accelerate latent diffusion model generation, decreasing overall inference time while maintaining favorable generation quality (please refer to Q2 in the common response).

We thank you for your reply, and would like to further clarify our claims.

Q1. To clarify our statement, it's important to note that our analysis focuses on the impact of different VAEs on inference time while keeping the LDM unchanged. This is because our comparisons (Tables 1 and 2 in the main paper) utilize the same LDM throughout, ensuring its inference time remains unchanged. In this controlled setting, the decoding time of the VAE becomes the remaining factor influencing variations in overall inference time.

While we acknowledge that the LDM generally dominates overall inference time, our statement aims to highlight the significant role of VAE decoding time within the above specific context. To make this clear, we propose rephrasing our statement to: “When utilizing a fixed LDM, the decoding time of the VAE emerges as the remaining primary factor affecting overall inference time.”

Q2. We respectfully argue that our statement does not inherently suggest that the issue is unsolvable within the current framework. Our statement was intended to highlight an area for future improvement, rather than suggest a fundamental limitation. To further clarify, we note that (1) our proposed optimizations based on model architecture and decoding method are complementary to the current framework; (2) since our decoder utilizes a diffusion model, any advancements in diffusion model efficiency will directly benefit our framework and contribute to addressing the issue.

We will revise our statement to better reflect this and avoid any misunderstanding.

Thank you for your constructive comments. Below we provide a point-by-point response to all of your questions. Please let us know if you have any further questions.

Q1: Aligning latents and image resolutions.

As described in L197–199 and illustrated in Figure 1, low-resolution latents obtained from the encoder are upsampled to match the resolution of the noisy image before being passed into the diffusion decoder. We also explored an alternative conditioning strategy that modulates the decoder layers through AdaGN (in Appendix C). However, due to its increased computational complexity, we chose to use resizing and concatenation for their simplicity and effectiveness.

Q2: Sampling efficiency.

Please refer to Q1 in the common response. To further emphasize this aspect, as shown in Table 3 and Figure 3 (left), our e-VAE achieves competitive performance with single-step decoding and delivers the best results with three steps, all without relying on distillation or consistency regularizations. This is made possible by our parameterization strategy (velocity-prediction) and the integration of LPIPS (applied to $\hat{x}_0$) and GAN losses (trajectory matching). Such sampling efficiency has not been demonstrated by any prior diffusion decoder-based VAE approaches.

Q3: Throughput.

As noted in L357, we also measure the actual throughput. Admittedly, our framework has lower throughput compared to traditional VAEs due to the use of a UNet-based diffusion decoder. However, we believe this limitation can be mitigated in the future by adopting lightweight architectures or implementing patch-based decoder, as described in Section 5.

Thank you for your reply and the opportunity to further clarify this aspect of our work.

To address your concern, we want to emphasize that the e-VAE is indeed designed as a better alternative to the standard VAE which are paired with latent diffusion models (LDMs) for image generation. As demonstrated in Table 2 of the main paper, e-VAEs’ generation qualities consistently outperform VAEs in this context. These comparisons were conducted by first denoising in the latent space using an LDM (with a fixed architecture, i.e., DiT-XL/2) and then decoding the resulting latents into images using decoders from either VAEs or e-VAEs.

To better illustrate the trade-off between generation quality and inference time when paired with LDMs, we provided the following table (reorganized from the second table in "Common response to all reviewers [Part 2]") with 'f8' and 'f16' denoting downsampling rates of 8 and 16 for the VAEs/e-VAEs, respectively.

This table highlights two key observations regarding the trade-off:

• (a) v.s. (b): Replacing the SD-VAE (f8) directly with e-VAE (f8) results in better generation quality but slightly longer inference time. This is attributed to the UNet-based diffusion decoder in the e-VAE, which improves quality while adding computational complexity. Our improvements are more significant when the setting is challenging, e.g., with higher downsampling rates and weaker encoder (see Table 2 of the main paper).

• (a) v.s. (c): Replacing the SD-VAE (f8) with e-VAE (f16) results in slightly worse generation quality but significantly reduced inference time. This reduction is achieved by decreasing the input token length to the LDM, rather than reducing the LDM's inference steps, as we maintain the original LDM design unchanged throughout our experiments. This is supported by the inherent ability of e-VAEs to provide significantly better reconstruction quality than VAEs when the downsampling rate is high (see Table 1 of the main paper).

It's important to note that our method is orthogonal to other techniques aimed at improving LDM inference time through distillation. These techniques, as mentioned in the review, could be directly applied to our framework since we do not modify the LDMs themselves.

We hope the above explanation addresses your concern. Please let us know if you have any further questions.

Dear Reviewer CU6D,

We sincerely appreciate your valuable feedback on our paper. In our last response, we have provided detailed explanations on the setup of our model for image generation and trade-off between generation quality and inference time.

We kindly inquire whether these clarifications have adequately addressed your concerns. Please feel free to let us know if you have any further questions.

Thank you a lot!

Best regards,

Authors of Paper 332

Improving the quality of generated images using $\epsilon$-VAE comes at the cost of increased inference time. Similarly, increasing the number of sampling steps can also enhance the quality of the generated images. Why use $\epsilon$-VAE instead of increasing the number of inference steps in the vanilla SD?.

The authors mention that "our method is orthogonal to other techniques aimed at improving LDM inference time through distillation." However, I have concerns about this claim. It is necessary to validate this through experiments by combining existing distillation models, such as SD-Turbo, LCM, LCM-LoRA, and SwiftBrush.

SD-Turbo: Adversarial Diffusion Distillation
LCM: Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference.
LCM-LoRA: LCM-LoRA: A Universal Stable-Diffusion Acceleration Module
SwiftBrush: One-Step Text-to-Image Diffusion Model with Variational Score Distillation

Thank you for your valuable comments. Below we provide a point-by-point response to all of your questions. Please let us know if you have any further questions.

Q1: Why not using DiT with a patch size of 2?

In our preliminary experiments, we explored the DiT/2 but found it unsuitable for high-resolution image decoding. For example, the Transformer's computational cost increases quadratically with token length, making it prohibitively expensive to decode a 256x256 image, which results in 128^2 tokens. Notably, the original DiT’s optimal patch size of 2 was designed for latent-level generation, where token lengths are typically small and manageable. In contrast, our decoder operates at the pixel level, leading us to consider only patch sizes of 4 and 8 for DiT in our experiments. Ultimately, we found ADM to produce better results and to be more suitable for this task.

Q2: It would strengthen the paper to include additional commonly used metrics, such as PSNR, SSIM, and LPIPS, to provide a more comprehensive assessment.

Under the Stable Diffusion configuration, e-VAE achieves 0.152 LPIPS, 25.11 PSNR, and 0.71 SSIM on ImageNet, performing comparably to the standard SD-VAE. As discussed in Section 5 of the main paper, our approach prioritizes preserving the overall perceptual distribution of images rather than achieving pixel-perfect reconstruction. This aligns with our focus on perception-based compression under high compression rates. As a result, e-VAE excels in metrics such as rFID, which capture differences in perceived image distributions, rather than in pixel-level metrics like PSNR and SSIM.

Q3: The paper lacks mention and discussion of previous diffusion-based autoencoders like DiffusionAE or DiVAE.

Please refer to Q1 in the common response for discussions on DiffusionAE and DiVAE. We also add discussions with these works in the paper revision.

Q4: Clarification on our contributions. What are the "new insights" you mention in the abstract?

Please refer to Q1 in the common response for justifications on our contributions and impact.

Q5: However, the reason that it not become popular, i believe, is time-intensive both training and inference, a challenge that this paper does not seem to overcome yet.

Our approach tackles the inefficacy of diffusion decoders by enhancing both training and inference speed. Unlike traditional diffusion models, e-VAEs could be trained on low-resolution images and well-generalized to high-resolution images (see Table 1). For example, our model achieves a training speed of 9 steps/sec on 128x128 images, comparable to VAEs (11 steps/sec). During inference, our approach can achieve fast decoding (1–3 steps) through the integration of our proposed objectives and parameterizations (Please see Q1 in the common response for detailed justifications). We believe further efficiency gains are anticipated with improved architectures and decoding methods, as discussed in Section 5.

Q6: Comment on the writing.

Thank you for your meticulous review. We will fix the typo in the paper revision.

Thank you for your valuable comments. Below we provide a point-by-point response to all of your questions. Please let us know if you have any further questions.

Q1: Clarification on the term "tokenization”.

As noted in L33, tokenization refers to the process of transforming data into compact representations, which can be either continuous or discrete. As suggested, we will further clarify that continuous latents are typically obtained through an encoder in vision tasks, while discrete tokens are usually derived using embeddings in language tasks.

Q2: Clarification on the naming “e-VAE”.

Thanks for pointing this out. We use the term “e-VAE” because the parameterizations in the diffusion decoder, whether epsilon-prediction or velocity-prediction, essentially transform into a traditional score-matching objective (L135–137), which can be interpreted as a reweighted version of the ELBO. Furthermore, our proposal directly builds upon the original VAE framework by replacing the reconstruction objective with the score-matching objective, while retaining LPIPS and GAN losses by ensuring compatibility with the diffusion decoder. For these reasons, we believe the term "e-VAE" is appropriate. However, if this naming remains unclear, we are open to considering alternative terms.

Q3: Comparison with plain diffusion ADM.

As suggested, under the same training setup, we directly trained a plain diffusion model (ADM) for comparison, which resulted in rFID score of 38.26. Its conditional form is already provided as a baseline in Table3, achieving 28.22. This demonstrates that our conditional form $p(x_{t-1}|x_t,z)$ offers a better approximation of the true posterior $q(x_{t-1}|x_t,x_0)$ compared to the standard form $p(x_{t-1}|x_t)$. By further combining LPIPS and GAN loss, we achieve rFID of 8.24, outperforming its VAE counterpart, which achieves 11.15. With better training configurations, our final rFID improves to 6.24. This progression, from plain diffusion ADM to e-VAE, underscores the significance of our proposals and their impact.

Q4: Sampling steps.

We apologize for the unclear description. As noted in Table 3 and Figure 3 (left), our e-VAE achieves the best results using three sampling steps, while still delivering competitive performance even with a single step. We will ensure this is clarified more explicitly in the text.

Q5: Typos.

Thank you for your meticulous review. We will make the suggested corrections in the manuscript.

Thanks for your insightful comments. Most of your questions are answered in the common response, and we provide a summary below. We sincerely hope you could reconsider your rating regarding our newly added results. Please let us know if you have any further questions.

Q1: Novelty.

Please refer to Q1 in the common response for justifications on our novelty and contributions.

Q2: Additional quantitative results on ImageNet and COCO.

Please refer to Q2 in the common response for the additional quantitative results that aligns with the SD-VAE configuration. We also included results on COCO as suggested. We leave reconstruction on video data as our future work since (1) the major focus of this paper is image reconstruction; (2) the rebuttal time and compute budget are limited.

Q3: Additional visual results for reconstruction.

Please refer to Q3 in the common response for the additional visual results on face details and text details. Overall, we find that our models perform significantly better than SD-VAE when the model and training configurations are matched.

Thanks for your positive feedback. Below we provide a point-by-point response to all of your questions. Please let us know if you have any further questions.

Q1: The LPIPS and GAN loss are applied on the estimated sample, which seems to be not accurate that may cause objective bias in theory.

This is indeed a good question. Unlike traditional diffusion models, our e-VAEs utilize a diffusion model conditioned on the encoder outputs. This conditioning provides strong prior information, significantly improving the estimation quality of the samples (see Q3 from bhvX), thereby reducing potential objective bias.

Q2: Visual comparisons at higher resolutions under the Stable Diffusion configuration.

Please refer to Q3 in the common response for visual comparisons between SD-VAE and e-VAE under x8 downsampling and 4 channels at resolutions of 256x256 and 512x512. Generally, e-VAE performs better when reconstructing complex regions such as human faces and small texts.

Q3: The generation FIDs are a bit high in Table 2.

This is because we perform unconditional image generation, which usually results in higher FID compared to class-conditional image generation. Additionally, the employed VAEs have a fraction of parameters of their SD-VAE counterparts (6M vs. 34M) and use a higher compression rate (16 vs. 8), being trained under a more challenging configuration (see Q2 in the common response).

Q4: A few related works [1, 2] that might be missing in the paper.

Thanks for pointing them out. Please refer to Q1 in the common response for our discussion. We also add discussions with these works in the paper revision.

Q5: How is the baseline VAE (GAN-based autoencoder) designed and scaled up in the paper? Are they all re-trained to match the setting of eps-VAE? Is the number of channels 8 for Figure 4 and 5?

The design of baseline VAEs is summarized in L277 of the main paper, and they are scaled up in the same way as e-VAEs (see Table 4). Yes, all the baseline VAEs in the original paper are re-trained to match the setting of eps-VAE. The number of channels is 8 for Figures 4 and 5.

Dear Reviewer 9E74,

We sincerely appreciate your recognition and valuable feedback on our paper. In our response, we have provided detailed quantitative and high-resolution visual comparisons of our model against SD-VAE. Additionally, we have clarified the loss functions and baseline setups, and discussed the related work you pointed out.

We kindly inquire whether these clarifications and results have adequately addressed your concerns. Please feel free to let us know if you have any further questions.

Thank you a lot!

Best regards,

Authors of Paper 332

Q2: Additional quantitative results aligning with the Stable Diffusion setting.

In the main paper, we use a light-weight encoder (with 6M parameters), the compression rate of 16, and the channel dimension of 8 for 128x128 image reconstruction, since we focus on handling high compression rates. This configuration is more challenging than the traditional setup of VAEs in Stable Diffusion (SD-VAE), which leads to the performance gap.

In this rebuttal, to address reviewers’ concerns on reconstruction quality (9E74, drdo), we provide additional results of our models under the same configuration with SD-VAE, i.e., with a standard encoder (with 34M parameters), the compression rate of 8, and the channel dimension of 4 for 256x256 image reconstruction. We report rFID on the full validation set of ImageNet and COCO2017. Comparisons with SD-VAE and SDXL-VAE are shown in the table below.

We find that Epsilon-VAE outperforms SD-VAEs when the decoder sizes are similar, and our results could be further improved when we scale up the decoder.

In addition, we want to highlight that when combined with latent diffusion models for class-conditional image generation, e-VAE is able to achieve comparable generation quality even when using only 25% of the typical token length required by SD-VAE. To show this, we train an additional e-VAE (M) under the Stable Diffusion configuration but double the downsampling rate. Then, we compare our models with SD-VAE by training DiT/2 under the class-conditional image generation setup (no classifier-free guidance) on ImageNet 256x256. We follow the setup in the DiT paper and all DiTs are trained with 1M steps. Results are shown in the table below.

Q3: Additional high-resolution visual results aligning with the Stable Diffusion setting.

As requested by reviewers (9E74, drdo), in the appendix of the revised paper (Appendix D, Pages 21-22), we provide additional apple-to-apple visual comparisons between e-VAE and SD-VAE under the SD-VAE configuration at the resolutions of 256x256 and 512x512. We observe that e-VAE achieves significantly better visual qualities than SD-VAE when reconstructing local regions with complex textures or structures, such as human faces and small texts.

Hey Xu,

Thanks for your interest and kind words. We answer your questions below.

1. We feel these lines of work are somehow different from ours in the main goal. They focus on training stronger decoders to take more duties on the generation side in autoencoding, so that the latents could be used for modeling other perspectives (e.g., semantic information) potentially for specific purposes (e.g., unification); while, we remain on developing a VAE for achieving better visual qualities in pure generative models. I believe the direction you mentioned is indeed an important topic where our method could be extended to.

2. (1) Generally, it largely depends on the use case. In my own view, for a pure generative model, low-level latents are usually beneficial as they alleviate the duty of a decoder (low-level latents are easier for reconstruction); high-level/semantic latents are usually required in understanding models or unified models, while a heavier decoder (like the models in Q1) might be essential for good generation qualities in this case. (2) We expected that our method could work with discrete representations out-of-the-box, since we do not make any assumptions about the latent forms in our model.

3. Thank you for your suggestion. While we agree that incorporating AR approaches could enhance our current results, we believe that using diffusion models alone still provides a sufficient showcase. This is because latent prediction and autoencoding are two orthogonal functionalities, even diffusion models are applied in both cases.

4. Yes, this is indeed a good point. We will consider incorporating it in the revision.

5. While we haven't specifically tested such scenarios, traditional VAEs have demonstrated a certain degree of robustness to corrupted latents (https://arxiv.org/pdf/2409.16211), which we believe could be transferred to our method. Furthermore, our method's emphasis on perception-compression (see the “Discussion” section in our paper) might offer additional advantages in handling these cases.

Thanks,

Long