Epsilon-VAE: Denoising as Visual Decoding

Thank you for your constructive comments. We will rephrase our claims and reduce the redundancy as suggested, and include suggested related work in the revision. Below we provide a point-by-point response to all of your questions. Please let us know if you have any further questions.

Q1: Because of the denoising procedure with 3 network forward evaluations used in the paper, the decoding is more expensive than with a standard VAE, which the paper misses to openly discuss. This limits its application in cases with real-time visualization, e.g., during the generation process of images.

Thank you for pointing out this limitation, and we will include a detailed discussion of it in the revised version. However, we would like to note that our denoising process demonstrates promising results even with a single iteration (please refer to Figure 3, left and right, in the original paper). Consequently, this allows our model to be adapted for scenarios with latency-sensitive requirements, such as real-time visualization during image generation as mentioned, by reducing the decoding step to a single pass.

Q2: How would the pipeline of eps-VAE and LDM benefit from even higher downsampling factors? Could the efficiency of the diffusion model be improved (both in terms of training and sampling) without hurting the generation quality significantly?

To address the questions regarding higher downsampling factors, we present additional results for a 32 x 32 downsampling factor in the table below, comparing them to our 16 x 16 results. Notably, Epsilon-VAE achieves a 25% improvement in generation quality over SD-VAE at the 32 x 32 factor, alongside a 3.2x inference speedup than SD-VAE at the 16 x 16 factor with comparable FID (highlighted in bold). We observed similar training speedups for latent diffusion models utilizing Epsilon-VAE at this higher downsampling rate. These gains are more pronounced than those observed when increasing the downsampling factor from 8 x 8 to 16 x 16 (where we observe 22% quality improvement and 2.3x inference speedup, respectively, as shown in the original paper). These findings strongly suggest that the benefits of the Epsilon-VAE and LDM pipeline are amplified with higher downsampling factors. We will include detailed results in the revised paper.

Q3: Why does the reconstruction performance degrade when increasing the number of steps starting from 3?

To enable large step sizes for the reverse process during inference, we introduced the denoising trajectory matching loss to implicitly model the conditional distribution $p(x_0|x_t)$, shifting the denoising distributions from traditional Gaussian to non-Gaussian multimodal forms (please refer to [1*] for a detailed discussion on this). However, the assumptions underpinning this approach are most effective when the total number of denoising steps remains small. Consequently, there appears to be an optimal range or "sweet spot" for the number of total inference steps. We will elaborate on this phenomenon and provide further clarification in the revised version.

[1*] Tackling the Generative Learning Trilemma with Denoising Diffusion GANs, ICLR 2022.

Q4: The qualitative comparisons in the paper and the appendix are not very convincing showing only very minor differences in high-frequency details.

This is because in the original paper, we show compressed high-resolution images with Epsilon-VAEs using low downsampling factors. We will include additional uncompressed visual results under more changeling settings (e.g., 128 x 128 images with 16 or 32 downsampling factors) in the revised version to highlight our advantages.

Q5: The proposed reversed logarithm mapping for the timesteps during sampling lacks intuition and would benefit from a visualization.

As detailed in our response to Q3, the denoising trajectory matching loss results in denoising distributions that deviate from the standard Gaussian. This deviation suggests that the conventional uniform sampling approach may no longer be optimal. Hence, we empirically investigated alternative sampling strategies and found the reversed logarithm mapping to yield the best performance. We will clarify the intuition and provide additional visualizations as suggested in the 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: The reported performance in Table 3 is worse than numbers reported in the original DiT paper.

We did the following modifications to the original DiT training recipe as mentioned in Section C of the original paper, leading to a slight drop in performance: (1) we reimplement the pipeline with JAX for training models on TPUs; (2) for simplicity and training stability, we remove the variational lower bound loss term; (3) we reduce the number of training steps to 1M to conserve compute.

Q2: It would greatly help validate the contribution of the work by showing how ϵ-VAE works with CFG.

We provided results with CFG under the 8 x 8 downsample factor in the table below, where we find that Epsilon-VAE (M) performs relatively 20% better than SD-VAE and further improvements are obtained after we scale up our model to Epsilon-VAE (H). These results are consistent with the results without CFG, confirming the effectiveness of our model. We will provide more detailed results of other models and under different downsample factors in the revised version.

Q3: In latent diffusion, is the DiT setting optimized for Epsilon-VAE or is it also benefitting standard DiT with SD-VAE?

Our training protocol for latent diffusion was consistently applied across all evaluated VAE models. The three minor modifications to the original DiT training recipe, as outlined in our response to Q1, were implemented solely to conserve computational resources. Our current setting does not favor either our Epsilon-VAE or the standard SD-VAE, but it ensures a fair basis for comparison.

Q4: Can authors provide more details of how inference is conducted on higher resolutions with models trained at 128 x 128 images?

Inference on higher resolutions works out-of-the-box. Since our model is a fully convolutional UNet, it can directly process images with resolutions exceeding the 128 x 128 training size without any modifications to the input.

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

Q1: Clarification on high-level insights.

Thank you for raising this concern. Rather than opposing latent diffusion, we offer a complementary perspective by demonstrating that the diffusion process can also be applied to the outer structure of the latent diffusion model – the autoencoder part. Our goal is to revisit the standard autoencoder framework and integrate diffusion into the decoding step to enhance the entire model’s generative capacity. We empirically show that a latent diffusion DiT integrated with Epsilon-VAE significantly outperforms its SD-VAE counterpart, with the performance gap being significant in high compression regime.

Q2: Diffusion-based iterative decoding.

While replacing single-step decoding with an iterative process may seem counter-intuitive due to increased computational cost, the diffusion-based decoder addresses this concern in three key ways. First, it offers scalable inference, where even a single-step variant already outperforms a plain VAE decoder, and additional steps further enhance quality (please refer to Figure 3, left and right, and Table 7 in the original paper, where a single-step Epsilon-VAE-lite achieves 7.18 rFID and a VAE obtains 11.15 rFID in the same setting). Second, it provides controllable trade-offs between computation and visual fidelity, allowing the number of steps to be adjusted at inference time based on application needs. Third, as shown in Section 4.2 of the main paper and our response to Q3 below, it enables training under higher compression ratios, which helps offset the added cost of iterative decoding by reducing the size of latent representations.

An additional advantage of scaling the autoencoder over the latent model lies in computational efficiency. Recent trends show latent diffusion models increasingly adopt Transformer architectures (e.g., DiT), where self-attention scales quadratically with input resolution. In contrast, our convolution-based UNet decoder offers more favorable linear scaling. As models grow, shifting complexity to the autoencoder helps reduce the burden on the latent model, leading to a more efficient overall system.

Q3: A benchmark on computation at various decoding resolutions.

In response to the reviewer's valuable point regarding computational cost at various decoding resolutions, we conducted experiments with Epsilon-VAE using a 16 x 16 downsampling factor. Our analysis reveals that: (1) the dominant factor in the overall memory footprint during the generation process is the LDM; (2) increasing the decoding resolution leads to a more substantial increase in the memory requirements of the LDM compared to the Epsilon-VAE decoder itself; (3) Epsilon-VAE outperforms SD-VAE under the same 8 x 8 downsampling factor at the 256 x 256 resolution, but with slightly worse throughput and memory efficiency; and (4) importantly, by employing a 16 x 16 downsampling factor instead of 8 x 8, Epsilon-VAE demonstrates a significant 2.3x inference speedup and a 3.3x reduction in total memory cost compared to SD-VAE, while maintaining comparable image generation performance (highlighted in bold). This demonstrates the efficiency gains achieved with our proposed approach in mitigating the potential for undesirable decoding memory scaling.

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

Q1: The LPIPS and GAN loss are applied on the estimated one-step sample, which may not be accurate and may cause objective bias in theory. Although this can potentially be addressed with finetuning a diffusion decoder with frozen $z$.

We acknowledge the reviewer's point regarding potential objective bias due to applying the LPIPS and GAN loss on the estimated one-step sample. However, we would like to emphasize that our Epsilon-VAE differs significantly from traditional diffusion models in that its diffusion decoder is conditioned on encoded latents. This conditioning provides a strong prior about the input image to reconstruct, resulting in a more accurate estimated one-step sample than in typical diffusion scenarios. Therefore, we believe the potential for objective bias is considerably reduced. We agree that finetuning the diffusion decoder with frozen $z$ is a promising avenue for further improvement and will explore this in future work.

Q2: Besides rFID, how much improvement it gets from the additional LPIPS and GAN losses for the actual visual quality in images?

We find that the LPIPS loss enhances textural fidelity and structural coherence in generated images, while the GAN loss contributes to sharper high-frequency details. These observations align with their established roles in traditional VAEs. Detailed visual comparisons demonstrating these improvements will be included in the revised version.

Q3: The zoom-in boxes in Fig. 4 seem to be not very clear even for the GT.

Thanks for pointing out this. This is because the compressed images were included in the original paper to ensure a reasonable file size. We will provide the original, uncompressed images in the revised version.