LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models
Authors: Seyedmorteza Sadat
Jakob Buhmann
Derek Bradley
Otmar Hilliges
Romann M. Weber
Venue: NeurIPS 2024 poster (2024)
Keywords: variational autoencoders
latent diffusion models
Primary area: diffusion_based_models
TL;DR: We introduce an efficiency-adapted variational autoencoder design for latent diffusion models with improved scalability and computational efficiency.

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
Seyedmorteza Sadat', Jakob Buhmann', Derek Bradley', Otmar Hilliges', Romann M. Weber
'ETH Zürich, siiliggee@@nnt wttee (seyedmorteza, sadat, otmar {jakob. buhmann, derek, bradley romann. weber aeiberrrdiineeyeseeaachh
ethz. ch
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, introduce VAte new au- toencoder design for LDMs, which leverages the 2D discrete wavelet transform enhance scalability and computational efficiency over standard variational au- toencoders (VAEs) with no sacrifice in output quality. We investigate the training methodologies and the decoder architecture of LiteVAE and propose several en- hancements that improve the training dynamics and reconstruction quality. Our base VAE 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 compa- rable complexity across all evaluated metrics (rFID, LPIPS, PSNR, and SSIM).
Encoder
Feature Extraction and Aggrogation
DWT
DWT
Medllook
Helllook
BisBlock
Multi-level
Wavelet Transform
Image Reconstruction
Fi
F
F
Decxlet
Figure 1: An overview of VAE. The input image is first decomposed into multi-level wavelet coefficients, and each wavelet sub-band is separately processed via a feature-cxtraction network. The features then combined via which is then transformed back into the image space by the decoder. We use a lightweight UNet architecture (top right) without spatial down/upsanpling for feature extraction and aggregation. The decoder is a fully convolutional network similar to that in the Stable Diffusion VAE [55]. Lite design allows it to be significantly more efficient than standard VAEs in LDMs while maintaining high reconstruction quality.
aggreertag module to compute the final latent code,
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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, 33J, 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 the resource requirements of querying large encoder network for computing the latent codes. For example, as the autoencoder operates on high high-resolution images, the VAE encoder of Stable Diffusion 2.1 uses 135.59 GFLOPS compared with 86.37 for the diffusion UNet. This becomes an even greater concern for video diffusion models, as the encoder then needs to provide the latents for 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 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 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 VAE 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, new autoencoder design for LDMs, which has superior compute/quality trade-offs compared with standard VAEs.
Lite VAE consists of lightweight coefficients and a feature-aggen agggregatonn module to combine these multiscale features into a unified latent code. A decoder then converts the latent code back to image. An overview of the VAE pipeline is shown in Figure 1.
module compute features from the wavelet
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 ViE in depth and propose several variations on the network architccture and training sctup that further boost reconstruction quality and training efficiency.
Through extensive experimentation, we show that VAE considerably reduces the computational cost of the standard VAE encoder while maintaining the same level of reconstruction quality. In addition, Lite VAE provides better reconstruction quality when compared with VAE of comparable complexity. We also perform an analysis on the latent space learned by LiteVAE and show that it is similar to that of regular VAE.
To summarize, our main contributions in this paper are as follows: (i) We introduce VAE, 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 VAE 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 Lite VAE and empirically verify its superior compute efficiency compared to regular VAE.
Result of processing single 256 x 256 x image and its corresponding 32 x 32 x 4 latent representation.

---

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 forward destruction process that gradually adds Gaussian noise to the data. These models have recently achieved kfte-of-the-aa 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 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 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 input to the VAE 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 's initial representation.
Wavelet transform The wavelet transformation [5, 42] is a classic spatial-frequency decomposition of 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 & 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 æ E R x3 convolutional autoencoders aim to find a latent vector (x)) E RgH//xW/ft xn. such that D(E(x)) Ro æ. where f is the spatial downsampling scale and n. is the number of latent channels.
The training of autoencoders mainly consists of reconstruction loss frecom (D(((a)) x)) between the input image and the reconstructed image, and regularization term Lieg (E (Exx) on the latents. Crecon is typically a combination of ε and perceptual loss [75], and the regularization Lreg can be enforced via Kullback-Leibler (KL) divergence [35] relative to 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 Lad that differentiates the real images x from the reconstructions (zE())) for more photorealistic outputs. The overall training loss is then equal to
Lrain Lrecon + AregGreg + Audy
where the X's are weighting hyperparameters. Esser et al. [12] also proposed an adaptive weighting strategy for Andv given by
1 1/ C
Auav
2 | + 8
for a small 8 > 0 to balance the relative gradient norm of the adversarial loss with that of the reconstruction loss.
3

---

Discrete wavelet transform Wavelet transforms are signal processing technique for extracting frequency information from imput data. Wavelets are characterized by a low-pass filter L and high-pass filter H. For 2DD signals, four filters are defined via T.I. L.H HT, and H H T_ Given an input image æ; the 2D wavelet transform decomposes a: into a now-freuueccy sub-band æy and three high-frequency sub-bands {æ FV, DD) capturing horizontal, vertical, and diagonal details. For an image of size И x W, each wavelet sub-band of size / x W/2. Multi-resolution analysis is achievable by iteratively applying the wavelet transform 2t, each level. Wavelet transforms are also invertible, and one can reconstruct the original image a from the sub-bands {ær, HH æV, up) 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 æ. Consistent with the recent literature [15, 49], we use Haar basis the wavelet filter.
Method
In this section, we describe our design of a more efficient VAE for LDMs and discuss modifications the net- work architectures and the training setup that lead to better reconstruction quality and training efficiency. To motivate our approach, Figure shows that when visualizing the la- tent code learned by the Stable Diffusion VAE (SD-VAB), the code is itself image-like, with a strong similarity to the input. This observation leads us explore whether the learning of these latent representations can be simplified by applying fast aaage-processing function the input images prior to encoding. We opt for the discrete wavelet transform (DWT) as the amage-processingg function due to its image-like structure, proven effective- ness in extracting rich, compact features from images, and wide applicability in tasks such as image compression.
(b) Latent code Figure 2: RGB visualization of the first three channels of SD- VAE latent code.
(a) Input image
4.1 Model design
We now propose Vite 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 {ar n. a ) at level 1. To achieve 8x downsampling, we use three wavelet levels (i.e., E [1, 2)).. These features extract multiscale information from æ.
Feature extraction and aggregation: The wavelet coefficients {#), xt, æty æ'p) are then sep- arately processed via a
module Fi to compute a multiscale set of feature maps Fi(f"', THH. . a/p }). The features are then combined via feature- agguurragggt module Fage that takes in the output of each Fi and computes the latent z. We use UNet-based architecture similar to the ADM model [10] without spatial down/upsampling layers for feature extraction and aggregation. (See Appendix B for discussion of the importance of these learned modules.)
Image reconstruction:
Finally, decoder network D processes the latent code z and computes the reconstructed image é- 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 {Fi). Fage 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, Lite VAE 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. 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
Figure 3: Two examples of the feature maps from final block of the decoder before and after removing group normalization layers. Using SMC blocks instead of group normalization allows the model learn more balanced feature maps. The image best viewed when zoomed in.
due to excessive group normalization layers [69] in 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 modified version of modulated convolution [31] instead of group normalization to avoid imbalances. Instead of modulating the convolution layers via 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 wiyk according
W sjk
8:Wijjk
E. (SrWejM) +
for e >> O, where 8; a learnable parameter, and {1, 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 hance the training dynamics and reconstruction quality of Lite We verify the effect of these modifications in Sections 5 and 6.
Training resolution While the autoencoders in LDMs are typically trained on 256x 256 data (similar to SD-VAE), observe that the bulk of the training of VAE can be effectively conducted at lower 128x 128 resolution. Our experiment suggests that pretraining at this lower resolution followed by 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 also generally applicable the standard VAE models.
We replace the PatchGAN discriminator used in Stable Diffusion with 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 mixeddpreccisonn setups.
Improving the adversarial setup
Additional loss functions We also introduce two high- nigh-freqeency reconstruction loss terms based on the wavelet transform and Gaussian blurring [74]. Let ær be the input image and the corresponding reconstruction. For the wavelet term, we compute the Charbonnier loss [2] between the high- frequency DWT sub-bands {I Rv, #D) and (#LH &HI, } For the Gaussian loss, given Gaussian filter h. compute the ε loss between h(x)) and ă h((²).
5

---

Table 1: Comparison between LiteVAE and VAE in terms of reconstruction quality across different datasets and latent dimensions. VAE achieves better or similar reconstruction quality while having considerably fewer parameters in the encoder (34. 16M for the VAE and 6.75M for LiteVAE). All models use a downscaling factor of 8 and are trained from scratch with similar training
configs (including the choice of loss functions and discriminator).
Dataset
Latent dim
Model
rFID
LPIPS
PSNR
SSIM
VAE
0.88
0.089
28.08
0.85
FFHQ 128
16x 16x4
Lit
0.74
0.085
28.36
0.85
VAE
0.47
0.109
28.16
0.81
FFHQ 256
x2
LiteVAE
0.41
0.117
28.33
0.82
VAE
4.54
0.164
24.25
0.69
16x 16x4
LiteVAE
4.40
0.164
24.49
ImageNet 128
VAE 0.94
0.069
29.25
0.86
16x16x
LiteVAE
0.94
0.069
29.45
0.87
VAE 0.89
0.160
25.83
0.73
32x32x4
Lite VAE
0.87
0.157
26.02
0.74
ImageNet 256
VAE 0.23
0.073
30.41
0.86
32x32x12
Lite VAE
0.23
0.072
30.91
0.88
(a) Input
(b) Reconstruction
(c) Latent
Figure 4: An example of the autoencoder reconstruction alongside the learned latent code by Lite VAE,
We observe that Lite VAE maintains the image-like structure of SD-VAE
Experiments
This section presents a comprchensive empirical evaluation ViE, demonstrating its superior trade-off between computational efficiency and quality relative to standard VAEs. We further explore the properties of LiteVAR along with the changes proposed in Section 4. For each experiment, all models in comparison 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 in Rombach 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 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 fair comparison between different models.
Main results We first demonstrate that VAE matches or exceeds the performance of standard VAEs across various datasets and latent dimensions, shown in Table 1. Notably, the model employed for this table utilizes approximately one-sixth of the encoder parameters compared VAE model (6.75M vs 34. 16M) and hence trains faster. Also, one example of the reconstruction quality and the learned latent representation by VAE given in Figure 4. We notice that Lite VAE maintains the image-like latent codes, similar to the SD-VAE latent in Figure 2.

---

Table 2: Comparison of the scalability of Vit with standard VAE across different model sizes. (a) ViE matches the performance of the VAE with significantly fewer parameters and outperforms VAEs of similar complexity. (b) A naive downscaling of the VAE performs worse than Lite VAE. All models use the same decoder. More architecture details are provided in Appendix
(a) Scaling VAE (n, 12 for all models)
(b) Downscaling the VAE (n, = 4 for all models)
Model
Params (M) rFID
LPIPS 4
PSNR
SSIM
Model
Params (M)
rFD
LPIPS
PSNR
SSIM
VAE
34.16
0.95
0.069
29.25
0.86
VAE
34,16
4.54
0.164
24.25
0.69
LAE
1.03
0.075
29.12
0.86
VAE-Small
6.75
5.27
0.175
24.10
0.69
VAE-B
6.75
0.94
0.069
29.55
0.87
LIeVAEE
6.75
4.40
0.164
24.49
0.71
LAEM
32.75
0.79
0.064
29.66 68
0.87
LAE-L
41.42
0.74
0.062
29.94
0.88
Table 3: Comparing the complexity of encoder with the encoder from the Stable Diffusion VAE
for a batch size of 32. The values are measured on one Quadro RTX 6000.
Encoder
Params (M)
GPU Memory (MB)
Throughput (img/sec)
VAE Lite
34.16
8860
68
1.03
1324
384
VAE-B
6.75
3155
129
LitMVAE-M
32.75
12130
42.24
LiteLAE-L
41.425
12130
41.6
Increasing model complexity In Table 2 show the scalability of VAE as 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 VAE 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 VAE takes better advantage of the encoder parameters than normal VAEs, mainly due to the wavelet processing step that provides the encoder with rich representation from the beginning.
Computational cost Table presents a comparison of the computational costs between VAE 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 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, 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 I VAE at lower resolu- tion of x28x 128 followed by fine-tuning step on 256x256 images. To illustrate this, we compare model trained for 150k steps at full resolution (256-full) with one trained for 100k steps 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 the higher resolution. We also note that fine-tuning is essential, as the model trained solely on 128 x 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 freuuunryr contents in the fine-tuning stage. This pretraining technique reduced the overall wall-clock time of our training runs 256 256 resolution by more than factor of two.
Figure demonstrates that compared the standard VAFs, LiteVAE less prone to performance degradation when evaluating the model different resolutions. We hypothesize
Scale dependency
7

---

Table 4: Effect of replacing group normalization with SMC on reconstruction quality based on the ImageNet 128 128 model.
Table 5: Effect of pretraining the autoencoder at lower resolutions. We observe that training at
128 128 followed by fine-tuning at 256 256 per- forms best.
Normalization rFID 4 I.PIPS
PSNR 1 SSIM f
Group Norm
1.01
0.07
29.24
0.86
Training Config
rFID
LPIPS
PSNR
SSIM |
SMC
0.97
0.07
29.32
0.87
256-full
0.75
0.153
26.10
0.73
128-full
0.97
0.162
25.90
0.72
128-tuned
0.73
0.147
26.22
0.74
SSIM
LPIPS
0.3
0.5
0.7
0.1
Resolution (kcgg)
Resotation (kig)
Resolation Oiagg
Resolution (ingg)
Figure 5: Comparing the performance of VAE with normal VAE across different resolutions.
Lite VAE shows less degradation in all metrics.
Table Comparing MMD between Lite VAE la- tent space and a standard Gaussian vs SD-VAE latent space for different RBF kernels. VAE is statistically closer to standard Gaussian.
Table 7: Comparison between diffusion models trained in the latent space of a standard VAE [55] vs the latent space of VAE. We observe that both models perform similarly in terms of genera- tion quality.
σ
SD-VAE
VAE
Dataset
Encoder
FID 4
25
8.67:L0.1
1.44
28.90:10.49
0.94
FFHQ [30] (256x256)
LDM
8.11
100
10.77
5.14=0.19
LiteVAE
8.03
250
1.78+00.06
1.09:10.04
LDM
5.92
CelebA-HQ [29] (256x256)
500
0.44
0.28:10.01
Lite
5.73
that as our model learns features on top of multi-resolution wavelet coefficients, 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 Vite latent space We also analyzed the characteristics of the latent space of VAE. 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 similar image-like structure. Separately, we also examined the statistical distance between our model's 's latent space and pure Gaussian noise. The intuition here that, since diffusion model will have form a path from pure Gaussian noise 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 Lite VAE 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 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-HO datasets and compared their performance with standard VAE-based LDMs. The diffusion model architecture used for this experiment a UNet identical to the original model from Rombach al. [55]. Table 7 shows that the diffusion models trained in the latent space of VAE perform similarly to (or slightly better than) standard LDMs. Additionally, Figure includes some generated examples from our FFHQ model. These results suggest that diffusion models also capable of modeling the latent space of Lite VAE.
8

---

10a
10a
Constant Adaptiee
10a
10"
100
160
Training Steps (x 10")
Figure 7: Relative gradient norm of adversarial and the reconstruction loss.
Figure 6: Generated samples from the FFHQ model.
Table 8: Reconstruction quality after using con- stant weight for the adversarial loss.
Table 9: Effect of using Gaussian and wavelet loss on final reconstruction quality.
Weight Type
rFID
LPIPS
PSNR f
SSIM f
Truining Config
rFID
LPIPS
PSNR
Adaptive
1.01
0.07
29.24
0.86
Baseline
Constant
1.01
0.07
29.33
0.87
Gaussian loss
0.99
0.070
29.64
0.87
+ wavelet loss
0.96
0.069
29.73
0.88
Table 10: Reconstruction quality for different discriminators.
Discriminator
rFID +
LPIPS 4
PSNR +
SSIM t
UNet
1.01
0.070
29.24
0.86
StyleGAN
1.38
0.065
29.51
0.87
PatchGAN
1.61
0.063
29.61
0.87
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 latent size of 32x32x x 12 the baselinc for all ablations. Further ablation studies on other design choices in Lite VAE are provided in Appendix D.
Removing adaptive weight for Areg Table 8 demonstrates that weight for the adversarial loss (Equation (2)) and still slightly improve the metrics. Figure also shows the relative norm of the gradient of the adversarial loss compared to the reconstruction loss for both adaptive and constant Andv. We observe that using adaptive Xadv 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.
safely remove the adaptive
High-frequency loss functions Table shows the effect of adding eigh-frequeccy 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 UNet-based discriminator [61] outper- forms both PatchGAN and StyleGAN discriminators used previous works [55, 72] in terms of rFID while having comparable performance for other metrics. We also empirically noted that using UNet discriminator resuited 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 Lite a new design concept for autoencoders based on the multi- resolution wavelet transform. VAE can match the performance of standard VAEs 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, Lite offers more flexibility in terms of performance/compure 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 VAE extend to other domains is a question we leave to follow-up work. Although we introduced Lite in the context of LDMs, we hypothesize that its application is not confined this scenario. We consider the extension of Lite to other autoencoder based generative modeling schemes (e.g., tokenization) promising avenue for further research.
References
[1] 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: (0.485500/rXiv.2211 1.01224. URL https://doi.org/ w1w55////....22111011222 1.11224.
[2] Jonathan T Barron. A general and adaptive robust loss function. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages
2019.
[3] 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.15117,, 2023. doi: 1044550/ARXXV22111..550
URL
[4] 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.
[5] M. E. Brewster. An introduction wavelets (charles k. chui). SIAM Rev., 35(2):312-313, 1993. doi: 10.1137/1035061. URL https://doi..rg org/10.1 1137/1035061.
[6] S Grace Chang, Bin Yu, and Martin Vetterli. Adaptive wavelet thresholding for image denoising and compression. JEEE transactions on image processing. 9()): 1092-1522-1444 1546, 2000.
[7] Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, and William Chan. Wavegrad: Estimating gradients for waveform generation. In 9th International Conference Learning Representations, ICLR 1021, Virtual Event, Austria, May 3-7, 2021. Openkeview.net, 2021. URL http:///eeeeeeee..e/////o///l/ts
[8] Xiaojie Chu, Liangyu Chen, and Wenqing Yu. Nafssr: Stereo image super-resolution us- ing nafnet. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, pages 1239-1248, June 2022.
[9] Xiaoliang Dai, Ji Hou, Chih-Yao Ma, Sam Tsai, Jialiang Wang, Rui Wang. Peizhao Zhang, Simon Vandenhende, Xiaofang Wang, Abhimanyu Dubey, al. Emu: Enhancing image generation models using photogenic needles in a haystack. arXiv preprint 27X 2023.
[10] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat gans image syn- thesis. 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, NeurlPS 2021, 1, December 6-14, 2021, virtual, pages 8780-8794, 2021. URL hash/49ad23d haa. 101022681ffcff.... Abstract.html.
[11] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly,
10

---

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. URI. httpss//openreview.nee/ forumPid=VicbFaNNTYy
[12] Patrick Esser, Robin Rombach, and Björn Ommer. Taming transformers for high- resolution image synthesis. In IEEE Conference on Computer Vision and Par- tern Recognition, CVPR 2021, virtual, June 19-25, 2021, pages 12873-12883. Com- puter Vision Foundation / IEEE, 2021. doi:
URL https://ppenaccesssteccccci wwwpe///ooo/ooeeee/////////////e Staraiiioooommmmeeaeetieiiiiiiee
Taming_ Transformers_
for
[13] 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: 00-48550/ARXXV.2443... 033006 URL https://doi.org wtt.:/oo.rr/111.000000001 000448550/aXx..24003. 200.03306.
[14] Mário AT Figueiredo and Robert Nowak. An em algorithm for wavelet-based image restora- tion. IEEE Transactions on Image Processing, 1 12(8))906 9166699966666 2003.
[15] Rinon Gal, Dana Cohen Hochberg, Amit Bermano, and Daniel Cohen-Or. SWAGAN: style- based wavelet-driven generative model. ACM Trans. Graph., 44444::11111111111111 2021. doi: 1145/3450666.35568888
0114453345062.34598883 URL https://doi.orgg/
[16] Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffiusion transformer is a strong image synthesizer. CoRR, abs/2303. 14389, 2023. doi: 10.48550//aXi..2333..444 URI. http://dd..oog1110055555///////
14389.
[17] 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.
[18] Tiantong Guo, Hojjat Seyed Mousavi, Tiep Huu Vu, and Vishal Monga. Deep wavelet pre- diction for image super-resolution. In 2017 IEEE Conference on Computer Vision and Pat- tern Recognition Workshops, CVPR Workshops 2017, Honolulu, HI, USA, July 21-26 2017, pages 1100-110000 109. IEEE Computer Society, 2017. doi: 10.1109/CVPR..201..1 148. URL
https:/doi..rr/// http://dd...rr/1...1111/////CC/ 109.///9....00.....
[19] 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.0666 2023.
[20] 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 2022, 2022. URL http::/paaprss.nip wttppe/rssspsscccppee.e.ee hcc/ppper_fles/pappeeeee CeatrrattCCoffeencee.te himl.
[21] 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 IEEF/CVE International Conference on Computer Vision, ICCV 2023, Paris, France, October 1-6, 2023, pages 7407- 7417. IEEE, 2023. doi: 10.1199/CCV5107002223.000
URL https://doi.or//00 109/
ICCV51070.2023.000 200.000..222..00 300644
[22] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vish- wanathan, 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 httpss//procceednngs httpss/pppccedenn..uuuus.ccpppp 2017/hash/8all (atashh//ald947777666 714477770e00ee115777777777449990 111119997779222--btrrtmm.m.mm
[23] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. CoRR, Abs/2207.11559. 2022. doi: 0...4850//.xii.227.... URL https://bb..gg//.48555wxxxxxxx.
11

---

[24] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: An- nual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, De- cember 6-12, 2020, virtual, 2020. URI. Abstract.html.
(25] 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.00933 URL https://doi..rgg 10.445550arxii.2221111
[26] Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffussio: End-to-end diffusion for high resolution images. CoRR, abs/2301. 1109,, 2023. doi: 10.48550/aXi.2231.11000
URL https://doi.org/ 10..0005///r.iiii 230t:a.Xiiv.331..... 1109.
[27] Huaibo Huang, Ran He, Zhenan Sun, and Tieniu Tan. Wavelet-smett 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: 00.100//CCV.2217.1188 URL https://doi.orr http///oo.org/..1111111 wtgg//..1199/ICC// 100.1109C/22221..11...
[28] 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, Quoe V. Le, William Chan, and Wei Han. Moise2musie: Text-conditioned music generation with diffusion models. CoRR, abs/2302.03917, 2023. doi: 10.485550/aXx..2302..333 arXiv.2302.03917.
URL https://doi.org// http:///doooog/10.485555
[29] 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 30-May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL http://ppwnnreeiie.e/ooooommmmm
[30] 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.1 10.1109/CVPR2201 00453. URL http:///penaaceess.oo wtw.::swwwweeee..ooooooooee Eimerator_Architecttue Biaiiieccure_forGGGeeeaaaveee Siceeraive_AAdversaaiall Networks CVPR (ii__201___pa_a 2019_paper.html.
CVPR CVPR_2019/htii 2019//tml//arraaa A Style- -
[31] 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. 10.1000CV2P42200022000 00813. URL https://ppenaccesstthewww http///ceeeecctc..cooooocoeettt Ctteiicotttttcccccceee LWtp__220/t/////////e1 improving_the_Immee theep_tthhhh_eeeeee Lfaimaae__alttooofSSSeSSo StyleG/NN_CVPR__O222
1e000htmi/Karra_Annaaa _nnd
[32] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. 2022. URL https://opnneeiee.nneffommm k7FuTOWMOc7.
[33] Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models, 2023.
[34] 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..rr/abb//444 441..9900.
[35] Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. 2014. URL http:
//axxiv.or//abs/1311 1ttp://abbb/31..1111...
[36] 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 hwtp://ewwwweeeeeeeeo//om/iaaaaa
m751J.
[37] 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.
12

---

[38] Chen- 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.
[39] Guan-Homg Liu, Arash Vahdat, De-An Huang, Evangelos A. Theodorou, Weili Nie, and Anima Anandkumar. 1bsb: Image-to-image schrödinger bridge. CoRR, abs/2302. 05872, 2023. doi: 48555VVV22Xi
10. 00.4450/a2xi..222..5555
URL https://doi.org// http::/dddbooorg/144455555///ii
[40] Pengju Liu, Hongzhi Zhang. Kai Zhang. Liang Lin, and Wangmeng Zuo. Multi-level wavelet-cnn for image restoration. In 2018 IEEE Conference Computer Vision and Pat- tern Recognition Workshops, CVPR Workshops 2018, Salt Lake City, UT. USA, June 18- 22, 2018, pages 773-782. Computer Vision Foundation / IEEE Computer Society, 2018. doi: 10. 109/CVPRR.2018.. 20....201..0022.... 00121. URL http://openaccess.theee wtteevf.com/coniiii cetppcootetttttrpp___e 2opr_2018__ workshops/w 13/htm//Liu Melti-Level__ie Sereelaaweeeetteeeeeeeeeeeee CNN CVPR 2018 paper.html.
[41] 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.
[42] Stéphane Mallat. A theory for multiresolution signal decomposition: The wavelet representation. IEEE Trans. Pattern Anal. Mach. Intell., 1177):674 693, 1989. doi: 10.1109/34.12246. URL
https://doi.ogge 000 109/34, 192463.
[43] Stéphane Mallat. A wavelet tour of signal processing. Elsevier, 1999.
[44] Stephane G Mallat. A theory for multiresolution signal decomposition: the wavelet repre- sentation. IEEE transactions on pattern analysis and machine intelligence, 11(7))674 693, 1989.
[45] 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
[46] 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: botboceedings.ml..rreee.
httph/iccho22a.htmm
[47] Gaurav Parmar, Richard Zhang, and Jun-Yan Zhu. On aliased resizing and surprising subtleties in GAN evaluation. In IEEE/CVF Conference Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 11400- 14400011440, IEEE, 2022. doi: 10.1109/CVPR52688.. 211.PS22688820220111222 122.. URL http:://oooo..../..... 1AO.11199VVPP288822220.000
[48] William Peebles and Saining Xie. Scalable diffusion models with transformers. CoRR, abs/2212.09748, 2022.
doi: 10.48550Waxxx.2222.09444 URL https://oo.orr//004555500aa
2212.09748.
[49] 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: 1099CCVP52229. 2023.00983.
10110095222229 0177729.22200088.. URL https:///oi.org/////
[50] 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- high-resolution image synthesis. CoRR, 0wt.222070019533 2023. doi: 10.48550/ARXIV.2307000 https://doi.or/110. http::/ww/.1..0000aaii22.0000000
URL
[51] 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://openneviewwne//iiii id-FiNysSe7VyY.
[52] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. 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://proceediings wttp:eeeeddding.m..prssss/// 19/rramesh2 la html.
13

---

[53] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with CLIP latents. CoRR, abs/2204.06 06112,, 2022. doi: 28.550/arXi..22440060
URL httpp:://ooi.or///..... 400
[54] Oren Rippel and Lubomir Bourdev. Real-time adaptive image compression. In International Conference on Machine Learning, pages 2922-2930. PMLR, 2017.
[55] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm 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 18- URL
24, 2022, pages 10674-10685. IEHE, 2022. doi: 10.11090CVPS2228822222.00 http://....rrro...11/////./..... 202.020000..02204022...
[56] Negar Rostamzadeh, Emily Denton, and Linda Petrini. Ethics and creativity in computer vision. CoRR, abs/2112.03111 2021. URL https:/arrivv.or//bb//11 12..03111..
[57] 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 Fei- Fei. Imagenet large scale visual recognition challenge. Int. J. Comput. Vis., (5:3)::21-2522 2015. doi: 100700 1263-015-0816- URL https://doi.org ort. 1000..00//s11 1263-015-0161-06666 16-y.
[58] Seyedmorteza Sadat, Jakob Buhmann, Derek Bradley, Otmar Hilliges, and Romann M. Weber. CADS: Unleashing the diversity of diffusion models through aanition--nnnalle sampling. In The Twelfth International Conference on Learning Representations, 2024. URL https: wtwp:r/wwieee.eeeeouummiiieeeeee
[59] 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 Munkht- setseg 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. 10.1145/528233333337777 URL https://ooi.org/10000 http://d...rr0001.14522222323333 44//5528333335307557
[60]
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- -mage diffusion models with deep language understanding. 2022. URL http::/papers..ij http::sspspssccpperreeeese Conference.html
[61] 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.
[62] Ke Shen and Edward J Delp. Wavelet based rate scalable video compression. IEEE transactions on circuits and systems for video technology, 91):1099-222 1999.
[63] Wenzhe Shi, Jose Caballero, Ferene 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.
[64] Jascha Sohl-Dickstein. Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. 37:2256-2265, 2015. URL p.r..presss/v37/shll http://poceeddinn....mrrpsss///s
http://proceedings.
-iicc//ttt/slliiiiktiiiin 15..tml.
[65] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data dis- tribution. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' A'AlhhéBBus Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, Decem- ber 8-14, 2019, Vancouver, BC, Canada, pages 11895-11907, 2019. URL https:/proceediigs. cwurips.cc/paper//// htwpperr201//aah/0001eeee 9/thahh/300eee255777555553333 1add Abstract.tttttttmm
[66] Yang Song, Jascha Sohl-Dicksteins 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 http:://pwprrrooeee.e/////mm/d/d 2RRHS.
14

---

[67] David S. Taubman and Michael W. Marcellin. JPEG2000 image compression fundamentals, standards and practice, volume 642 of The Kluwer international series engineering and computer science. Kluwer, 2002. ISBN 978-0-7923-7519-7. doi:
URL. https:/ddb.orr//1110007997/... 1-007779788---115--7999--4
[68] 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.
[69] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference
on computer vision (ECCV), pages 3-19, 2018.
[70] 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.007900 2022. doi: 10448500/axX1.22200007799
URL http:://oi..or// 10. 10.44455/a2....2220.00000
[71] Mengping Yang, Zhe Wang. Ziqiu Chi, and Yanbing Zhang. Fregan: Exploiting fre- quency components for training gans under limited data. In Sanmi Koyejo, S. Mo- hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Informa- lion Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL Abstra--neccttenCCeeeeeeeee wwtppppcppappreiiiiiee http://ssssscccpppeeeeeee/pppppp 17
[72] 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 http:pppwwwwee..mmooo/oomdddd//o
[73] Yingehen Yu, Fangneng Zhan, Shijian Lu, Jianxiong Pan, Feiying Ma, Xuansong Xie, and Chunyan Miao. Wavefill: wavelet-based generation network for image inpainting. In Proceedings of the IEEE/CVF international conference computer vision, pages 1411 14114-
14123, 2021.
[74] Eduard Zamfir, Marcos V Conde, and Radu Timofte. Towards real-time 4k image super- resolution. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 1522-1532, 2023.
[75] Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, and Oliver Wang. The un- reasonable effectiveness of deep features as a perceptual metric. In 2018 IEEE Conference on Computer Vision and Puttern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pages 586-595. Computer Vision Foundation IEEE Computer Society, 2015 8. doi: 10. 10.1109/CVPR..O 8.0066.. URL http://ppenaccess. wtwpp//nnccess.eee....ooooooeeee thhecf.ccom/comeettcryy Zhang The Unreasonable Effectiveness CVVR 2018 1at18ppppr.html..
8/html/
[76] Zixin Zhu, Xuelu Feng, Dongdong Chen, Jianmin Bao, Le Wang, Yinpeng Chen, Lu Yuan, and Gang Hua. Designing better asymmetric VQGAN for stablediffusion. CoRR, 04422306. 2023. doi:
httpp:://ooi.rgg////// http:://w....gg//0.555//aaa/aaaa 0466..
15

---

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, approach holds promise for positive environmental impacts and sigmificant advancements generative modeling. It is important to note that while Al-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 VAE. We investigate using a non-learned (i.e., fixed) encoder in two settings: (I) 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 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 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. LiteVAA. essentially combines the computational benefits of the non-learned encoder with the learned latent space of a regular VAF.
Simple datasets For simpler datasets like FFHQ, is possible to completely replace the encoder & with a predefined function and get similar reconstruction quality. our case, we used a three-level DWT and only kept the sub-bands of the lowest level. We then trained decoder 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 similar reconstruction quality to 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-VAB by reducing the complexity of the encoder.
Table 11: The performance of the DWT-based encoder on simple datasets.
Datnset
Encoder
n2
rFID J
SD-VAE
FFHQ
DWT
12
0,70
DeepFashion
SD-VAE
1.64
DWT
12
1.71
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 non- learned encoder is effective in simpler scenarios, it falls short of the quality of normal VAEs in more complex settings. This indicates that the information present in the higher frequency of the wavelet transform is essential for the decoder to reconstruct more diverse images with higher quality. To incorporate the information from higher frequency sub-bands via
Table 12: The performance of the
non-learned encoder ImageNet.
Dataset
Encoder
1)g
rFID
SD-VAE
4
0.80
ImageNet
DWT
12
1.65
DWT-2
48
0.32
validate this hypothesis, space-to-depth operation [63] in the encoder and observe that we can recover the high reconstruction quality of the learned encoder (DWTT-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 VAE with non-learned wavelet- based latent representation for the FFHQ dataset, the learned latent space in VAE offers better structure for training diffusion models. Table 13 indicates that training the diffusion model on learned latent code of Lite- VAE outperforms the non-learned DWT representation. We argue that the sparse nature of wavelets harmful to generation quality compared to the dense representation Iearned by the encoder of Lite VAE.
Table 13: Comparison between a non- learned encoder and Lite VAE for train- ing diffusion models.
Dataset
Eneoder
FID
non-learned
12.51
FFHQ (256x256)
LiteVAE
8.03
16

---

C
Summary of diffusion models
Diffusion models learn the data distribution Pdata by reversing noising process that gradually converts a data point æ into random Gaussian noise. More specifically, diffusion models define forward process via æ; estimate the clean signal a from the current noisy sample 2t. has been shown that this process corresponds to the following stochastic differential equation (SDE) [66, 32]
o()),, where e (O, 1). Then, they train denoiser network Do to
dæ. = (()))(() Vx, (og))((;)) dt (1)) Vx, log))(#;)) dt + v (2(())(()) dwr,
(4)
where duy is the standard Wiener process, Pt (xxz is the distribution of noisy samples at time t, and B8() is term that controls the influence of noise during the sampling process. The denoiser network Do effectively approximates the score function V log (a,). Given that Po = Pdata and P1 = (0, 4) 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 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 diffusion models.
D Additional ablation studies
This section contains additional ablation studies on the design space and training dynamic of VAE.
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 diffcrent 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 VAES 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, used the LPIPS loss for the experiments in the main text have a similar training setup with commonly used VAES in LDMs.
Including the locally discriminative learning (LDL) loss 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 does not have any noticeable impact on the reconstruction quality of Lite VAE, as shown in Table 15.
Liang al. [37] introduced the LDL
Choosing different adversarial loss functions We also ablated the adversarial loss function for two different setups: a hinge loss, and 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 1x1 convolution layers
Ramesh al. [52] showed that using 1x1 x1 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 VAt and found that restricting the receptive field of the latent space with these 1 X 1 convolution layers might be harmful to the reconstruction quality by enforcing too much KL regularization. Table shows that removing these convolution blocks leads to much better reconstruction quality in our 256x256 model. Accordingly, we suggest removing these 1xl convolutions from the model (or, equivalently, adjusting the weight of the KL loss) to get better reconstruction.
17

---

Table 14: Ablation on using different VGG loss functions for the perceptual loss.
Perceptual Loss Type
rFID +
LPIPS 4
PSNR t
SSIM 1
LPIPS [75]
0.99
0.071
29.33
0.86
ESRGAN [68]
0.78
0.071
28.53
0.84
Table 15: Ablation on the effect of adding the LDL loss [37].
with LDL rFID +
LPIPS
PSNR 1
SSIM
0.99
0.071
29.33
0.86
0.98
0.071
29.33
0.87
Table 16: Ablation on using different adversarial loss functions.
Adversarial Loss
rFID
LPIPS | PSNR
SSIM t
Hinge Logistic
0.99
0.071
29.33
0.86
1.00
0.068
29.67
0.88
Table 17: Ablation on the effect of 1x1 convolution layers.
1x1 Conv
rFID
LPIPS
PSNR
SSIM
2.04
0.190
25.61
0.72
X
0.87
0.157
26.02
0.74
Table 18: Ablation using NAFNet [8] for feature extraction.
Fi
rFID 1
LPIPS
PSNR f
SSIM
UNet 0.94
0.069
29.55 29.36
0.87
NAFNet
0.93
0.069
0.87
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 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 Lite VAE is flexible w.r.t. this design choice. We chose the UNet to keep the setup as close as possible to the standard VAE design in LDMs.
D.4 Sharing the weights of the feature-extraction UNet
We next investigated whether single UNet could be shared across different wavelet sub-bands to further reduce the encoder's 's trainable parameters. Table 19 demonstrates that it is indeed possible to share between different sub-bands. shared UNet might lead 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 Fagg. As indicated in Table 20, employing ViT achieves comparable reconstruction quality to that fully-convolutionaal 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 VAE LDMss usually required to operate on data with varying resolutions. Hence, we side with the UNet models to make the encoder iesolution-indepennde
18

---

Table 19: Ablation on parameter sharing for the exature-extraction module.
Shared UNet
rFID 4
LPIPS
PSNR t
SSIM 1
0.75
0.153
26.10
0.73
0.78
0.154
26.05
0.73
Table 20: Ablation using VIT for feature aggregation.
F Fagg
Params (M)
rFID 4
LPIPS 4
PSNR
SSIM f
UNet
1.69
0.94
0.069
29.25
0.86
VIT
0.84
0.92
0.070
29.44
0.87
Table 21: Ablation on removing the highest resolution wavelets from feature extraction.
Config
rFID L
LPIPS
PSNR f SSIM
All sub-bands
0.87
0.157
26.02
0.74
Last two sub-bands
1.20
0.17
26.04
0.74
13
8:19
8:14
Sopa)
Figure 8: Comparing the performace of ViE with and without Group Normalization. Using SMC instead of Group Norm makes the autoencoder less scale-dependent.
D.6 Importance of using all wavelet levels
We also explored the possibility of performing feature extraction on only a subset of wavelet coeffi- cients 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 Lite The results in Figure 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 128x 128 followed by fine-tuning at 256 also works for the standard VAEs. The results of this experiment are given in Table 22. Similar to LiteVAE. 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.
19

---

Table 22: Effect of pretraining the autoencoder at lower resolutions for standard VAE model. The
128-tuned model performs similarly to the model trained solely on 2 256 data.
Training Config
rFID !
LPIPS +
PSNR t
SSIM t
256-full
0.67
0.150
26.01
0.74
128-full
0.89
0.161
25.83
0.73
128-tuned
0.69
0.151
25.97
0.73
Figure 9: Additional uncurated generations from the FFHQ diffusion model
Table 23: Details of the feature-extraction module for different Lite VAE models.
Feature extraction
Model
Input dim
Output dim
Channels
Channels multiple
VAE-S
12
12
16
(1, 21, 2)
VAE-B
12
12
32
(1, 3)
LiE-
12
12
64
(1, 4)
LAE-L
12
12
64
2,
Table 24: Details of the feature-aggregation module for different Lit models.
Feature aggregation
Model
Input dim
Output dim
Channels
Channels multiple
Lite VAE-S
36
Latent dim (ng)
16
(1, 2, 2)
VAE-B
36
Latent dim (ng)
32
(1, 2, 3)
Lit-
36
Latent dim (ng)
32
(1, 2, 3)
VAE-L
36
Latent dim (nzg)
64
(1,
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 256x 256 and FFHQ models. We use Adam optimizer [34] with learning rate of 10-4 and (81, 32) (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] ensure fair comparison.
20

---

G Pseudocode for different LiteVAE blocks
In this section, present additional pseudocode for various Lite VAE components. The core element of Lite is the Haar wavelet transform, which can be implemented in PyTorch as shown below:
class HaarTransform (nn..odule)
def __ (self, level=3, mode = wode="symment"" "symmetric", with-grad www.g-rdddaaes)) False) -> None:
super () in
self wavelet pywt. Wavelet ("haaarr""""" ("haar")
belf.level level
self mode mode
self welf.with.gras grad = with_grad
def (self, x, level=None):
10
with torch get__rrd____ lew.grrddenaaeedeeeeee (sel. (tteellllttegrddl
level = level or
low, *x_high ptwt. w
13
f)
14
wavelet=self. wtw.eeeeellww.eeee....e......e,e
level level,
mode=self mode,
17
18
X_ combined torch cat
[x_low, x_high[0] [o], x x_high [1], x_high(0) [2]], dim=1
31
return x_combined
22
23
def idvt (self, x):
with torch.set_pppp..
with_grad):
x_low, x_high x[:, :3], x[:, 3:]
26
x_high torch. chunk (x_high, 3, dim=1)
27
z_recon wwwtt.aaveee2(( ([x_[or float(), x_high. float ()),
vavelet=self wavelet)
return x_recon
29
30
def forward(self, x, inverse=False):
31
if inverse:
return self (x)
return velf.dw((x)
The PyTorch implementation of the self-modulated convolution block introduced in Section 4.2 is
provided below:
class MoCC(nnMModull def __init__(
self, channell: int, out_channels: int = None, kernel_size: int 3,
21

---

stride: int = None,
padding: int None,
bias: bool True,
10
super .._..i.t__iii
13
setting the default values
14
out_channels out_channels or in channels
padding- int (kernel_size // 2) if padding is None else padding stride_ = 1 if stride is None else stride
18
self padding padding_
self conv = nn. Conv2d
cn_hnennl,,,
22
out_channels,
23
kernel_sizen sizeesereeeeeeesize
padding-paddingg
stride=stride.,
26
bias=bias,
27
28
self gain nn Parameter (torch. ones (1))
self scales Pa.Paaameee ones (in channels))
31
32
def forward(self, x: torch. Tensor) -> torch. Tensor: scales = self scales. expand (x. shape [0], -1)
out modulateddconv2d
35
x=,, seeeff www.l..cov...e.. ww.v.wwight weight, s=scales,
36
sel padding,
39
input_gain=self. gain,
40
41
if self.conv conv. bias is not None: out out + self conv. wonv.bias. view(1, -1, 1, 1)
return out
Next, we present the code for the residual blocks utilized in the Lite VAE UNet networks:
class ResBlock :
def __init__
self, in_channels: int, dropout: float 0.0,
out_channels: int = None, use_conv: bool = False,
activation: str "swish",
22

---

Corm_nun_groupp int = 32,
scale_ float 1,
11
12
super in__ini___
self channels tn_channels
15
self out_channels out_channels or in channels
16
self in = crouplormmiin,hhmmeel. norm corm_num-groups)
self act_in SiLU()
self.conv_it Cinnaael out_channels, 3)
20
self norm_out (rouplorm(out_chhnnell norm_num_grev groups)
21
act_outt SiLU()
dropout = Dropout (dropout)
self.covv_ce
ConvLayer2Doutt. (oahaacnne, 3)
24
25
if self. out channell == in channell:
self akiip_cooeectto = Identity()
elif use_conv:
28
self comppcoo
out_channels, 3)
29
else:
30
self. con = (in.layer2Dii_chaamee out_channels; 1)
self scate_ ocale_factor
32
33
def (sel,, x):
34
input layers
h self norm_in(x)
h self (n
37
h self conv_in
38
ou output layers
h self nut (h)
h self act_out (h)
41
h self dropout (h)
42
h self. conv_oootuutt (h)
43 return (self wtw.ssciiccoometiionf..... + h) self scale factor
46
class ResBlockWithSMCCnn Module))
47
def __init__
self,
in_channels: int, dropout: float 0.0,
30
51
int None, use_conv: bool = False, activation: str "swish", norm_num_groups: int 32,
52
55
scale_factor: float 1,
56 ):
23

---

super () -0_ init__0
59
self in_cheneels channelss
60
self.out_chee channellseels = out_channels or in, channels
61
self in SiLU()
colf.covv.. in SMC channelss out_channels,
64
self.act_owd out SiLU()
65
self dropout Dropout (dropout)
self conv_out = SMC foout_chhanees,, 3)
68
if self out_channels in_channels:
60
self aakip_connection. Identity()
elif use_conv:
71
self skip_ connection (onvLayer2Diin_chanaee out channels, 3)
72
else:
73
self wwtssskiippppciciieeeeee connection (inpaaeeee((innannnaaa out channels, 1)
self iicaleffacton = scale_factor
input layers 76
def forward(self, x):
77
78
h = self act_in(x)
= self conv_in (h)
output layers
81
self act_out (h)
82
h = self dropout (h)
h = self. conv_out (h)
return (self comecciio h) / self scale factor
85
36
class MidBlock2D (oduul)))))
def __init__
self,
90
in_channels: int,
91
out_channels: int,
dropout: float = 0.0,
use_amc: bool True,
94
None:
95
(opper()..iiit_..) in
cassoockcclass ResBlockWithSMC if use_smc else ResBlock
self res0 resblock_ class
98
in mm__chamees-in_chhmmll..
99
out out_chammels
100
dropout-dropout,
102
self res1 in.chhemnels-oottt.ccmmmeee.
resblock.class(
103
101
out_ mnt_chammelss-oottcchmmmelll
24

---

dropout-dropout,
107
def fsel,
ice
x self res0
self res.reell (x)
return
Additionally, the feature-extraction and crature-aggregation UNets can be implemented as follows:
class LiteVAAUDeetBlockkmn Module)
def __init__
self, in_channels: int,
out_channells: int,
model_channels: int,
ch_multiplies: list[int] = [1, 2, 4],
num_res_ aum_kes_block:: int 2,
use_smc: bool False,
super () init__ self in layer ConvLayer2D(in_ channellss,,
3)
13
self out_laerr ConvLayer2D (model (model_channels,, channels, out_channels,
14
resblock_class RosBlockWithSMC if use_smc else ResBlock
17
18
Ulltt encoder path
channel = model_channell]
in co__hhnneel_11 _ist [model_channels] self encoder_blocks
22
23
for level, ch ch_mult in (ch for in range (hum_res_blocke
self encoder wwwoder_blocks.co. append resblock_clase in nhacchnnells-cameel,
26
27
out_cchamels-model channels ch_mult
30
31
channel model_channels ch_mult
32
in_ laatnnel_listttie apwpeeennddlaappppnnpccea self. elooder_blocks nn wwM.Molllllitteelll.... encoder =mcokkr_blocks)
35
Ueltt middle block
36
self midd mid_block = MidBlock2D
in_nnennlls channneellchnnl,,,
39
out_ channeel,
25

---

chibed_channel=0,,
legacy=legacy
42
43
# UNet decoder path
46
self dww.derr_bockk.
47
for level, ch_mult in reversed (list (enumerate (ch_ (ch_multiplees) for in range thum_res_blocke
self decoder decoder_blocks. www.bbloockkpeoo append(
resblock.clase,
51
ch_chanmell=channel mh_channee_list pop
52
out_ pat__chamellsmmode_.ccmee channels ch_mult
55
channel model_channels ch_mult
56
self. decoder decooerrbblocks blocks nn. wwd.Mooulllst (self (self.decode....ee. decoder decoder_blocks)
def forward(self, x)
59
self in__ayer (x)
60
skip_features [x]
61
the encoder path
for enc_block in self. ancoder_blockes
x (x)_
64
skip_features (ttp.ffaauuessaapen(()))) (x)
65
the middle block
= self (id_blockkx)
the decoder path
for deccblock in self decooder_boockss
69
cat torch.cat cat [x, skip_features pop(), dim=1)
(xc_blocc(x_catt return self.out out_ (x)
The VAE encoder can be implemented as shown below:
class LiteVAEEncoder (nn.Modin Module)
def __init__
self, in_channels: int,
out_channell: int, wavelet_fn: HaarTransform, feature_extractor_paat dict,
dict,
self wavelet_fn wavelet_fn self www.urreextractorr.o L1 LiteVAEDNetBLock
13
channels, channels,,, *ffeture__errratoor.aaaass
26

---

)
self tnatburr_xttrattosss. L2 LiteVAEUWeeBllock in_chelel,, in_hhennlls, *
16
17
)
self
LiteVAEUNetBLock
in channels,,, in channelels,
20
21
out_channels out_ channels 2 for VAE mean and log_var var
cggrogl = in_ * 3
self reggatree = LiteVAEUWettBlock
aggregated,chameela. out_channels, #
25
self downsample_ b_ock___l_____ L1 Jovnaample2l(inccaanmee. scale_factor= 36
self www.ddaopllellleoooo block_L__ook_2_ L2
(attappeeeeinnnannnaae scale_factor=2)
29
def forward image)
30
dwt_L1 self wavelet fn. dwt (image, level=1)
dwt_L2 self wavelet_fn. dwt (image, level=2) 4
dwt_L3 self wavelet_fn dwt (image, level=3)
33
features_L1 self www.sampee..llc........
34
self fft.ture_eexrrattto_l[(((t..
35
features_L2 = self downsample_blooc__22
37
self feature_ axtreeerrrrrttt___1[tt1__..1
38
39
features_Ls L3 = self. Latiuure_extraatrr___C (dwtt_Ld
fea_uee [features_I] _1, features_Ly features_L3]
latent www.feereagggrgggeeee (torch. cat (features, dim=1))
42
return latent
Finally, the code for Lite is also provided below.
class LiteVAE(nn Module))
def _nint___ self, encoder: LiteVAEEncoder, SDVAEDecoder, config: DictConfig, output_type: str "image",
decoder:
assert output_type in ["image", "wavelet"] self encoder encoder
12.
self decodrr decoder
13
welf.wavelet.ce fn encoder encoder.waveless fn
cilffoutput_tyyee = output_type
pra_nhannels config latent_din for VAE mean and log_ var
27

---

channels = config laatetttttiil
if config get ["use_ixi_con'"' False)
19
self. pre_covv nn connooderree.ccamelll pre_channels,
30
self.post_come mn. Conv2d (post_ channels, post_ channel,, 1)
else:
self. pte_ppre_con Identityy)
23
self colf.oostcomi mn. Identity( ()
24
def pacode(self. image): return self pre_ (self. encoder encoder (image)) )
27
28
def (self, latent):
29
latent = self post_conv (latent)
if self wwt.ttuupptttpppeeee == "image":
image_recon self decoder (latent)
32
wavelet_recon self wavelet_fnm f_ dut (image_recon, level=1)
33
elif self wwtpftuututtyyyy type == "wavelet":
wavelet_recon self decoder (latent)
image_recon = self return image_recon, wavelet_recon
(wavelet_recon, level=1) 2
36
37
38
def (Belf, image, sample=True):
latent self encode encode
latent_dist Diaggaa..Guusssaaalliitttrbttroa
41
latent latent_dist ) sample() if sample else aatent_diste mode kl_reg latent_dist kl mean
42
image_recon havelet_recom = self decode (latent)
return Dict(
45
{
46
"sample": image_recon,
"wavelet": vavelet_recon,
"latent": latent,
"kl_reg": kl_reg,
50
"latent_dist"; latent_dist,
51
}
)
28

---

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 complete (and correct) proof?
Answer: [NA]
Justification: The paper does not include theoretical results.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The hyperparameters, algorithms, and implementation details are provided in the appendix for proper reproducibility our work.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions 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, al- gorithms, and implementation details are provided in the appendix ensure reproducibility.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide the details 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 com- puter 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
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https://eeeiii..cccppbii/ttcccct
Question:
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.
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
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?
Question:
Answer: [Yes]
Justification: We 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 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.
30


Reviews:
[{"review_id": "9MddaccgOJ", "reviewer": "Reviewer_rH5p", "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.", "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?\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?", "limitations": "The limitations are not discussed.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "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.\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.", "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?\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We wish to thank the reviewer for the helpful comments and for finding our work novel with detailed evaluations, good presentation, and significant contribution. Please find our answers to the comments below.\n\n### **Impact of the work**\nThe VAE component in latent diffusion models is responsible for processing high-resolution images, which can be computationally intensive. We observed that replacing the standard VAE with LiteVAE during the training of DiT models results in approximately 30-35% increase in the speed of each training step. Additionally, as highlighted in the introduction of our submission, the GFLOPs required for querying the VAE encoder exceed those required by the Stable Diffusion UNet. Consequently, enhancing the performance of the VAE significantly improves the training efficiency of latent diffusion models. Moreover, in applications such as score distillation (SDS) [1], the algorithm requires backpropagation through the encoder, meaning that optimizing the effeciency of the encoder can have a substantial impact on performance in these contexts as well. \n\n[1] Poole B, Jain A, Barron JT, Mildenhall B. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209.14988. 2022 Sep 29.\n\n### **Code availability**\nWe agree with the reviewer that making the work reproducible is an important aspect of experimental papers. Unfortunately, due to internal copyright policies, we are unable to share the source code of this work. For computing wavelets, we used the `fbcotter/pytorch_wavelets` library. We have also included detailed implementation details in the appendix and would be happy to add more information and detailed pseudocode to ensure that the results are easily reproducible.\n\n\n### **Application to other scenarios**\nWe thank the reviewer for pointing out this somewhat ambiguous terminology. Our work focuses on the application of latent diffusion and autoencoders for natural images, and we do not have any claims for other domains. What we mean by line 287 is that the use of wavelets can be explored further in other autoencoder-based generative models for natural images (such as vector-quantized models). Hence, the focus of the paper is solely on the natural image domain only. That being said, we agree with the reviewer that our method assumes wavelets are well-suited for the particular application in which LiteVAE is being used. We thank the reviewer for pointing this out and would be happy to acknowledge this as a limitation or assumption of our method.\n\n### **Error bars for Table 4**\nWe agree with the reviewer that including error bars in table 4 makes the result more convincing. However, doing so would require multiple training of different autoencoders, which is outside our compute budget. We should also note that other than the slight improvements in reconstruction quality, removing feature imbalances improves training stability as observed in previous works (e.g., [1, 2, 3]). \n\n[1] Karras T, Laine S, Aittala M, Hellsten J, Lehtinen J, Aila T. Analyzing and improving the image quality of stylegan. InProceedings of the IEEE/CVF conference on computer vision and pattern recognition 2020 (pp. 8110-8119).\n\n[2] Karras T, Aittala M, Lehtinen J, Hellsten J, Aila T, Laine S. Analyzing and improving the training dynamics of diffusion models. InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition 2024 (pp. 24174-24184).\n\n[3] Salimans T, Kingma DP. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems. 2016;29.\n\n### **Question about the GFLOPs**\nThe GFLOPs are reported for a single forward pass through the encoder and the diffusion UNet. These are the forward calls that have to be made at each training step of LDMs. As this step is used during training, we can conclude that a noticeable improvement can be made by switching the VAE to a more efficient version since its computational complexity is even larger than the UNet part. \n\n\n### **Question about using LiteVAE instead of SD-VAE**\nLiteVAE can be used as a drop-in replacement for SD-VAE in applications that only utilize the autoencoder from Stable Diffusion. However, because changing the autoencoder alters the latent space, the Stable Diffusion model is not compatible with LiteVAE out of the box. However, as pointed out in [1], it is possible to fine-tune the diffusion UNet with a new VAE for rapid adaptation of existing pretrained models.\n\n[1] Chen J, Ge C, Xie E, Wu Y, Yao L, Ren X, Wang Z, Luo P, Lu H, Li Z. Pixart-\\sigma: Weak-to-strong training of diffusion transformer for 4k text-to-image generation. arXiv preprint arXiv:2403.04692. 2024 Mar 7.\n\n### **Question about using wavelets with SD-VAE**\nWe believe that using wavelets with SD-VAE will also enhance quality and performance. However, please note that doing so deviates from the original SD-VAE model and makes the setup more similar to LiteVAE. In case we misunderstood the question, we would be happy to provide more discussion on this.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2UsveRHLVM", "reviewer": "Reviewer_xykJ", "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.", "questions": "See the weaknesses part.", "limitations": "None", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper integrates multi-scale VAE and discrete wavelet transform to reduce computational cost and boost reconstruction performance.\n2. The paper offers a detailed training pipeline for VAE in diffusion models, along with ablation studies that readers will find beneficial.\n3. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer's helpful suggestions, as well as the positive assessment of the influence and quality of our work.  Below, we provide detailed responses to the reviewer’s comments.\n\n### **Question about multiscale VAE and wavelets**\nWe would like to note that the multiscale structure of our model and the wavelet transforms are in fact relatively coupled, as the multi-scale part arises exactly because wavelet transforms are multiscale operations in nature. We are happy to make this relationship more clear in the paper.  We did extra ablations in Appendix D on the importance of using all wavelet levels, as well sharing the feature-extraction networks for each part and other properties of this specific architecture. We would be happy to include additional ablations in the final version if requested.\n\n### **Code availability**\nAs pointed out in our general response, due to internal copyright policies, we are unfortunately unable to share the source code of this work. However, we will do our best to make sure that the results are reproducible by providing more implementation details and detailed pseudocode in the final version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cvNYqLXiSQ", "reviewer": "Reviewer_U437", "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.\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.", "questions": "Table 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\".", "limitations": "The checklist points to Section 7 for the limitations but the section only shows the conclusion. Please address this.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "ORIGINALITY:\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.", "weaknesses": "Table 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\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing constructive comments and for recognizing our paper as high-quality with numerous strengths and significant contribution. Please find our answers to the comments below.\n\n### **Throughput of other LiteVAE models**\nWe thank the reviewer for pointing out this question. Please find the throughput of other LiteVAE models in the following table, and we will include this updated result in the final version of the paper.\n\n| Model     | Throughput  (img/sec)   | GPU Memory (MB) |\n|-----------|:----------------:|:-------------:|\n| LiteVAE-S | 384     | 1324  |\n| LiteVAE-M | 42.24  | 12130  |\n| LiteVAE-L | 41.6  | 12170  |\n\n### **Compute cost for wavelet transforms**\nThe computational cost of wavelet transform is linear in terms of the number of pixels and is negligible compared to querying the neural network. For a tensor of shape (32, 3, 256, 256), computing the wavelets takes 829 microsecond on an RTX 3090 GPU while querying the LiteVAE-B encoder takes 55.1 millisecond for the same data.\n\n### **Error in line 207**\nWe thank the reviewer for mentioning this error. The text indeed means Table 2, and we will fix this issue in the final version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R1roMhBu0P", "reviewer": "Reviewer_86Z2", "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.", "questions": "* Did the authors try to expand the bottleneck channel from 4 to other numbers, such as 16, 32, or more?\n\n* Could the authors provide some distribution analysis about the latent features and the differences between LiteVAE and VAE?\n\n* Could the authors provide the values of LiteVAE-S/M/L in Table 3?\n\n* 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.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The new encoder could reduce more parameters than the original one with less GPU memory.\n\n2. The paper provides some interesting analysis of the feature map.\n\n3. 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? \n\n* 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.)\n\n* Did the paper have the ablation study on the structure modules in Fig.1?\n\n* 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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's helpful comments and the positive reaction to our work. Please find our responses to the individual comments below.\n\n### **Information about the decoder**\nAs mentioned in line 149, we use the same decoder architecture as SD-VAE, and the encoder and decoder networks are trained together. While the decoder is the component mainly used during inference, our work targets the efficiency of LDM training. During each training step, only the encoder is used; therefore, the complexity of the encoder directly affects the training efficiency of LDMs. Additionally, the encoder is used in other LDM applications, such as image editing and score distillation (SDS) [1]. Thus, improving the efficiency of the encoder will also enhance performance in those applications. Please also note that the encoder is the component that must be fixed before training the diffusion part. The decoder can be distilled into a more lightweight network after training the diffusion and the autoencoder, as it does not change the latent space of the VAE.\n\n[1] Poole B, Jain A, Barron JT, Mildenhall B. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209.14988. 2022 Sep 29.\n\n### **Question about compatibility with latent diffusion models**\nWe thank the reviewer for mentioning this interesting question. As mentioned in the general comment, the focus of this work is offering a new encoder architecture for latent diffusion models that results in faster training of the first stage and more efficient training of the second stage of LDMs. As changing the encoder changes the latent space of the VAE, the model cannot be directly used with pretrained diffusion models such as Stable Diffusion. We show that it is possible to train latent diffusion models in the latent space of LiteVAE, but training another Stable Diffusion is unfortunately not in the scope of our compute budget. As the Stable Diffusion models also train a new autoencoder for each version (e.g., SD 2.1, SDXL, and SD3 all have different VAEs), we hope that our findings will be useful when developing new SD models. We did not try direct fine-tuning of existing SD models due to limited computational resources, but [1] demonstrates that it is possible to fine-tune the diffusion UNet with a new VAE to adapt existing pretrained models.\n\n[1] Chen J, Ge C, Xie E, Wu Y, Yao L, Ren X, Wang Z, Luo P, Lu H, Li Z. Pixart-\\sigma: Weak-to-strong training of diffusion transformer for 4k text-to-image generation. arXiv preprint arXiv:2403.04692. 2024 Mar 7.\n\n\n\n### **Ablation studies regarding LiteVAE structure**\nBesides the scaling experiments mentioned in the main text, we provided several ablation studies in Appendix D regarding the feature-extraction network architecture (D.3), sharing feature-extraction modules (D.4), using ViT for feature aggregation (D.5), and the importance of using all wavelet levels (D.6). We would be happy to include more ablations in the final version if requested.\n\n### **More visual comparisons for the VAE-based LDMs**\nFigure 6 and 9 include generated images based on our latent diffusion model. Hence, there is not a direct correspondence between these images and the images generated by another LDM trained on the same data. As the FIDs of these two models are close to each other, we expect the VAE generations to also have similar characteristics. We would be happy to also provide generations for the VAE-based LDM models in the final version.\n\n### **Question about the bottleneck channel**\nWe have included experiments for $n_z=4$ and $n_z=12$ in the paper, and the conclusions are similar. Our internal experiments also showed that similar results hold for $n_z=16$ and $n_z=48$. Hence, we conclude that our findings are independent of the number of channels used in the encoder bottleneck layer.\n\n\n### **Distribution analysis**\nThe submission includes a distribution analysis between SD-VAE and LiteVAE latent spaces in Table 6, where we concluded that the latent space of LiteVAE is closer to a standard Gaussian distribution in terms of MMD metrics. \n\n### **Throughput for other LiteVAEs**\nWe thank the reviewer for pointing out this missing detail. Please find the throughputs in the following table, and we will include these results in the final version of the paper.\n\n| Model     | Throughput  (img/sec)   | GPU Memory (MB) |\n|-----------|:----------------:|:-------------:|\n| LiteVAE-S | 384     | 1324  |\n| LiteVAE-M | 42.24  | 12130  |\n| LiteVAE-L | 41.6  | 12170  |\n\n### **Comment about the checklist**\nWe would be happy to double check the checklist in the final version to make sure that it is fully compatible with NeurIPS requirements.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for recognizing our paper as well-structured and easy to read, and for highlighting its interesting ideas and detailed evaluations. 

We would like to clarify that the primary goal of LiteVAE is to study the efficiency and reconstruction capabilities of the autoencoder, as we believe this is an important area that has received comparably less attention from the LDM community than the diffusion model itself. Our main contribution is offering a more efficient encoder that achieves the same reconstruction quality with significantly less parameters. We demonstrate that our method achieves comparable reconstruction quality to that of a standard VAE while requiring significantly less compute. This leads to faster (>2x) training of the first stage and higher throughput (up to 35%) in the second stage of LDMs, as well as having the potential to improve efficiency for LDM-based applications such as score distillation (SDS) [1].

Also, since some reviewers asked for the code of our work, we should unfortunately mention that due to internal copyright policies, we are not able to share the full code of the paper. However, we would be happy to add more information regarding the implementation details and detailed pseudocode for different LiteVAE components to ensure reproducibility. In response to the reviewers' suggestions, we are happy to also expand the limitation section of the paper in the final version.

We have also prepared individual responses to each reviewer and welcome any follow-up discussions. Given that there are no major concerns with our work and reviewers agree that it is a novel and well-presented paper with strong motivation/evaluation and relatively significant contribution, we hope that our rebuttal motivates the reviewers to adjust their scores accordingly.


[1] Poole B, Jain A, Barron JT, Mildenhall B. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209.14988. 2022 Sep 29.

Meta-review: The paper proposes a new lightweight multiscale encoder by incorporating 2D DWTs, thereby significantly decreasing the encoder parameters leading to faster training and lower GPU memory requirements as compared to traditional VAEs. A set of model sizes are presented leading to desirable GPU memory, throughput and reconstruction quality tradeoffs. 

The paper received four reviews with unanimity in their (post discussion) ratings. The reviewers carefully engaged with the paper and participated in subsequent discussions. The paper is readable and well-organized, the key ideas well motivated, and the reviewers liked the reasoning accompanying the main approach. Though not entirely novel, the use of wavelets to address the computational costs of training LDMs is seen to have a potentially significant impact and the proposed approach novel. The training details and the ‘tricks’ shared are expected to be valuable for improving VAE training. The experiments and ablations are found to be sufficient to support the main claims of the work. 

The initial reviews were split around the borderline. The primary concerns were around the actual computational and quality impact when the LightVAE is incorporated into LDMs, lack of error bars when the statistical results are close, reproducibility, ablation studies and sensitivity analysis, and some typos and missing details that needed clarification. The authors addressed many of these concerns during the discussion resulting in an upward revision of ratings by two reviewers resulting in unanimity. 

While concerns around the actual impact when LightVAE is incorporated into LDMs, statistical significance, and reproducibility persist, the paper can be significantly improved by incorporating the clarifications and material provided by the authors in their rebuttal and responses.
Decision: Accept (poster)