Title: LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: variational autoencoders

latent diffusion models
Abstract: Advances in latent diffusion models (LDMs) have revolutionized high-resolution image generation, but the design space of the autoencoder that is central to these systems remains underexplored. In this paper, we introduce LiteVAE, a new autoencoder design for LDMs, which leverages the 2D discrete wavelet transform to enhance scalability and computational efficiency over standard variational autoencoders (VAEs) with no sacrifice in output quality. We investigate the training methodologies and the decoder architecture of LiteVAE and propose several enhancements that improve the training dynamics and reconstruction quality. Our base LiteVAE model matches the quality of the established VAEs in current LDMs with a six-fold reduction in encoder parameters, leading to faster training and lower GPU memory requirements, while our larger model outperforms VAEs of comparable complexity across all evaluated metrics (rFID, LPIPS, PSNR, and SSIM).

LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models

Abstract

Advances in latent diffusion models (LDMs) have revolutionized high-resolution image generation, but the design space of the autoencoder that is central to these systems remains underexplored. In this paper, we introduce LiteVAE, a new autoencoder design for LDMs, which leverages the 2D discrete wavelet transform to enhance scalability and computational efficiency over standard variational autoencoders (VAEs) with no sacrifice in output quality. We investigate the training methodologies and the decoder architecture of LiteVAE and propose several enhancements that improve the training dynamics and reconstruction quality. Our base LiteVAE model matches the quality of the established V AEs in current LDMs with a six-fold reduction in encoder parameters, leading to faster training and lower GPU memory requirements, while our larger model outperforms VAEs of comparable complexity across all evaluated metrics (rFID, LPIPS, PSNR, and SSIM).

DWT-L1

1 Introduction

Latent diffusion models (LDMs) [55] have recently assumed dominance in the field of high-resolution image generation, primarily due to their scalability and training stability over pixel-space diffusion. The training process of LDMs involves two separate stages. In the first, an expressive variational autoencoder (VAE) is trained to transform the raw pixels of an image into a more compact latent representation. In the second, a diffusion model is trained on the latent representations of training images. While numerous studies have investigated the scalability and dynamics of the diffusion component in LDMs [48, 33], the autoencoder element has received far less attention.

The VAE in LDMs is not only computationally demanding to train but also affects the efficiency of the diffusion training phase due to the resource requirements of querying a large encoder network for computing the latent codes. For example, as the autoencoder operates on high-resolution images, the VAE encoder of Stable Diffusion 2.1 uses 135 . 59 GFLOPS compared with 86 . 37 for the diffusion UNet. 1 This becomes an even greater concern for video diffusion models, as the encoder then needs to provide the latents for a batch of frames instead of a single image [3].

A common workaround for this resource burden is to precompute and cache the latent codes for the entire dataset to avoid having to use the autoencoder during diffusion training. However, in addition to its initial overhead, this approach eliminates the possibility of using on-the-fly techniques, such as data augmentation, which have been shown to improve the training and performance of diffusion models [32]. Using a large encoder also adds noticeable overhead in applications that are based on pretrained latent diffusion models. For example, when training 3D models through score distillation of 2D LDMs [51], the process necessitates backpropagating gradients through the LDM encoder, which is computationally intensive [38]. Beyond the computational aspects, improving the reconstruction quality of the autoencoder also improves the quality of generated images, as the autoencoder provides an upper bound on the generation quality [50, 13].

With these issues in mind, we investigated improving the efficiency of LDMs through their core V AE component with the goal of preserving overall quality. We show that with the help of the 2D discrete wavelet transform (DWT), we can considerably simplify the encoder network in LDMs. This leads to our proposal of LiteVAE, a new autoencoder design for LDMs, which has superior compute/quality trade-offs compared with standard VAEs.

LiteVAE consists of a lightweight feature-extraction module to compute features from the wavelet coefficients and a feature-aggregation module to combine these multiscale features into a unified latent code. A decoder then converts the latent code back to an image. An overview of the LiteV AE pipeline is shown in Figure 1.

We chose the wavelet transform due to its proven ability to represent rich, compact image features [43], and we argue that the wavelet decomposition simplifies the encoder's task by facilitating the learning of meaningful features. We examine the design space of LiteVAE in depth and propose several variations on the network architecture and training setup that further boost reconstruction quality and training efficiency.

Through extensive experimentation, we show that LiteVAE considerably reduces the computational cost of the standard VAE encoder while maintaining the same level of reconstruction quality. In addition, LiteV AE provides better reconstruction quality when compared with a V AE of comparable complexity. We also perform an analysis on the latent space learned by LiteV AE and show that it is similar to that of a regular V AE.

To summarize, our main contributions in this paper are as follows: (i) We introduce LiteV AE, a more efficient and lightweight VAE for LDMs with similar reconstruction quality. This leads to faster training of the autoencoder and higher throughput when training latent diffusion models. (ii) We explore the design space of LiteVAE and propose variations that further enhance reconstruction quality and improve its training dynamics. (iii) We perform extensive experimental analyses on the design choices and computational efficiency of LiteV AE and empirically verify its superior compute efficiency compared to a regular V AE.

2 Related work

Diffusion models and LDMs Score-based diffusion models [64, 65, 24, 66] are a class of generative models that learn the data distribution by reversing a forward destruction process that gradually adds Gaussian noise to the data. These models have recently achieved state-of-the-art generation performance on a number of diverse tasks, including unconditional and conditional image generation [46, 10, 32], text-to-image synthesis [53, 60, 55, 1, 13], video generation [4, 3, 19], image-to-image translation [59, 39], and audio generation [7, 36, 28].

While diffusion models were originally proposed for operating in the ambient image space, Rombach et al. [55] advocated for following the same methodology in the latent space of a frozen, pre-trained VAE. Following this, a number of advancements have been proposed to enhance latent diffusion models, including architecture improvements [48, 16, 33], training setups [21, 32], and sampling techniques [23, 25, 58]. In contrast to these proposed methods, our work focuses on the first stage of LDMs and aims at improving the architecture and efficiency of the VAE component.

Zhu et al. [76] recently proposed an improved decoder for the Stable Diffusion VAE that better preserves the details of conditional inputs for tasks such as in-painting. In contrast, our focus in this paper is mainly on the efficiency and properties of the entire VAE in LDMs, and our method is not restricted to conditional scenarios. Dai et al. [9] also introduced FFT features as input to the V AE for better reconstruction quality. However, their work does not address efficiency, and it can be seen as complementary to ours since FFT features can be combined with our DWT approach to further refine the encoder's initial representation.

Wavelet transform The wavelet transformation [5, 42] is a classic spatial-frequency decomposition of a signal that has gained popularity in numerous computer vision tasks, including denoising [6, 45], image and video compression [62, 67, 54, 41], super-resolution [18, 27], and image restoration [14, 40, 73]. More recently, wavelets have been integrated into generative adversarial networks [15, 71] and pixel-space diffusion models for high-resolution image synthesis [26, 49, 20]. Building on these advancements, we investigate the use of DWT to enhance the efficiency and characteristics of VAEs in LDMs, addressing an underexplored area in the literature.

3 Background

This section includes a brief overview of deep autoencoders and the wavelet transform. A summary of diffusion models is given in Appendix C.

Deep autoencoders Deep autoencoders consist of an encoder network E that maps an image to a latent representation and a decoder D that reconstructs the data from the latent code. More specifically, given an input image x ∈ R H × W × 3 , convolutional autoencoders aim to find a latent vector E ( x ) ∈ R H/f × W/f × n z such that D ( E ( x )) ≈ x , where f is the spatial downsampling scale and n z is the number of latent channels.

The training of autoencoders mainly consists of a reconstruction loss L recon ( D ( E ( x )) , x ) between the input image and the reconstructed image, and a regularization term L reg ( E ( x )) on the latents. L recon is typically a combination of ℓ 1 and perceptual loss [75], and the regularization L reg can be enforced via Kullback-Leibler (KL) divergence [35] relative to a reference distribution, typically the standard Gaussian. The regularization term forces the latent space to have a better structure for other applications, such as generative modeling. Following Esser et al. [12], it is also common to train a discriminator D with an adversarial loss L adv that differentiates the real images x from the reconstructions D ( E ( x )) for more photorealistic outputs. The overall training loss is then equal to

where the λ 's are weighting hyperparameters. Esser et al. [12] also proposed an adaptive weighting strategy for λ adv given by

for a small δ > 0 to balance the relative gradient norm of the adversarial loss with that of the reconstruction loss.

Discrete wavelet transform Wavelet transforms are a signal processing technique for extracting spatial-frequency information from input data. Wavelets are characterized by a low-pass filter L and a high-pass filter H . For 2D signals, four filters are defined via LL ⊤ , LH ⊤ , HL ⊤ , and HH ⊤ . Given an input image x , the 2D wavelet transform decomposes x into a low-frequency sub-band x L and three high-frequency sub-bands { x H , x V , x D } capturing horizontal, vertical, and diagonal details. For an image of size H × W , each wavelet sub-band is of size H/ 2 × W/ 2 . Multi-resolution analysis is achievable by iteratively applying the wavelet transform to x L at each level. Wavelet transforms are also invertible, and one can reconstruct the original image x from the sub-bands { x L , x H , x V , x D } using the inverse wavelet transform. Additionally, the Fast Wavelet Transform (FWT) [44] enables the computation of wavelet sub-bands with linear complexity relative to the number of pixels in x . Consistent with the recent literature [15, 49], we use Haar basis as the wavelet filter.

4 Method

In this section, we describe our design of a more efficient VAE for LDMs and discuss our modifications to the network architectures and the training setup that lead to better reconstruction quality and training efficiency. To motivate our approach, Figure 2 shows that when visualizing the latent code learned by the Stable Diffusion V AE (SD-V AE), the code is itself image-like, with a strong similarity to the input. This observation leads us to explore whether the learning of these latent representations can be simplified by applying a fast image-processing function to the input images prior to encoding. We opt for the discrete wavelet

(a) Input image (b) Latent code Figure 2: RGB visualization of the first three channels of a SD-VAE latent code.

transform (DWT) as the image-processing function due to its image-like structure, proven effectiveness in extracting rich, compact features from images, and wide applicability in image-processing tasks such as image compression.

4.1 Model design

We now propose LiteVAE, a wavelet-based autoencoder that reaches the reconstruction quality of standard VAEs with much lower complexity. Our method consists of three main components (see also Figure 1):

Wavelet processing: Each image x is first processed via a multi-level DWT to get the corresponding wavelet coefficients { x l L , x l H , x l V , x l D } at level l . To achieve an 8 × downsampling, we use three wavelet levels (i.e., l ∈ { 1 , 2 , 3 } ). These features extract multiscale information from x .

Feature extraction and aggregation: The wavelet coefficients { x l L , x l H , x l V , x l D } are then separately processed via a feature-extraction module F l to compute a multiscale set of feature maps F l ( { x l L , x l H , x l V , x l D } ) . The features are then combined via a feature-aggregation module F agg that takes in the output of each F l and computes the latent z . We use a UNet-based architecture similar to the ADM model [10] without spatial down/upsampling layers for feature extraction and aggregation. (See Appendix B for a discussion of the importance of these learned modules.)

Image reconstruction: Finally, a decoder network D processes the latent code z and computes the reconstructed image ˆ x = D ( z ) . We use the same decoder network as in SD-VAE for D .

The model is then trained end-to-end to learn the parameters of {F l } , F agg, and D . Because different wavelet levels already contain enough information about the images, we can use lightweight networks for the feature extraction and aggregation steps. Hence, LiteVAE essentially combines the computational benefits of DWT with the expressiveness of a learned encoder. Please refer to Appendices F and G for implementation details.

4.2 Self-modulated convolution

In addition to improving the encoder, we observe that the intermediate feature maps learned by the decoder are relatively imbalanced, with certain areas having significantly stronger magnitudes. An example of this issue is shown in Figure 3. Consistent with Karras et al. [30], we argue that this issue

Group Norm

SMC

Group Norm

SMC

O

O

O

is due to excessive group normalization layers [69] in the decoder architectures typically used in autoencoders, since such layers potentially destroy any information found in the magnitudes of the features relative to each other [31].

We propose a modified version of modulated convolution [31] instead of group normalization to avoid imbalances. Instead of modulating the convolution layers via a data-dependent style vector, we allow the convolution layer to learn the corresponding scales for each feature map. We call this operation self-modulated convolution (SMC). SMC modifies the convolution weights w ijk according to

for ϵ > 0 , where s i is a learnable parameter, and { i, j, k } spans the input feature maps, output feature maps, and the spatial dimension of the convolution. Our experiments show that using SMC in the decoder balances the feature maps and also improves the final reconstruction quality due to better training dynamics. Two examples of the decoder feature maps after using SMC are shown in Figure 3.

4.3 Training improvements

Besides the network architecture, we also introduce the following modifications that further enhance the training dynamics and reconstruction quality of LiteVAE. We verify the effect of these modifications in Sections 5 and 6.

Training resolution While the autoencoders in LDMs are typically trained on 256 × 256 data (similar to SD-V AE), we observe that the bulk of the training of LiteV AE can be effectively conducted at a lower 128 × 128 resolution. Our experiment suggests that pretraining at this lower resolution followed by a fine-tuning stage at the full resolution achieves similar reconstruction quality while requiring significantly less compute for most of the training. We later show in Appendix D.8 that this improvement is also generally applicable to the standard V AE models.

Improving the adversarial setup We replace the PatchGAN discriminator used in Stable Diffusion with a UNet-based model for pixel-wise discrimination [61]. We also notice that the adaptive weight (Equation (2)) for the adversarial loss update does not introduce any benefit and can be removed for more stable training, especially in mixed-precision setups.

Additional loss functions We also introduce two high-frequency reconstruction loss terms based on the wavelet transform and Gaussian blurring [74]. Let x be the input image and ˆ x the corresponding reconstruction. For the wavelet term, we compute the Charbonnier loss [2] between the highfrequency DWT sub-bands { x H , x V , x D } and { ˆ x LH , ˆ x HL , ˆ x HH } . For the Gaussian loss, given a Gaussian filter h , we compute the ℓ 1 loss between x -h ( x ) and ˆ x -h (ˆ x ) .

(a) Input

(b) Reconstruction

(c) Latent

OM

5 Experiments

This section presents a comprehensive empirical evaluation of LiteVAE, demonstrating its superior trade-off between computational efficiency and quality relative to standard V AEs. We further explore the properties of LiteV AE along with the changes proposed in Section 4. For each experiment, all models in comparison are trained with the exact same training setup, including the loss functions and the discriminator, to ensure a fair comparison.

Evaluation metrics We follow the same evaluation pipeline as in Rombach et al. [55] and use reconstruction Fréchet Inception Distance (rFID) [22] as the main metric to measure the quality and realism of autoencoder outputs due to its alignment with human judgment. For completeness, we also report PSNR, SSIM, and LPIPS [75]. As FID is sensitive to small implementation details [47], we recompute the metrics as much as possible based on released checkpoints to have a fair comparison between different models.

Main results We first demonstrate that LiteV AE matches or exceeds the performance of standard VAEs across various datasets and latent dimensions, as shown in Table 1. Notably, the model employed for this table utilizes approximately one-sixth of the encoder parameters compared to the VAE model (6.75M vs 34.16M) and hence trains faster. Also, one example of the reconstruction quality and the learned latent representation by LiteV AE is given in Figure 4. We notice that LiteV AE maintains the image-like latent codes, similar to the SD-V AE latent in Figure 2.

n

(a) Scaling LiteV AE (

z

= 12

for all models)

(b) Downscaling the VAE ( n z = 4 for all models)

Increasing model complexity In Table 2 we show the scalability of LiteV AE as we increase the complexity of the feature-extraction and feature-aggregation blocks. We note that the reconstruction performance strictly improves by using more encoder parameters, and our large models outperform a standard VAE of similar complexity across all metrics. Hence, we conclude that LiteVAE offers superior scalability w.r.t. the model size.

Scaling down the encoder in VAEs Table 2b also indicates that the naïve approach of scaling down the encoder in standard VAEs does not perform on par with our method in terms of reconstruction quality. Thus, we conclude that LiteVAE takes better advantage of the encoder parameters than normal VAEs, mainly due to the wavelet processing step that provides the encoder with a rich representation from the beginning.

Computational cost Table 3 presents a comparison of the computational costs between LiteVAE and the Stable Diffusion VAE encoder. LiteVAE-B requires considerably less GPU memory and offers nearly double the throughput. This reduction in computational complexity allows the usage of larger batch sizes when training the autoencoder, as shown to be beneficial by Podell et al. [50], and leads to better hardware utilization for diffusion training in the second stage of LDMs since fewer resources should be devoted to computing the latent input for the diffusion model.

Removing group normalization in the decoder We qualitatively showed in Figure 3 that group normalization in the decoder causes imbalanced feature maps in the network and that SMC can remove such artifacts. Here we also quantitatively show in Table 4 that replacing group normalization with SMC leads to better reconstruction quality. Additionally, we demonstrate in Appendix D.7 that removing the imbalanced feature maps results in less scale dependency in the final model.

Training resolution We next demonstrate the feasibility of pretraining LiteV AE at a lower resolution of 128 × 128 followed by a fine-tuning step on 256 × 256 images. To illustrate this, we compare a model trained for 150k steps at full resolution (256-full) with one trained for 100k steps at 128 and an additional 50k steps at 256 (128-tuned). As shown in Table 5, the 128-tuned model even slightly outperforms the model fully trained at the higher resolution. We also note that fine-tuning is essential, as the model trained solely on 128 × 128 images for 150k steps (128-full) performs worse than the other two. This experiment implies that the model can learn most of the semantics at lower resolutions and recover additional higher-frequency contents in the fine-tuning stage. This pretraining technique reduced the overall wall-clock time of our training runs at 256 × 256 resolution by more than a factor of two.

Scale dependency Figure 5 demonstrates that compared to the standard VAEs, LiteVAE is less prone to performance degradation when evaluating the model at different resolutions. We hypothesize

5

6

7

8

9

0

10

20

30

40

50

Resolution (

log

2

)

rFID

5

6

7

8

9

20

25

30

Resolution (

log

2

)

PSNR

5

6

7

8

9

0

.

7

0

.

8

0

.

9

Resolution (

log

2

)

SSIM

5

6

7

8

9

0

.

1

0

.

15

0

.

2

0

.

25

Resolution (

log

2

)

LPIPS

LiteVAE

VAE

that as our model learns features on top of multi-resolution wavelet coefficients, it is able to learn more scale-independent features compared to a standard encoder and leave the specific details of each scale to the initial wavelet processing step.

Analysis of the LiteVAE latent space We also analyzed the characteristics of the latent space of LiteVAE. Qualitative inspection of Figures 2 and 4, which are representative of the results that hold across our data, show that our latent space and SD-VAE share a similar image-like structure. Separately, we also examined the statistical distance between our model's latent space and pure Gaussian noise. The intuition here is that, since a diffusion model will have to form a path from pure Gaussian noise to our model's latent space, we do not want that path to be longer than the path a diffusion model has to form between Gaussian noise and the Stable Diffusion latent space. To this end, we compute the maximum mean discrepancy (MMD) [17] between latent codes from LiteVAE and samples from a standard Gaussian and compare the result with that observed for the SD-VAE (See Table 6). Here the MMD serves as a proxy measure for the path length between these distributions. In all tested cases, over a variety of RBF kernel bandwidths, our latent space is closer to Gaussian noise than that of SD-VAE.

Lastly, we trained two diffusion models on the FFHQ and CelebA-HQ datasets and compared their performance with standard VAE-based LDMs. The diffusion model architecture used for this experiment is a UNet identical to the original model from Rombach et al. [55]. Table 7 shows that the diffusion models trained in the latent space of LiteV AE perform similarly to (or slightly better than) the standard LDMs. Additionally, Figure 6 includes some generated examples from our FFHQ model. These results suggest that diffusion models are also capable of modeling the latent space of LiteV AE.

40

60

80

100

120

140

160

180

200

10

0

10

1

10

2

10

3

10

4

Training Steps (

×

10

3

)

Constant

Adaptive

6 Ablation studies

We next present our main ablation studies to determine the individual impact of the changes proposed in Section 4. We use the ImageNet 128 × 128 model with a latent size of 32 × 32 × 12 as the baseline for all ablations. Further ablation studies on other design choices in LiteV AE are provided in Appendix D.

Removing adaptive weight for λ reg Table 8 demonstrates that we can safely remove the adaptive weight for the adversarial loss (Equation (2)) and still slightly improve the metrics. Figure 7 also shows the relative norm of the gradient of the adversarial loss compared to the reconstruction loss for both adaptive and constant λ adv. We observe that using adaptive λ adv leads to more imbalanced gradient ratios, and hence less stable training, especially for mixed-precision scenarios. Accordingly, we exclusively use a constant weight for the adversarial loss in our experiments.

High-frequency loss functions Table 9 shows the effect of adding high-frequency losses based on Gaussian filtering and the wavelet transform. The addition of these high-frequency loss terms during training consistently improves all reconstruction metrics.

Choice of the discriminator We finally show that using a UNet-based discriminator [61] outperforms both PatchGAN and StyleGAN discriminators used in previous works [55, 72] in terms of rFID while having comparable performance for other metrics. We also empirically noted that using a UNet discriminator resulted in more stable training across different runs and hyperparameters. The full comparison for this experiment is given in Table 10.

7 Conclusion

In this paper, we presented LiteVAE, a new design concept for autoencoders based on the multiresolution wavelet transform. LiteV AE can match the performance of standard V AEs while requiring significantly less compute. We also analyzed the design space and training of this proposed family of autoencoders and offered several modifications that further improve the final reconstruction quality and training dynamics of the base model. Overall, LiteVAE offers more flexibility in terms of performance/compute trade-off and outperforms the naïve approach of making the VAE encoder smaller. Our current work is focused on improving efficiency in the models responsible for encoding the latent representation of natural images, and whether the efficiency benefits of LiteV AE extend to other domains is a question we leave to follow-up work. Although we introduced LiteVAE in the context of LDMs, we hypothesize that its application is not confined to this scenario. We consider the extension of LiteVAE to other autoencoder-based generative modeling schemes (e.g., tokenization) a promising avenue for further research.

References

Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, Bryan Catanzaro, Tero Karras, and Ming-Yu Liu. ediff-i: Text-to-image diffusion models with an ensemble of expert denoisers. CoRR , abs/2211.01324, 2022. doi: 10.48550/arXiv.2211.01324. URL https://doi.org/10.48550/arXiv.2211.01324.

Jonathan T Barron. A general and adaptive robust loss function. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 4331-4339, 2019.

Andreas Blattmann, Tim Dockhorn, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Dominik Lorenz, Yam Levi, Zion English, Vikram Voleti, Adam Letts, Varun Jampani, and Robin Rombach. Stable video diffusion: Scaling latent video diffusion models to large datasets. CoRR , abs/2311.15127, 2023. doi: 10.48550/ARXIV.2311.15127. URL https://doi.org/10.48550/arXiv.2311.15127.

Andreas Blattmann, Robin Rombach, Huan Ling, Tim Dockhorn, Seung Wook Kim, Sanja Fidler, and Karsten Kreis. Align your latents: High-resolution video synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 22563-22575, 2023.

M. E. Brewster. An introduction to wavelets (charles k. chui). SIAM Rev. , 35(2):312-313, 1993. doi: 10.1137/1035061. URL https://doi.org/10.1137/1035061.

S Grace Chang, Bin Yu, and Martin Vetterli. Adaptive wavelet thresholding for image denoising and compression. IEEE transactions on image processing , 9(9):1532-1546, 2000.

Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, and William Chan. Wavegrad: Estimating gradients for waveform generation. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL https://openreview.net/forum?id=NsMLjcFaO8O.

Xiaojie Chu, Liangyu Chen, and Wenqing Yu. Nafssr: Stereo image super-resolution using nafnet. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops , pages 1239-1248, June 2022.

Xiaoliang Dai, Ji Hou, Chih-Yao Ma, Sam Tsai, Jialiang Wang, Rui Wang, Peizhao Zhang, Simon Vandenhende, Xiaofang Wang, Abhimanyu Dubey, et al. Emu: Enhancing image generation models using photogenic needles in a haystack. arXiv preprint arXiv:2309.15807 , 2023.

Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat gans on image synthesis. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 8780-8794, 2021. URL https://proceedings.neurips.cc/paper/2021/ hash/49ad23d1ec9fa4bd8d77d02681df5cfa-Abstract.html.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly,

Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL https://openreview.net/ forum?id=YicbFdNTTy.

Patrick Esser, Robin Rombach, and Björn Ommer. Taming transformers for highresolution image synthesis. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021 , pages 12873-12883. Computer Vision Foundation / IEEE, 2021. doi: 10.1109/CVPR46437.2021.01268. URL https://openaccess.thecvf.com/content/CVPR2021/html/Esser_Taming_Transformers_ for_High-Resolution_Image_Synthesis_CVPR_2021_paper.html.

Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, Kyle Lacey, Alex Goodwin, Yannik Marek, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis. CoRR , abs/2403.03206, 2024. doi: 10.48550/ARXIV.2403.03206. URL https://doi.org/10.48550/arXiv.2403.03206.

Mário AT Figueiredo and Robert D Nowak. An em algorithm for wavelet-based image restoration. IEEE Transactions on Image Processing , 12(8):906-916, 2003.

Rinon Gal, Dana Cohen Hochberg, Amit Bermano, and Daniel Cohen-Or. SWAGAN: a stylebased wavelet-driven generative model. ACM Trans. Graph. , 40(4):134:1-134:11, 2021. doi: 10.1145/3450626.3459836. URL https://doi.org/10.1145/3450626.3459836.

Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is a strong image synthesizer. CoRR , abs/2303.14389, 2023. doi: 10.48550/arXiv.2303.14389. URL https://doi.org/10.48550/arXiv.2303.14389.

Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research , 13(1):723-773, 2012.

Tiantong Guo, Hojjat Seyed Mousavi, Tiep Huu Vu, and Vishal Monga. Deep wavelet prediction for image super-resolution. In 2017 IEEE Conference on Computer Vision and Pattern Recognition Workshops, CVPR Workshops 2017, Honolulu, HI, USA, July 21-26, 2017 , pages 1100-1109. IEEE Computer Society, 2017. doi: 10.1109/CVPRW.2017.148. URL https://doi.org/10.1109/CVPRW.2017.148.

Agrim Gupta, Lijun Yu, Kihyuk Sohn, Xiuye Gu, Meera Hahn, Li Fei-Fei, Irfan Essa, Lu Jiang, and José Lezama. Photorealistic video generation with diffusion models. arXiv preprint arXiv:2312.06662 , 2023.

Florentin Guth, Simon Coste, Valentin De Bortoli, and Stéphane Mallat. Wavelet score-based generative modeling. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022 , 2022. URL http://papers.nips.cc/paper_files/paper/2022/ hash/03474669b759f6d38cdca6fb4eb905f4-Abstract-Conference.html.

Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Baining Guo. Efficient diffusion training via min-snr weighting strategy. In IEEE/CVF International Conference on Computer Vision, ICCV 2023, Paris, France, October 1-6, 2023 , pages 74077417. IEEE, 2023. doi: 10.1109/ICCV51070.2023.00684. URL https://doi.org/10.1109/ ICCV51070.2023.00684.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA , pages 6626-6637, 2017. URL https://proceedings.neurips.cc/paper/ 2017/hash/8a1d694707eb0fefe65871369074926d-Abstract.html.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. CoRR , abs/2207.12598, 2022. doi: 10.48550/arXiv.2207.12598. URL https://doi.org/10.48550/arXiv.2207.12598.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020. URL https://proceedings.neurips.cc/paper/2020/hash/ 4c5bcfec8584af0d967f1ab10179ca4b-Abstract.html.

Susung Hong, Gyuseong Lee, Wooseok Jang, and Seungryong Kim. Improving sample quality of diffusion models using self-attention guidance. CoRR , abs/2210.00939, 2022. doi: 10.48550/ arXiv.2210.00939. URL https://doi.org/10.48550/arXiv.2210.00939.

Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. CoRR , abs/2301.11093, 2023. doi: 10.48550/arXiv.2301.11093. URL https://doi.org/10.48550/arXiv.2301.11093.

Huaibo Huang, Ran He, Zhenan Sun, and Tieniu Tan. Wavelet-srnet: A wavelet-based CNN for multi-scale face super resolution. In IEEE International Conference on Computer Vision, ICCV 2017, Venice, Italy, October 22-29, 2017 , pages 1698-1706. IEEE Computer Society, 2017. doi: 10.1109/ICCV.2017.187. URL https://doi.org/10.1109/ICCV.2017.187.

Qingqing Huang, Daniel S. Park, Tao Wang, Timo I. Denk, Andy Ly, Nanxin Chen, Zhengdong Zhang, Zhishuai Zhang, Jiahui Yu, Christian Havnø Frank, Jesse H. Engel, Quoc V. Le, William Chan, and Wei Han. Noise2music: Text-conditioned music generation with diffusion models. CoRR , abs/2302.03917, 2023. doi: 10.48550/arXiv.2302.03917. URL https://doi.org/10.48550/ arXiv.2302.03917.

Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings . OpenReview.net, 2018. URL https://openreview.net/forum?id=Hk99zCeAb.

Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019 , pages 4401-4410. Computer Vision Foundation / IEEE, 2019. doi: 10.1109/CVPR.2019. 00453. URL http://openaccess.thecvf.com/content_CVPR_2019/html/Karras_A_Style-Based_ Generator_Architecture_for_Generative_Adversarial_Networks_CVPR_2019_paper.html.

Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020 , pages 8107-8116. Computer Vision Foundation / IEEE, 2020. doi: 10.1109/CVPR42600.2020. 00813. URL https://openaccess.thecvf.com/content_CVPR_2020/html/Karras_Analyzing_and_ Improving_the_Image_Quality_of_StyleGAN_CVPR_2020_paper.html.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. 2022. URL https://openreview.net/forum?id= k7FuTOWMOc7.

Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models, 2023.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings , 2015. URL http://arxiv.org/abs/1412.6980.

Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. 2014. URL http: //arxiv.org/abs/1312.6114.

Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL https://openreview.net/forum?id=a-xFK8Ymz5J.

Jie Liang, Hui Zeng, and Lei Zhang. Details or artifacts: A locally discriminative learning approach to realistic image super-resolution. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , 2022.

Chen-Hsuan Lin, Jun Gao, Luming Tang, Towaki Takikawa, Xiaohui Zeng, Xun Huang, Karsten Kreis, Sanja Fidler, Ming-Yu Liu, and Tsung-Yi Lin. Magic3d: High-resolution text-to-3d content creation. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2023.

Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A. Theodorou, Weili Nie, and Anima Anandkumar. I 2 sb: Image-to-image schrödinger bridge. CoRR , abs/2302.05872, 2023. doi: 10.48550/arXiv.2302.05872. URL https://doi.org/10.48550/arXiv.2302.05872.

Pengju Liu, Hongzhi Zhang, Kai Zhang, Liang Lin, and Wangmeng Zuo. Multi-level wavelet-cnn for image restoration. In 2018 IEEE Conference on Computer Vision and Pattern Recognition Workshops, CVPR Workshops 2018, Salt Lake City, UT, USA, June 1822, 2018 , pages 773-782. Computer Vision Foundation / IEEE Computer Society, 2018. doi: 10.1109/CVPRW.2018.00121. URL http://openaccess.thecvf.com/content_cvpr_2018_ workshops/w13/html/Liu_Multi-Level_Wavelet-CNN_for_CVPR_2018_paper.html.

Haichuan Ma, Dong Liu, Ning Yan, Houqiang Li, and Feng Wu. End-to-end optimized versatile image compression with wavelet-like transform. IEEE Transactions on Pattern Analysis and Machine Intelligence , 44(3):1247-1263, 2020.

Stéphane Mallat. A theory for multiresolution signal decomposition: The wavelet representation. IEEE Trans. Pattern Anal. Mach. Intell. , 11(7):674-693, 1989. doi: 10.1109/34.192463. URL https://doi.org/10.1109/34.192463.

Stéphane Mallat. A wavelet tour of signal processing . Elsevier, 1999.

Stephane G Mallat. A theory for multiresolution signal decomposition: the wavelet representation. IEEE transactions on pattern analysis and machine intelligence , 11(7):674-693, 1989.

S Kother Mohideen, S Arumuga Perumal, and M Mohamed Sathik. Image de-noising using discrete wavelet transform. International Journal of Computer Science and Network Security , 8 (1):213-216, 2008.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 8162-8171. PMLR, 2021. URL http: //proceedings.mlr.press/v139/nichol21a.html.

Gaurav Parmar, Richard Zhang, and Jun-Yan Zhu. On aliased resizing and surprising subtleties in GAN evaluation. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022 , pages 11400-11410. IEEE, 2022. doi: 10.1109/CVPR52688.2022.01112. URL https://doi.org/10.1109/CVPR52688.2022.01112.

William Peebles and Saining Xie. Scalable diffusion models with transformers. CoRR , abs/2212.09748, 2022. doi: 10.48550/arXiv.2212.09748. URL https://doi.org/10.48550/arXiv. 2212.09748.

Hao Phung, Quan Dao, and Anh Tran. Wavelet diffusion models are fast and scalable image generators. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2023, Vancouver, BC, Canada, June 17-24, 2023 , pages 10199-10208. IEEE, 2023. doi: 10.1109/CVPR52729.2023.00983. URL https://doi.org/10.1109/CVPR52729.2023.00983.

Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. SDXL: improving latent diffusion models for high-resolution image synthesis. CoRR , abs/2307.01952, 2023. doi: 10.48550/ARXIV.2307.01952. URL https://doi.org/10.48550/arXiv.2307.01952.

Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023 . OpenReview.net, 2023. URL https://openreview.net/pdf? id=FjNys5c7VyY.

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , volume 139 of Proceedings of Machine Learning Research , pages 8821-8831. PMLR, 2021. URL http://proceedings.mlr.press/v139/ramesh21a.html.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with CLIP latents. CoRR , abs/2204.06125, 2022. doi: 10.48550/arXiv.2204.06125. URL https://doi.org/10.48550/arXiv.2204.06125.

Oren Rippel and Lubomir Bourdev. Real-time adaptive image compression. In International Conference on Machine Learning , pages 2922-2930. PMLR, 2017.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 1824, 2022 , pages 10674-10685. IEEE, 2022. doi: 10.1109/CVPR52688.2022.01042. URL https://doi.org/10.1109/CVPR52688.2022.01042.

Negar Rostamzadeh, Emily Denton, and Linda Petrini. Ethics and creativity in computer vision. CoRR , abs/2112.03111, 2021. URL https://arxiv.org/abs/2112.03111.

Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael S. Bernstein, Alexander C. Berg, and Li FeiFei. Imagenet large scale visual recognition challenge. Int. J. Comput. Vis. , 115(3):211-252, 2015. doi: 10.1007/s11263-015-0816-y. URL https://doi.org/10.1007/s11263-015-0816-y.

Seyedmorteza Sadat, Jakob Buhmann, Derek Bradley, Otmar Hilliges, and Romann M. Weber. CADS: Unleashing the diversity of diffusion models through condition-annealed sampling. In The Twelfth International Conference on Learning Representations , 2024. URL https: //openreview.net/forum?id=zMoNrajk2X.

Chitwan Saharia, William Chan, Huiwen Chang, Chris A. Lee, Jonathan Ho, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models. In Munkhtsetseg Nandigjav, Niloy J. Mitra, and Aaron Hertzmann, editors, SIGGRAPH '22: Special Interest Group on Computer Graphics and Interactive Techniques Conference, Vancouver, BC, Canada, August 7 - 11, 2022 , pages 15:1-15:10. ACM, 2022. doi: 10.1145/3528233.3530757. URL https://doi.org/10.1145/3528233.3530757.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Seyed Kamyar Seyed Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. 2022. URL http://papers.nips.cc/paper_files/paper/ 2022/hash/ec795aeadae0b7d230fa35cbaf04c041-Abstract-Conference.html.

Edgar Schonfeld, Bernt Schiele, and Anna Khoreva. A u-net based discriminator for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 8207-8216, 2020.

Ke Shen and Edward J Delp. Wavelet based rate scalable video compression. IEEE transactions on circuits and systems for video technology , 9(1):109-122, 1999.

Wenzhe Shi, Jose Caballero, Ferenc Huszár, Johannes Totz, Andrew P Aitken, Rob Bishop, Daniel Rueckert, and Zehan Wang. Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 1874-1883, 2016.

Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. 37:2256-2265, 2015. URL http://proceedings.mlr.press/v37/sohl-dickstein15.html.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d'Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada , pages 11895-11907, 2019. URL https://proceedings. neurips.cc/paper/2019/hash/3001ef257407d5a371a96dcd947c7d93-Abstract.html.

Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL https://openreview.net/forum?id=PxTIG12RRHS.

David S. Taubman and Michael W. Marcellin. JPEG2000 - image compression fundamentals, standards and practice , volume 642 of The Kluwer international series in engineering and computer science . Kluwer, 2002. ISBN 978-0-7923-7519-7. doi: 10.1007/978-1-4615-0799-4. URL https://doi.org/10.1007/978-1-4615-0799-4.

Xintao Wang, Ke Yu, Shixiang Wu, Jinjin Gu, Yihao Liu, Chao Dong, Yu Qiao, and Chen Change Loy. Esrgan: Enhanced super-resolution generative adversarial networks. In The European Conference on Computer Vision Workshops (ECCVW) , September 2018.

Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV) , pages 3-19, 2018.

Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Yingxia Shao, Wentao Zhang, Ming-Hsuan Yang, and Bin Cui. Diffusion models: A comprehensive survey of methods and applications. CoRR , abs/2209.00796, 2022. doi: 10.48550/arXiv.2209.00796. URL https://doi.org/10.48550/arXiv.2209.00796.

Mengping Yang, Zhe Wang, Ziqiu Chi, and Yanbing Zhang. Fregan: Exploiting frequency components for training gans under limited data. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 -December 9, 2022 , 2022. URL http://papers.nips.cc/paper_files/paper/2022/hash/ d804cef41362be39d3972c1a71cfc4e9-Abstract-Conference.html.

Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. Vector-quantized image modeling with improved VQGAN. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=pfNyExj7z2.

Yingchen Yu, Fangneng Zhan, Shijian Lu, Jianxiong Pan, Feiying Ma, Xuansong Xie, and Chunyan Miao. Wavefill: A wavelet-based generation network for image inpainting. In Proceedings of the IEEE/CVF international conference on computer vision , pages 1411414123, 2021.

Eduard Zamfir, Marcos V Conde, and Radu Timofte. Towards real-time 4k image superresolution. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 1522-1532, 2023.

Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018 , pages 586-595. Computer Vision Foundation / IEEE Computer Society, 2018. doi: 10.1109/CVPR.2018.00068. URL http://openaccess.thecvf.com/content_cvpr_2018/html/ Zhang_The_Unreasonable_Effectiveness_CVPR_2018_paper.html.

Zixin Zhu, Xuelu Feng, Dongdong Chen, Jianmin Bao, Le Wang, Yinpeng Chen, Lu Yuan, and Gang Hua. Designing a better asymmetric VQGAN for stablediffusion. CoRR , abs/2306.04632, 2023. doi: 10.48550/ARXIV.2306.04632. URL https://doi.org/10.48550/arXiv.2306.04632.

A Broader impact statement

Our work can significantly reduce the training time and memory requirements of autoencoders in latent diffusion models (LDMs). Given the rising popularity of LDMs, our approach holds promise for positive environmental impacts and significant advancements in generative modeling. It is important to note that while AI-generated content can enhance productivity and creativity, we must remain mindful of the potential risks and ethical concerns involved. For a deeper discussion of ethics and creativity in computer vision, readers are directed to [56].

B Using a non-learned encoder

This section provides further motivation behind the design of LiteVAE. We investigate using a non-learned (i.e., fixed) encoder in two settings: (1) for simple datasets such as FFHQ [30] and (2) for more diverse datasets such as ImageNet [57]. We use the reconstruction FID (rFID) [22] as our measure of reconstruction quality, aiming to achieve the SD-VAE's downsampling factor of f = 8 . The analysis leads to two observations. First, the non-learned autoencoder (although efficient) can provide high-quality reconstructions only if we use a larger channel depth for the encoder network compared to SD-VAE. Secondly, the dense latent space learned by the autoencoder provides a better structure for generative modeling. LiteV AE essentially combines the computational benefits of the non-learned encoder with the learned latent space of a regular V AE.

Simple datasets For simpler datasets like FFHQ, it is possible to completely replace the encoder E with a predefined function and get similar reconstruction quality. In our case, we used a three-level DWT and only kept the sub-bands of the lowest level. We then trained a decoder to convert the lowest-level sub-bands back to the image. Table 11 shows the results of this approach on two relatively restricted datasets. We observe that this wavelet representation offers a similar reconstruction quality to a learned encoder while reducing the number of encoder parameters from

about 34M to zero. This experiment indicates that with the help of rich image representations from the wavelet transform, we can speed up SD-VAE by reducing the complexity of the encoder.

Complex datasets The next step is to explore whether this non-learned encoder setup is scalable to more diverse datasets such as ImageNet. Table 12 demonstrates that while the nonlearned encoder is effective in simpler scenarios, it falls short of the quality of normal V AEs in more complex settings. This indicates that the information present in the higher frequency sub-bands of the wavelet transform is essential for the decoder to reconstruct more diverse images with higher quality. To

validate this hypothesis, we incorporate the information from higher frequency sub-bands via a space-to-depth operation [63] in the encoder and observe that we can recover the high reconstruction quality of the learned encoder (DWT-2 in Table 12). However, this approach is not preferable because the channel dimension is now too high for generative modeling.

Importance of having a learned latent space Finally, we demonstrate that although it is possible to completely replace the encoder of the V AE with a non-learned waveletbased latent representation for the FFHQ dataset, the learned latent space in LiteV AE offers a better structure for training diffusion models. Table 13 indicates that training the diffusion model on the learned latent code of LiteVAE outperforms the non-learned DWT representation. We argue that the sparse nature of wavelets is harmful to

generation quality compared to the dense representation learned by the encoder of LiteV AE.

C Summary of diffusion models

Diffusion models learn the data distribution p data by reversing a noising process that gradually converts a data point x into random Gaussian noise. More specifically, diffusion models define a forward process via x t = x + σ ( t ) ϵ , where ϵ ∼ N (0 , I ) . Then, they train a denoiser network D θ to estimate the clean signal x from the current noisy sample x t . It has been shown that this process corresponds to the following stochastic differential equation (SDE) [66, 32]

where d ω t is the standard Wiener process, p t ( x t ) is the distribution of noisy samples at time t , and β ( t ) is a term that controls the influence of noise during the sampling process. The denoiser network D θ effectively approximates the score function ∇ x t log p t ( x t ) . Given that p 0 = p data and p 1 = N ( 0 , σ 2 max I ) , sampling new data points is then possible by starting from random Gaussian noise and solving the corresponding SDE reverse in time.

Latent diffusion models [55] follow the same methodology, but instead of performing the forward and backward process in the pixel space, they first convert the data into the latent codes via a pretrained VAE and employ the diffusion process in the latent space. Please refer to Karras et al. [32] and Yang et al. [70] for more details on diffusion models.

D Additional ablation studies

This section contains additional ablation studies on the design space and training dynamic of LiteV AE.

D.1 Training loss functions

We experimented with the following changes to the loss functions during training of the autoencoder to measure whether they lead to any improvement in reconstruction quality.

Changing the VGG loss Wang et al. [68] proposed a different VGG loss function based on the features before the activation layer. We also ablated this choice against the standard LPIPS loss typically used in the V AEs of LDMs. Table 14 indicates that this change has a considerable boost to the reconstruction FID at the cost of lower PSNR. As the perceptual quality is more important for LDMs compared to distortion metrics, we recommend switching to this loss function instead of the LPIPS loss. However, we used the LPIPS loss for the experiments in the main text to have a similar training setup with commonly used VAEs in LDMs.

Including the locally discriminative learning (LDL) loss Liang et al. [37] introduced the LDL loss function to reduce the artifacts caused by the discriminator in the super-resolution context. We also experimented with this loss term and found that it does not have any noticeable impact on the reconstruction quality of LiteV AE, as shown in Table 15.

Choosing different adversarial loss functions We also ablated the adversarial loss function for two different setups: a hinge loss, and a non-saturating (logistic) loss. As depicted in Table 16, we observe that the hinge loss generally leads to slightly better rFID while the logistic loss achieves slightly better PSNR. Since the adversarial loss in the autoencoder training is only responsible for increasing the photorealism of the outputs, we conclude that both loss terms work equally well.

D.2 Role of the 1 × 1 convolution layers

Ramesh et al. [52] showed that using a 1 × 1 convolution after the output of the encoder and before the input of the decoder improves the approximation accuracy of the evidence lower bound (ELBO) term in the loss function. We ablated this design choice in the context of LiteV AE and found that restricting the receptive field of the latent space with these 1 × 1 convolution layers might be harmful to the reconstruction quality by enforcing too much KL regularization. Table 17 shows that removing these convolution blocks leads to much better reconstruction quality in our 256 × 256 model. Accordingly, we suggest removing these 1 × 1 convolutions from the model (or, equivalently, adjusting the weight of the KL loss) to get better reconstruction.

D.3 Different networks for feature extraction

We also experimented with NAFNet [8] instead of the UNet for extracting features from wavelet sub-bands and observed that it performs similarly to the UNet architecture mentioned in the main text. The results are given in Table 18. This experiment indicates that other network choices for the feature-extraction module are indeed possible, and LiteVAE is flexible w.r.t. this design choice. We chose the UNet to keep the setup as close as possible to the standard V AE design in LDMs.

D.4 Sharing the weights of the feature-extraction UNet

We next investigated whether a single UNet could be shared across different wavelet sub-bands to further reduce the encoder's trainable parameters. Table 19 demonstrates that it is indeed possible to share F l between different sub-bands. A shared UNet might lead to the post hoc usage of the encoder across different wavelet levels and resolutions at inference. However, as the computational cost (in terms of GFLOPS) does not change with parameter sharing, we did not use this technique for the main experiments.

D.5 Using ViT for feature aggregation

We also explore the use of non-convolutional vision transformer (ViT) blocks [11] for feature aggregation F agg . As indicated in Table 20, employing ViT achieves comparable reconstruction quality to that of a fully-convolutional encoder, but with fewer parameters. However, it is important to note that incorporating ViT makes the model resolution-dependent. This is a drawback, as the V AE in LDMs is usually required to operate on data with varying resolutions. Hence, we side with the UNet models to make the encoder resolution-independent.

5

6

7

8

0

5

10

15

Resolution (

log

2

)

rFID

5

6

7

8

22

24

26

28

30

Resolution (

log

2

)

PSNR

5

6

7

8

0

.

75

0

.

8

0

.

85

Resolution (

log

2

)

SSIM

5

6

7

8

6

·

10

-

2

8

·

10

-

2

0

.

1

0

.

12

0

.

14

0

.

16

0

.

18

Resolution (

log

2

)

LPIPS

with SMC

with Group Norm

D.6 Importance of using all wavelet levels

We also explored the possibility of performing feature extraction on only a subset of wavelet coefficients rather than across all wavelet levels. As shown in Table 21, this approach negatively impacts reconstruction performance on ImageNet, indicating that incorporating information from all wavelet levels is essential for high-quality reconstruction, particularly with complex datasets.

D.7 Scale dependency of SMC

This section shows that using SMC improves the scale dependency of LiteV AE. The results in Figure 8 indicate that using SMC instead of group normalization leads to less degradation in performance as we change the resolution of the evaluation dataset. We argue that removing the imbalanced feature maps aids the network in learning features that are less scale-dependent.

D.8 Training resolution for the standard VAEs

This experiment validates that the idea of pretraining the autoencoder at 128 × 128 followed by fine-tuning at 256 × 256 also works for the standard V AEs. The results of this experiment are given in Table 22. Similar to LiteV AE, the 128-tuned model matches the performance of the 256-full model while requiring considerably less training compute.

E Additional generated samples

Figure 9 provides additional generated samples from our latent diffusion model trained on FFHQ.

F Implementation details

All models were trained with a batch size of 16 on two GPUs until the autoencoder could produce high-quality reconstructions. The training duration was 200k steps for the ImageNet 128 × 128 models, and 100k for the ImageNet 256 × 256 and FFHQ models. We use Adam optimizer [34] with a learning rate of 10 -4 and ( β 1 , β 2 ) = (0 . 5 , 0 . 9) . The details of the model architecture for feature-extraction and feature-aggregation modules are given in Tables 23 and 24. Our implementation of the UNet used for feature extraction and aggregation closely follows the ADM model [10] without spatial down/upsampling layers. The decoder in LiteVAE exactly follows the implementation of the decoder from Stable Diffusion VAE [55], except for the SMC experiment. For training the latent diffusion and the standard VAE models, we closely follow Rombach et al. [55] to ensure a fair comparison.

G Pseudocode for different LiteVAE blocks

In this section, we present additional pseudocode for various LiteV AE components. The core element of LiteVAE is the Haar wavelet transform, which can be implemented in PyTorch as shown below:

The PyTorch implementation of the self-modulated convolution block introduced in Section 4.2 is provided below:

Next, we present the code for the residual blocks utilized in the LiteV AE UNet networks:

Additionally, the feature-extraction and feature-aggregation UNets can be implemented as follows:

The LiteVAE encoder can be implemented as shown below:

Finally, the code for LiteV AE is also provided below.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claims are supported via detailed empirical analysis.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations are discussed in Section 7.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: The paper does not include theoretical results.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The hyperparameters, algorithms, and implementation details are provided in the appendix for proper reproducibility of our work.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: While we do not provide open access to our codebase, the hyperparameters, algorithms, and implementation details are provided in the appendix to ensure reproducibility.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide the details of the experiments in the main text and the appendix.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Computing the error bar in the context of our work is too computationally demanding.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide hardware details whenever appropriate.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?

Answer: [Yes]

Justification: The paper follows the NeurIPS Code of Ethics as stated on the website.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The broader impact statements are discussed in Appendix A.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our work does not release such models or datasets.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We use well-established datasets and models with proper academic licenses.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing or research with human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing or research with human subjects.

Peer reviews:
comments: 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.

LIMITATIONS:
The limitations of their work and any potential negative societal impacts should be strengthened.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nThese 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.\n\nFinally, 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.\n\nSTRENGTHS:\nORIGINALITY:\nThe 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. \n\n\nPRESENTATION:\nThis 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. \nThe inclusion of relevant figures further improves the readability of this work.\n\n\nQUALITY:\nThis 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.\n \n\nSIGNIFICANCE:\nThe reduction in computational cost with no tradeoff in reconstruction quality points to the significance of this work.\n\nWEAKNESSES:\nTable 3 only compares the throughput of VAE and LiteVAE-B. Please provide the throughput for LiteVAE-S and LiteVAE-L.\n\nAlso, the model did not discuss the additional computational cost for the 2D wavelet transform. Please provide some information about this.\n\nThere is a typographical error in line 207. I think \"Table 3\" should be \"Table 2\".\n\nQUESTIONS:\nTable 3 only compares the throughput of VAE and LiteVAE-B. Please provide the throughput for LiteVAE-S and LiteVAE-L.\n\nAlso, the model did not discuss the additional computational cost for the 2D wavelet transform. Please provide some information about this.\n\nThere is a typographical error in line 207. I think \"Table 3\" should be \"Table 2\".\n\nLIMITATIONS:\nThe checklist points to Section 7 for the limitations but the section only shows the conclusion. Please address this.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}

comments: 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.

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-14

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\n2. 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.\n\n3. All claims are supported by experimentation and ablation studies.\n\nWEAKNESSES:\n1. 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?\n\n2. 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.\n\n3. 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.\n\n4. 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.\n\nQUESTIONS:\n1. 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?\n\n2. Can LiteVAE be used as a drop-in replacement of SD-VAE in real world applications?\n\n3. 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? \n\n4. Could you comment on whether using the wavelet basis and reconstruction loss with the original SD-VAE might yield similar performance?\n\nLIMITATIONS:\nThe limitations are not discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-26"}