LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models

Summary

This paper provides a LiteVAE structure to replace the original one, aiming to reduce the computations when training on large-scale datasets, which can boost the performance of latent diffusion models with more potential augmentations during training.

Strengths

1. The new encoder could reduce more parameters than the original one with less GPU memory. 2. The paper provides some interesting analysis of the feature map. 3. The first half of the paper is easy to read.

Weaknesses

* There is less information about the decoder. If without the lite version of the decoder, the innovation of the work may be limited because the encoder would not be used during inference. What is the structure of the decoder? Will the decoder be trained with an encoder? Did the authors implement a new lite decoder? * Have the authors verify their structure in some known Latent diffusion models with fine-tuning strategies. Without such validation, this work may not contribute effectively to modern latent diffusion models. (There is no need to provide new experiments to verify this.) * Did the paper have the ablation study on the structure modules in Fig.1? * There is less visual comparison between LiteVAE and VAE to verify its effectiveness. (Fig.6 and Fig.9 didn't show the results of VAE)

Questions

* Did the authors try to expand the bottleneck channel from 4 to other numbers, such as 16, 32, or more? * Could the authors provide some distribution analysis about the latent features and the differences between LiteVAE and VAE? * Could the authors provide the values of LiteVAE-S/M/L in Table 3? * The format of the Checklist does not fully meet the requirements.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The limitations of their work and any potential negative societal impacts should be strengthened.

Summary

This paper presents LiteVAE. LiteVAE is an efficient and lightweight modification to latent diffusion models (LDMs) that incorporate 2D wavelet transform into the encoding structure. It then uses a feature aggregating model (UNet-based architecture) to fuse multiscale wavelet coefficients into a unified latent code. It uses a decoder to transform the latent code into an image. This work also provides other modifications such as self-modulated convolution, pixel-wise discrimination, removing adaptive weight for adversarial loss, and additional loss functions. These modifications further enhance the training dynamics and reconstruction quality of LiteVAE. These modifications lead to a considerable reduction in the computational cost compared to the standard VAE encoder without sacrificing the quality of image construction. Bse LiteVAE with a six-fold reduction in encoder parameter count matches the reconstruction quality of standard VAE encoders. Large LiteVAE models provide better reconstruction quality than the standard VAE. Finally, this paper presents experimental results to support its claims. The result shows that large LiteVAE models outperform standard VAEs of similar size based on the following performance metrics: rFID, LPIPS, PSNR, and SSIM.

Strengths

ORIGINALITY: The main focus of this work is to use the traditional signal processing method (2D wavelet transform) to improve the performance and reduce the computational cost of deep learning methods. It explores the fact that the latent code of Stable Diffusion VAE (SD-VAE) is itself image-like. So they opted for a traditional signal processing method or transformation that preserves the image-like structure of the latent code for SD-VAE. Although there are prior works that use 2D wavelet transform to improve the performance of generative models. This work incorporates 2D wavelet transform and adds some other useful modifications. PRESENTATION: This paper is well organized and provides sufficient background information to understand the central claim of this work. For example, it section provides relevant background information to understand the main components of the new LiteVAE. The inclusion of relevant figures further improves the readability of this work. QUALITY: This is a high-quality. It clearly explains the motivation for this work and provides a detailed explanation and justification for the modifications presented in this work. It also provides a detailed explanation of how the modifications improve efficiency and scalability. Finally, it provides sufficient experiments to support the main claims of this work. SIGNIFICANCE: The reduction in computational cost with no tradeoff in reconstruction quality points to the significance of this work.

Weaknesses

Table 3 only compares the throughput of VAE and LiteVAE-B. Please provide the throughput for LiteVAE-S and LiteVAE-L. Also, the model did not discuss the additional computational cost for the 2D wavelet transform. Please provide some information about this. There is a typographical error in line 207. I think "Table 3" should be "Table 2".

Questions

Table 3 only compares the throughput of VAE and LiteVAE-B. Please provide the throughput for LiteVAE-S and LiteVAE-L. Also, the model did not discuss the additional computational cost for the 2D wavelet transform. Please provide some information about this. There is a typographical error in line 207. I think "Table 3" should be "Table 2".

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The checklist points to Section 7 for the limitations but the section only shows the conclusion. Please address this.

Summary

This paper introduces LiteVAE, a novel approach that combines multi-scale VAE and discrete wavelet transform to reduce computational cost and enhance reconstruction capabilities. Both components are well-grounded and supported by experimental results. Additionally, the paper provides a detailed pipeline and ablation studies on training VAE for diffusion models, which will also benefit readers.

Strengths

1. This paper integrates multi-scale VAE and discrete wavelet transform to reduce computational cost and boost reconstruction performance. 2. The paper offers a detailed training pipeline for VAE in diffusion models, along with ablation studies that readers will find beneficial. 3. The paper presents several interesting tricks for improving VAE training, including (1) removing group normalization from the decoder, (2) using a U-Net-based discriminator, and (3) eliminating the adaptive weight $\lambda_{reg}$.

Weaknesses

1. The multi-scale VAE and discrete wavelet transform are two relatively independent improvements for VAE, so I recommend that the authors conduct an ablation study on these two components. 2. The new VAE training pipeline is a valuable resource for the community, so I suggest that the authors release not only the VAE checkpoint but also the entire training code.

Questions

See the weaknesses part.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

None

Summary

The authors propose LiteVAE, a novel architecture for the VAE decoding step of latent diffusion. They show that LiteVAE can achieve comparative perfomance to SD-VAE, the default latent diffusion decoder, while using fewer parameters. The efficiency gain comes from using a more lightweight network, and a wavelet feature representation. The paper provides evaluation showing that LiteVAE outperforms the naïve approach of simply scaling down SD-VAE. This gives evidence that LiteVAE offers an architectural improvement.

Strengths

1. The paper provides an extensive experimental investigation of how to improve the efficiency of the VAE upscaling step, and is clearly presented. To me the contribution is significant, due to the dominance of diffusion modelling for image generation. A performance gain in even one step can improve the efficiency many real world uses. 2. The use of wavelets is relatively underexplored in the literature. Using them to improve perceptual quality in generated images is novel to the best of my knowledge. 3. All claims are supported by experimentation and ablation studies.

Weaknesses

1. Main concern is with the impact of this work. The motivation does not state how much more efficient their improvements would make the full diffusion pipeline. Is the VAE step such a significant performance bottleneck? 2. No code available for a paper whose main contribution is experimental. I would have liked to see how the wavelet features are handled, as Pytorch does not provide official modules that deal with wavelet decompositions. I am a bit concerned about how easy this model is to build and deploy. 3. I have some concerns about whether LiteVAE can offer a performance boost in domains other than natural images. If this is the case, maybe it should be mentioned as a limitation. My concern stems from the wavelet decomposition allowing to explicitly target high frequency content reconstruction, which is especially important for natural image upscaling. 4. Table 4 (difference between group norm and SMC) should ideally have error bars, as it claims an improvement in favour of SMC against group norm.

Questions

1. In line 24 the computational burden of the VAE is given and compared to the diffusion Unet. Is the 86 GFLOP figure for the latter given per diffusion step, or for the entire diffusion process? In general, how much of a performance gain would we expect for the full image generation pipeline when using LiteVAE instead of SD-VAE? 2. Can LiteVAE be used as a drop-in replacement of SD-VAE in real world applications? 3. In line 287 it is mentioned that LiteVAE could be applied to different scenarios. Wavelets are known to offer a good basis for representing natural images. What is the motivation behind hypothesising that this basis will yield good performance in other domains? 4. Could you comment on whether using the wavelet basis and reconstruction loss with the original SD-VAE might yield similar performance?

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

The limitations are not discussed.