Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models

## Abstract

Abstract We present Deep Compression Autoencoder (DC-AE), a new family of autoencoders for accelerating high-resolution diffusion models. Existing autoencoders have demonstrated impressive results at a moderate spatial compression ratio (e.g., 8 × \times × ), but fail to maintain satisfactory reconstruction accuracy for high spatial compression ratios (e.g., 64 × \times × ). We address this challenge by introducing two key techniques: (1) Residual Autoencoding , where we design our models to learn residuals based on the space-to-channel transformed features to alleviate the optimization difficulty of high spatial-compression autoencoders; (2) Decoupled High-Resolution Adaptation , an efficient decoupled three-phase training strategy for mitigating the generalization penalty of high spatial-compression autoencoders. With these designs, we improve the autoencoder’s spatial compression ratio up to 128 while maintaining the reconstruction quality. Applying our DC-AE to latent diffusion models, we achieve significant speedup without accuracy drop. For example, on ImageNet 512 × 512 512 512 512\times 512 512 × 512 , our DC-AE provides 19.1 × \times × inference speedup and 17.9 × \times × training speedup on H100 GPU for UViT-H while achieving a better FID, compared with the widely used SD-VAE-f8 autoencoder.

## 1 Introduction

Latent diffusion models (Rombach et al., ) have emerged as a leading framework and demonstrated great success in image synthesis (Labs, ; Esser et al., ). They employ an autoencoder to project the images to the latent space to reduce the cost of diffusion models. For example, the predominantly adopted solution in current latent diffusion models (Rombach et al., ; Labs, ; Esser et al., ; Chen et al., ; ) is to use an autoencoder with a spatial compression ratio of 8 (denoted as f8), which converts
images of spatial size $H\times W$ to latent features of spatial size $\frac{H}{8}\times\frac{W}{8}$. This spatial compression ratio is satisfactory for low-resolution image synthesis (e.g., $256\times 256$). However, for high-resolution image synthesis (e.g., $1024\times 1024$), further increasing the spatial compression ratio is critical, especially for diffusion transformer models (Peebles & Xie, ; Bao et al., ) that have quadratic computational complexity to the number of tokens.

The current common practice for further reducing the spatial size is downsampling on the diffusion model side. For example, in diffusion transformer models (Peebles & Xie, ; Bao et al., ), this is achieved by using a patch embedding layer with patch size $p$ that compresses the latent features to $\frac{H}{8p}\times\frac{W}{8p}$ tokens. In contrast, little effort has been made on the autoencoder side. The main bottleneck hindering the employment of high spatial-compression autoencoders is the reconstruction accuracy drop. For example, Figure 2 (a) shows the reconstruction results of SD-VAE (Rombach et al., ) on ImageNet $256\times 256$ with different spatial compression ratios. We can see that the rFID (reconstruction FID) degrades from 0.90 to 28.3 if switching from f8 to f64.

This work presents Deep Compression Autoencoder (DC-AE), a new family of high spatial-compression autoencoders for efficient high-resolution image synthesis. By analyzing the underlying source of the accuracy degradation between high spatial-compression and low spatial-compression autoencoders, we find high spatial-compression autoencoders are more difficult to optimize (Section 3.1) and suffer from the generalization penalty across resolutions (Figure 3 b). To this end, we introduce two key techniques to address these two challenges. First, we propose Residual Autoencoding (Figure 4) to alleviate the optimization difficulty of high spatial-compression autoencoders. It introduces extra non-parametric shortcuts to the autoencoder to let the neural network modules learn residuals based on the space-to-channel operation. Second, we propose Decoupled High-Resolution Adaptation (Figure 6) to tackle the other challenge. It introduces a high-resolution latent adaptation phase and a low-resolution local refinement phase to avoid the generalization penalty while maintaining a low training cost.

With these techniques, we increase the spatial compression ratio of autoencoders to 32, 64, and 128 while maintaining good reconstruction accuracy (Table 2). The diffusion models can fully focus on the denoising task with our DC-AE taking over the whole token compression task, which delivers better image generation results than prior approaches (Table 3). For example, replacing SD-VAE-f8 with our DC-AE-f64, we achieve 17.9$\times$ higher H100 training throughput and 19.1$\times$ higher H100 inference throughput on UViT-H (Bao et al., ) while improving the ImageNet $512\times 512$ FID from 3.55 to 3.01.
We summarize our contributions as follows:

- •
We analyze the challenges of increasing the spatial compression ratio of autoencoders and provide insights into how to address these challenges.
- •
We propose Residual Autoencoding and Decoupled High-Resolution Adaptation that effectively improve the reconstruction accuracy of high spatial-compression autoencoders, making their reconstruction accuracy feasible for use in latent diffusion models.
- •
We build DC-AE, a new family of autoencoders based on our techniques. It delivers significant training and inference speedup for diffusion models compared with prior autoencoders.

Figure: Figure 1: DC-AE accelerates diffusion models by increasing autoencoder’s spatial compression ratio.
Refer to caption: x1.png

Figure: Figure 2: (a) Image Reconstruction Results on ImageNet 256$\times$256. f denotes the spatial compression ratio. When the spatial compression ratio increases, SD-VAE has a significant reconstruction accuracy drop (higher rFID) while DC-AE does not have this issue. (b) ImageNet 512$\times$512 Image Generation Results on UViT-S with Various Autoencoders. p denotes the patch size. Shifting the token compression task to the autoencoder enables the diffusion model to focus more on the denoising task, leading to better FID. (c) Comparison to SD-VAE-f8 on ImageNet 512$\times$512 with UViT Variants. DC-AE-f64p1 provides 19.1$\times$ higher inference throughput and 0.54 better ImageNet FID than SD-VAE-f8p2 on UViT-H.
Refer to caption: x2.png

## 2 Related Work

#### Autoencoder for Diffusion Models.

Training and evaluating diffusion models directly in high-resolution pixel space results in prohibitive computational costs. To address this issue, Rombach et al. () proposes latent diffusion models that operate in a compressed latent space produced by pretrained autoencoders. The proposed autoencoder with $8\times$ spatial compression ratio and $4$ latent channels has been widely adopted in subsequent works (Peebles & Xie, ; Bao et al., ). Since then, follow-up works mainly focus on enhancing the reconstruction accuracy of the f8 autoencoder by increasing the number of latent channels (Esser et al., ; Dai et al., ; Labs, ). In contrast, our work focuses on an orthogonal direction, increasing the spatial compression ratio of the autoencoders (e.g., f64). To the best of our knowledge, our work is the first study in this critical but underexplored direction.

#### Diffusion Model Acceleration.

Diffusion models have been widely used for image generation and showed impressive results (Labs, ; Esser et al., ). However, diffusion models are computationally intensive, motivating many works to accelerate diffusion models. One representative strategy is reducing the number of inference sampling steps by training-free few-step samplers (Song et al., ; Lu et al., ; ; Zheng et al., ; Zhang & Chen, ; Zhang et al., ; Zhao et al., ; Shih et al., ; Tang et al., ) or distilling-based methods (Meng et al., ; Salimans & Ho, ; Yin et al., ; ; Song et al., ; Luo et al., ; Liu et al., ). Another representative strategy is model compression by leveraging sparsity (Li et al., ; Ma et al., ) or quantization (He et al., ; Fang et al., ; Li et al., ; Zhao et al., ). Designing efficient diffusion model architectures (Li et al., ; Liu et al., ; Cai et al., ) or inference systems (Li et al., ; Wang et al., ) is also an effective approach for boosting efficiency. In addition, improving the data quality (Chen et al., ; ) can boost the training efficiency of diffusion models.

All these works focus on diffusion models while the autoencoder remains the same. Our work opens up a new direction for accelerating diffusion models, which can benefit both training and inference.

Figure: Figure 3: (a) High spatial-compression autoencoders are more difficult to optimize. Even with the same latent shape and stronger learning capacity, it still cannot match the f8 autoencoder’s rFID. (b) High spatial-compression autoencoders suffer from significant reconstruction accuracy drops when generalizing from low-resolution to high-resolution.
Refer to caption: x3.png

## 3 Method

In this section, we first analyze why existing high spatial-compression autoencoders (e.g., SD-VAE-f64) fail to match the accuracy of low spatial-compression autoencoders (e.g., SD-VAE-f8). Then we introduce our Deep Compression Autoencoder (DC-AE) with *Residual Autoencoding* and *Decoupled High-Resolution Adaptation* to close the accuracy gap. Finally, we discuss the applications of our DC-AE to latent diffusion models.

### 3.1 Motivation

We conduct ablation study experiments to get insights into the underlying source of the accuracy gap between high spatial-compression and low spatial-compression autoencoders. Specifically, we consider three settings with gradually increased spatial compression ratio, from f8 to f64.

Each time the spatial compression ratio increases, we stack additional encoder and decoder stages upon the current autoencoder. In this way, high spatial-compression autoencoders contain low spatial-compression autoencoders as sub-networks and thus have higher learning capacity.

Additionally, we increase the latent channel number to maintain the same total latent size across different settings. We can then convert the latent to a higher spatial compression ratio one by applying a space-to-channel operation: $H\times W\times C\rightarrow\frac{H}{p}\times\frac{W}{p}\times p^{2}C$.

We summarize the results in Figure 3 (a, gray dash line). Even with the same total latent size and stronger learning capacity, we still observe degraded reconstruction accuracy when the spatial compression ratio increases. It demonstrates that *the added encoder and decoder stages (consisting of multiple SD-VAE building blocks) work worse than a simple space-to-channel operation*.

Based on this finding, we conjecture *the accuracy gap comes from the model learning process: while we have good local optimums in the parameter space, the optimization difficulty hinders high spatial-compression autoencoders from reaching such local optimums.*

Figure: Figure 4: Illustration of Residual Autoencoding. It adds non-parametric shortcuts to let the neural network modules learn residuals based on the space-to-channel operation.
Refer to caption: x4.png

### 3.2 Deep Compression Autoencoder

#### Residual Autoencoding.

Motivated by the analysis, we introduce Residual Autoencoding to address the accuracy gap. The general idea is depicted in Figure 4. The core difference from the conventional design is that we explicitly let neural network modules learn the downsample residuals based on the space-to-channel operation to alleviate the optimization difficulty. Different from ResNet (He et al., ), the residual here is not identity mapping, but space-to-channel mapping.

In practice, this is implemented by adding extra non-parametric shortcuts on the encoder’s downsample blocks and decoder’s upsample blocks (Figure 4 b, left). Specifically, for the downsample block, the non-parametric shortcut is a space-to-channel operation followed by a non-parametric channel averaging operation to match the channel number. For example, assuming the downsample block’s input feature map shape is $H\times W\times C$ and its output feature map shape is $\frac{H}{2}\times\frac{W}{2}\times 2C$, then the added shortcut is:

$$ $\displaystyle H\times W\times C$ $\displaystyle\xrightarrow{\text{space-to-channel}}\frac{H}{2}\times\frac{W}{2} \times 4C$ $\displaystyle\underbrace{\xrightarrow{\text{split into two groups}}[\frac{H}{2 }\times\frac{W}{2}\times 2C,\frac{H}{2}\times\frac{W}{2}\times 2C]\xrightarrow {\text{average}}\frac{H}{2}\times\frac{W}{2}\times 2C.}_{\text{channel averaging}}$ $$

Accordingly, for the upsample block, the non-parametric shortcut is a channel-to-space operation followed by a non-parametric channel duplicating operation:

$$ $\displaystyle\frac{H}{2}\times\frac{W}{2}\times 2C$ $\displaystyle\xrightarrow{\text{channel-to-space}}H\times W\times\frac{C}{2}$ $\displaystyle\underbrace{\xrightarrow{\text{duplicate}}[H\times W\times\frac{C }{2},H\times W\times\frac{C}{2}]\xrightarrow{\text{concat}}H\times W\times C.} _{\text{channel duplicating}}$ $$

In addition to the downsample and upsample blocks, we also change the middle stage design following the same principle (Figure 4 b, right).

Figure 3 (a) shows the comparison with and without our Residual Autoencoding on ImageNet $256\times 256$. We can see that Residual Autoencoding effectively improves the reconstruction accuracy of high spatial-compression autoencoders.

Figure: Figure 5: Autoencoder already learns to reconstruct content and semantics without GAN loss, while GAN loss improves local details and removes local artifacts. We replace the GAN loss full training with lightweight local refinement training which achieves the same goal and has lower training cost.
Refer to caption: x5.png

Figure: Figure 6: Illustration of Decoupled High-Resolution Adaptation.
Refer to caption: x6.png

#### Decoupled High-Resolution Adaptation.

Residual Autoencoding alone can address the accuracy gap when handling low-resolution images. However, when extending it to high-resolution images, we find it not sufficient. Due to the large cost of high-resolution training, the common practice for high-resolution diffusion models is directly using autoencoders trained on low-resolution images (e.g., $256\times 256$) (Chen et al., ; ). This strategy works well for low spatial-compression autoencoders. However, high spatial-compression autoencoders suffer from a significant accuracy drop. For example, in Figure 3 (b), we can see that f64 autoencoder’s rFID degrades from 0.50 to 7.40 when generalizing from $256\times 256$ to $1024\times 1024$. In contrast, the f8 autoencoder’s rFID improves from 0.51 to 0.19 under the same setting. Additionally, we also find this issue more severe when using a higher spatial compression ratio. In this work, we refer to this phenomenon as the *generalization penalty of high spatial-compression autoencoders*. A straightforward solution to address this issue is conducting training on high-resolution images. However, it suffers from a large training cost and unstable high-resolution GAN loss training.

We introduce Decoupled High-Resolution Adaptation to tackle this challenge. Figure 6 demonstrates the detailed training pipeline. Compared with the conventional single-phase training strategy (Rombach et al., ), our Decoupled High-Resolution Adaptation has two key differences.

First, we decouple the GAN loss training from the full model training and introduce a dedicated local refinement phase for the GAN loss training. In the local refinement phase (Figure 6, phase 3), we only tune the head layers of the decoder while freezing all the other layers. The intuition of this design is based on the finding that the reconstruction loss alone is sufficient for learning to reconstruct the content and semantics. Meanwhile, the GAN loss mainly improves local details and removes local artifacts (Figure 5). Achieving the same goal of local refinement, only tuning the decoder’s head layers has a lower training cost and delivers better accuracy than the full training.

Moreover, the decoupling prevents the GAN loss training from changing the latent space. This approach enables us to conduct the local refinement phase on low-resolution images without worrying about the generalization penalty. This further reduces the training cost of phase 3 and avoids the highly unstable high-resolution GAN loss training.

Second, we introduce an additional high-resolution latent adaptation phase (Figure 6, phase 2) that tunes the middle layers (i.e., encoder’s head layers and decoder’s input layers) to adapt the latent space for alleviating the generalization penalty. In our experiments, we find only tuning middle layers is sufficient for addressing this issue (Figure 3 b) while having a lower training cost than high-resolution full training (memory cost: 153.98 GB $\rightarrow$ 67.81 GB)(^1^11Assuming the input resolution is $1024\times 1024$ and the batch size is 12.) (Cai et al., ).

**Table 1: Ablation Study on Patch Size and Autoencoder’s Spatial Compression Ratio.**
| ImageNet 512$\times$512 (Class-Conditional) |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Diffusion Model | Autoencoder | Patch Size | #Tokens | FID (w/o CFG) $\downarrow$ | FID (w/ CFG) $\downarrow$ |
|  | SD-VAE-f8 [] | 4 | 256 | 37.42 | 15.61 |
|  | SD-VAE-f16 [] | 2 | 256 | 36.22 | 15.17 |
| DiT-XL [] | SD-VAE-f32 [] | 1 | 256 | 32.98 | 12.95 |
|  | SD-VAE-f8 | 8 | 64 | 125.08 | 95.93 |
|  | SD-VAE-f16 | 4 | 64 | 115.32 | 88.06 |
|  | SD-VAE-f32 | 2 | 64 | 107.33 | 76.57 |
| UViT-S [] | DC-AE-f64 | 1 | 64 | 67.30 | 35.96 |

### 3.3 Application to Latent Diffusion Models

Applying our DC-AE to latent diffusion models is straightforward. The only hyperparameter to change is the patch size (Peebles & Xie, ). For diffusion transformer models (Peebles & Xie, ; Bao et al., ), increasing the patch size $p$ is the common approach for reducing the number of tokens. It is equivalent to first applying the space-to-channel operation to reduce the spatial size of the given latent by $p\times$ and then using the transformer model with a patch size of 1.

Since combining a low spatial-compression autoencoder (e.g., f8) with the space-to-channel operation can also achieve a high spatial compression ratio, a natural question is how it compares with directly reaching the target spatial compression ratio with DC-AE.

We conduct ablation study experiments and summarize the results in Table 1. We can see that directly reaching the target spatial compression ratio with the autoencoder gives the best results among all settings. In addition, we also find that shifting the spatial compression ratio from the diffusion model to the autoencoder consistently leads to better FID. We conjecture it is because *the latent diffusion models need to simultaneously learn denoising and token compression when using a patch size $>1$. With the autoencoder taking over the whole token compression task, the diffusion model can fully focus on the denoising task and thus achieves better results.*

## 4 Experiments

**Table 2: Image Reconstruction Results.**
| ImageNet 256$\times$256 | Latent Shape | Autoencoder | rFID $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| f32c32 | 8$\times$8$\times$32 | SD-VAE [] | 2.64 | 22.13 | 0.59 | 0.117 |
| DC-AE | 0.69 | 23.85 | 0.66 | 0.082 |  |  |
| f64c128 | 4$\times$4$\times$128 | SD-VAE [] | 26.65 | 18.07 | 0.41 | 0.283 |
| DC-AE | 0.81 | 23.60 | 0.65 | 0.087 |  |  |
| ImageNet 512$\times$512 | Latent Shape | Autoencoder | rFID $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ |
| f64c128 | 8$\times$8$\times$128 | SD-VAE [] | 16.84 | 19.49 | 0.48 | 0.282 |
| DC-AE | 0.22 | 26.15 | 0.71 | 0.080 |  |  |
| f128c512 | 4$\times$4$\times$512 | SD-VAE [] | 100.74 | 15.90 | 0.40 | 0.531 |
| DC-AE | 0.23 | 25.73 | 0.70 | 0.084 |  |  |
| FFHQ 1024$\times$1024 | Latent Shape | Autoencoder | rFID $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ |
| f64c128 | 16$\times$16$\times$128 | SD-VAE [] | 6.62 | 24.55 | 0.68 | 0.237 |
| DC-AE | 0.23 | 31.04 | 0.83 | 0.061 |  |  |
| f128c512 | 8$\times$8$\times$512 | SD-VAE [] | 179.71 | 18.11 | 0.63 | 0.585 |
| DC-AE | 0.41 | 31.18 | 0.83 | 0.062 |  |  |
| MapillaryVistas 2048$\times$2048 | Latent Shape | Autoencoder | rFID $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ |
| f64c128 | 32$\times$32$\times$128 | SD-VAE [] | 7.55 | 22.37 | 0.68 | 0.262 |
| DC-AE | 0.36 | 29.57 | 0.84 | 0.075 |  |  |
| f128c512 | 16$\times$16$\times$512 | SD-VAE [] | 152.09 | 17.82 | 0.67 | 0.594 |
| DC-AE | 0.38 | 29.70 | 0.84 | 0.074 |  |  |

### 4.1 Setups

#### Implementation Details.

We use a mixture of datasets to train autoencoders (baselines and DC-AE), containing ImageNet (Deng et al., ), SAM (Kirillov et al., ), MapillaryVistas (Neuhold et al., ), and FFHQ (Karras et al., ). For ImageNet experiments, we exclusively use the ImageNet training split to train autoencoders and diffusion models. The model architecture is similar to SD-VAE (Rombach et al., ) except for our new designs discussed in Section 3.2. In addition, we use the original autoencoders instead of the variational autoencoders for our models, as they perform the same in our experiments and the original autoencoders are simpler. We also replace transformer blocks with EfficientViT blocks (Cai et al., ) to make autoencoders more friendly for handling high-resolution images while maintaining similar accuracy.

For image generation experiments, we apply autoencoders to diffusion transformer models including DiT (Peebles & Xie, ) and UViT (Bao et al., ). We follow the same training settings as the original papers.
We consider three settings with different resolutions, including ImageNet (Deng et al., ) for $512\times 512$ generation, FFHQ (Karras et al., ) and MJHQ (Li et al., ) for $1024\times 1024$ generation, and MapillaryVistas (Neuhold et al., ) for $2048\times 2048$ generation.

**Table 3: Class-Conditional Image Generation Results on ImageNet 512$\times$512. ^† represents the model is trained for 4$\times$ training iterations (i.e., 500K $\rightarrow$ 2,000K iterations).**
| Diffusion |  | Patch | Throughput (image/s) $\uparrow$ | Latency | Memory | FID $\downarrow$ |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Autoencoder | Size | Training | Inference | (ms) $\downarrow$ | (GB) $\downarrow$ | w/o CFG | w/ CFG |
|  | Flux-VAE-f8 [] | 2 | 352 | 2984 | 3.8 | 13.8 | 106.07 | 84.73 |
|  | SD-VAE-f8 [] | 2 | 352 | 2991 | 3.8 | 13.8 | 51.96 | 24.57 |
|  | SD-VAE-f16 [] | 2 | 1550 | 12881 | 1.3 | 4.0 | 76.86 | 44.22 |
|  | SD-VAE-f32 [] | 1 | 1551 | 12883 | 1.3 | 4.0 | 70.23 | 38.63 |
|  | DC-AE-f32 | 1 | 1553 | 12850 | 1.3 | 4.0 | 46.12 | 18.08 |
|  | DC-AE-f64 | 1 | 6295 | 53774 | 0.7 | 1.5 | 67.30 | 35.96 |
| UViT-S [] | DC-AE-f64^† | 1 | 6295 | 53774 | 0.7 | 1.5 | 61.84 | 30.63 |
|  | SD-VAE-f8 [] | 2 | 54 | 424 | 31.7 | 56.2 | 12.03 | 3.04 |
| DiT-XL [] | DC-AE-f32 | 1 | 241 | 2016 | 7.8 | 20.9 | 9.56 | 2.84 |
|  | Flux-VAE-f8 [] | 2 | 55 | 349 | 30.4 | 54.2 | 30.91 | 12.63 |
|  | SD-VAE-f8 [] | 2 | 55 | 351 | 30.3 | 54.1 | 11.04 | 3.55 |
|  | DC-AE-f32 | 1 | 247 | 1622 | 8.2 | 18.6 | 9.83 | 2.53 |
|  | DC-AE-f64 | 1 | 984 | 6706 | 3.5 | 10.6 | 13.96 | 3.01 |
| UViT-H [] | DC-AE-f64^† | 1 | 984 | 6706 | 3.5 | 10.6 | 12.26 | 2.66 |
|  | SD-VAE-f8 [] | 2 | 27 | 157 | 74.8 | OOM | 9.73 | 3.57 |
|  | DC-AE-f32 | 1 | 112 | 665 | 19.7 | 42.0 | 8.13 | 2.30 |
|  | DC-AE-f64 | 1 | 450 | 2733 | 8.6 | 30.2 | 7.78 | 2.47 |
| UViT-2B [] | DC-AE-f64^† | 1 | 450 | 2733 | 8.6 | 30.2 | 6.50 | 2.25 |

**Table 4: 1024$\times$1024 and 2048$\times$2048 Image Generation Results.**
| FFHQ 1024$\times$1024 (Unconditional) & MJHQ 1024$\times$1024 (Class-Conditional) |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Diffusion |  | Patch | Throughput (image/s) $\uparrow$ | Latency | Memory | FFHQ FID $\downarrow$ | MJHQ FID $\downarrow$ |  |  |
| Model | Autoencoder | Size | Training | Inference | (ms) $\downarrow$ | (GB) $\downarrow$ | w/o CFG | w/o CFG | w/ CFG |
|  |  | 2 | 84 | 833 | 14.1 | 41.2 | 16.98 | 48.05 | 38.19 |
|  | SD-VAE-f8 [] | 4 | 470 | 5566 | 2.5 | 10.7 | 23.81 | 60.94 | 51.29 |
|  | DC-AE-f32 | 1 | 475 | 5575 | 2.5 | 10.7 | 13.65 | 34.35 | 27.20 |
| DiT-S [] | DC-AE-f64 | 1 | 2085 | 25259 | 1.0 | 3.1 | 26.88 | 61.30 | 53.38 |
| MapillaryVistas<br>2048$\times$2048 (Unconditional) |  |  |  |  |  |  |  |  |  |
| Diffusion |  | Patch | Throughput (image/s) $\uparrow$ | Latency | Memory | MapillaryVistas FID $\downarrow$ |  |  |  |
| Model | Autoencoder | Size | Training | Inference | (ms) $\downarrow$ | (GB) $\downarrow$ | w/o CFG |  |  |
|  | SD-VAE-f8 [] | 4 | 84 | 810 | 14.3 | 41.4 | 69.50 |  |  |
| DiT-S [] | DC-AE-f64 | 1 | 459 | 5435 | 2.6 | 11.0 | 59.55 |  |  |

**Table 5: Text-to-Image Generation Results.**
| Diffusion |  | Patch | Throughput (image/s) $\uparrow$ | Latency | Memory | MJHQ 512$\times$512 |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Autoencoder | Size | Training | Inference | (ms) $\downarrow$ | (GB) $\downarrow$ | FID $\downarrow$ | CLIP Score $\uparrow$ |
|  | SD-VAE-f8 [] | 2 | 43 | 312 | 37.1 | 60.45 | 6.3 | 26.36 |
| PIXART-$\alpha$ [] | DC-AE-f32 | 1 | 173 | 1251 | 10.4 | 23.77 | 6.1 | 26.41 |

Figure: Figure 7: Autoencoder Image Reconstruction Samples.
Refer to caption: x7.png

Figure: Figure 8: Images Generated by Diffusion Model using Our DC-AE.
Refer to caption: x8.png

#### Efficiency Profiling.

We profile the training and inference throughput on the H100 GPU with PyTorch and TensorRT respectively. The latency is measured on the 3090 GPU with batch size 2. The training memory is profiled using PyTorch, assuming a batch size of 256. We use fp16 for all cases. For simplicity, we assume the number of sampling steps is 1.

### 4.2 Image Compression and Reconstruction

Table 2 summarizes the results of DC-AE and SD-VAE (Rombach et al., ) under various settings (f represents the spatial compression ratio and c denotes the number of latent channels). DC-AE provides significant reconstruction accuracy improvements than SD-VAE for all cases. For example, on ImageNet $512\times 512$, DC-AE improves the rFID from 16.84 to 0.22 for the f64c128 autoencoder and 100.74 to 0.23 for the f128c512 autoencoder.

In addition to the quantitative results, Figure 7 shows image reconstruction samples produced by SD-VAE and DC-AE. Reconstructed images by DC-AE demonstrate a better visual quality than SD-VAE’s reconstructed images. In particular, for the f64 and f128 autoencoders, DC-AE still maintains a good visual quality for small text and the human face.

### 4.3 Latent Diffusion Models

We compare DC-AE with the widely used SD-VAE-f8 autoencoder (Rombach et al., ) on various diffusion transformer models. For DC-AE, we always use a patch size of 1 (denoted as p1). For SD-VAE-f8, we follow the common setting and use a patch size of 2 or 4 (denoted as p2, p4). The results are summarized in Table 3, Table 4, and Table 5.

Figure: Figure 9: Model Scaling Results on ImageNet 512$\times$512 with UViT. DC-AE-f64 benefits more from scaling up than SD-VAE-f8.
Refer to caption: x9.png

#### ImageNet 512 × \times × 512.

As shown in Table 3, DC-AE-f32p1 consistently delivers better FID than SD-VAE-f8p2 on all diffusion transformer models. In addition, it has 4$\times$ fewer tokens than SD-VAE-f8p2, leading to 4.5$\times$ higher H100 training throughput and 4.8$\times$ higher H100 inference throughput for DiT-XL. We also observe that larger diffusion transformer models seem to benefit more from our DC-AE. For example, DC-AE-f64p1 has a worse FID than SD-VAE-f8p2 on UViT-S but a better FID on UViT-H. We conjecture it is because DC-AE-f64 has a larger latent channel number than SD-VAE-f8, thus needing more model capacity (Esser et al., ).

#### 1024 × \times × 1024 and 2048 × \times × 2048 Image Generation.

Apart from ImageNet 512$\times$512, we also test our models for higher-resolution image generation. As shown in Table 4, we have a similar finding where DC-AE-f32p1 achieves better FID than SD-VAE-f8p2 for all cases.

#### Text-to-Image Generation.

Table 5 reports our text-to-image generation results. All models are trained for 100K iterations from scratch. Similar to prior cases, we observe DC-AE-f32p1 provides a better FID and a better CLIP Score than SD-VAE-f8p2. Figure 8 demonstrates samples generated by the diffusion models with our DC-AE, showing the capacity to synthesize high-quality images while being significantly more efficient than prior models.

## 5 Conclusion

We accelerate high-resolution diffusion models by designing deep compression autoencoders to reduce the number of tokens. We proposed two techniques: residual autoencoding and decoupled high-resolution adaptation to address the challenges brought by the high compression ratio. The resulting new autoencoder family DC-AE demonstrated satisfactory reconstruction accuracy with a spatial compression ratio of up to 128. DC-AE also demonstrated significant training and inference efficiency improvements when applied to latent diffusion models.

## Acknowledgements

We thank NVIDIA for donating the DGX machines. We thank MIT-IBM Watson AI Lab, MIT and Amazon Science Hub, MIT AI Hardware Program, and National Science Foundation for supporting this research.

## Appendix A Additional Samples

Figure: Figure 10: Additional Autoencoder Image Reconstruction Samples.
Refer to caption: x10.png

Figure: Figure 11: Additional Autoencoder Image Reconstruction Samples.
Refer to caption: x11.png

## Appendix B Diffusion Model Architecture Details

In addition to existing UViT models, we scaled the model up to 1.6B parameters, with a depth of 28, a hidden dimension of 2048, and 32 heads. We denote this model as UViT-2B.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 6
presentation: 2
questions: * While increasing the patch size is one method to reduce spatial resolution, one can also use a pixel shuffle operation in order to achieve a lower spatial dimension with a higher channel dimension. How does this simple operation, using SD-VAE compare with your method ?
* Training LDMs requires a shift and scale factor to be applied to the latents before the diffusion process, how are these values computed in your experiments ?
soundness: 2
strengths: * The paper is well written and easy to follow.
* The method allows to significantly reduce the computational requirements to train large scale diffusion models at high resolution, making it more accessible as a research topic.
* Evaluation of multiple metrics across 4 different datasets are reported.
* Evaluations on downstream diffusion training is also provided.
* Qualitative examples clearly showcase the improvements brought about by the proposed model.
summary: This paper presents a family of autoencoders that achieve comparable or better reconstructions compared to the SD autoencoders while having a much stronger spatial compression rate, resulting in efficiency gains for downstream tasks such as latent diffusion training.
Using the techniques or residual autoencoding and decoupling high-resolution adaptation, the model is able to generalize well to higher resolutions while having a very large compression factor (x32, 64 vs 8 for SD-VAE).
weaknesses: * What parameters were used for sampling the diffusion models (sampler, number of sampling steps, guidance scale) ? A more through investigation on the impact on sampling quality would be useful to get a better grasp of the limitations of this method.
* Missing ablations on the constituent parts of DC-AE.
* In tables 3 and 4, missing comparisons with more recent autoencoders such as the one from SD-XL, SD3, and Asymetric Autoencoder.
* In table 5. the PixArt model trained with DC-AE achieves a lower CLIP score than the one trained with SD-VAE while the text says otherwise (maybe a typo?). What if the memory usage is equalized with the SD-VAE by increasing the batch size, do the improvements get larger or do they stagnate ?

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 4
questions: Regarding generalization across resolutions, I'm wondering how much changes to the training procedure can close the gap even with single-stage training. In particular, it makes sense that the statistics for a fixed-size patch (say 32x32 pixels) is very different between a 256x256 image and 1024x1024 pixel image of the same scene. You could train on 256x256 *crops* from a 1024x1024 dataset (rather than downscaling), but then you'd have a generalization gap in the other direction.

In other work, researchers try to avoid this problem by generating low-res training sets by randomly downscaling high-res sets. With enough training, this approach lets the model see the full range of statistics. Was this done for the experiments in the DC-AE paper? If not, do you think it's sufficient or perhaps complementary to DHRA? Either way, I still see potential benefits of DHRA, e.g., maybe it's good that the GAN doesn't influence the encoder (and thus the latent space), but this isn't investigated, in isolation, in the paper.

You can also frame higher spatial-compression rates as meaning that the model sees less training data (as measured by number of latents) per fixed-size image. This means that any negative affects from the boundary are exacerbated (e.g., if zero-padding is used for conv layers). In the extreme, imagine using f256 with 256x256 training data -- we certainly wouldn't expect this model to generalize to higher resolutions. Some discussion of these effects would help, and investigation of whether the generalization issue is mostly one of different pixel statistics or of boundary artifacts would strengthen the paper.

Did you experiment training different numbers of layers in Stage 2 and Stage 3? A graph showing the impact (training speed vs. evaluation metrics) as the number of layers increases would be interesting.
soundness: 3
strengths: The primary strength of this paper is the development of an effective and relatively simple solution to a real-world deficiency that impacts all research and products using autoencoders. Specifically, the use of residual autoencoding and DHRA allow either better visual quality (e.g., in Table 3 look at rows witht he same number of tokens like SD-VAE-f16 with patch-size=2 or SD-VAE-f32 with patch-size=1 to DC-AE-f32) or similar visual quality with lower latency and higher throughput (e.g., Table 3 comparing SD-VAE-f32 to DC-AE-f64).

Another strength is that the authors evaluate their method on both reconstruction and generation problems. Improvements on just reconstruction still have some value, but many improvements to reconstruction quality do not translate to gains on the generation side.
summary: This paper looks at the problem of scaling image autoencoders, used with generative models, to higher spatial compression rates without reducing visual quality. The authors identify a problem with existing AEs: visual quality degrades with higher spatial compression rates (Fig. 2) even when increasing the number of channels in the latent representation to maintain a fixed total size (see the "latent shape" column in Table 2). They also observe a quality degradation when generalizing from low resolutions (used for training) to higher-resolutions.

These observations imply that "high spatial-compression autoencoders are more difficult to optimize" so the paper introduces an architecture change ("residual autoencoding") and a training change ("decoupled high-resolution adaptation") to improve results with high spatial compression rates and to reduce generalization error across resolutions.

Residual autoencoding (Fig. 4) replaces strided conv blocks used for up/downscaling with residual blocks that learn a residual added to space-to-depth (for downscaling) or space-to-depth (for upscaling) blocks. The model uses channel averaging to adjust the number of channels when there is a mismatch, and authors also use channel averaging for a new kind of residual block even when there is no up/downscaling.

Decoupled high-resolution adaptation (DHRA) switched from a single-stage approach for training the AE to a three-stage appraoch (Fig. 6). In the standard approach, the AE is trained on low-res data using all reconstruction losses (typically mse + perceptual + adversarial). In DHRA, the stages are:
 1. low-res images optimized without the adversarial loss
 2. finetune with high-res images with only the "inner" layers (end of encoder, beginning of the decoder) trained. Other layers are frozen and, again, no adversarial loss is used.
 3. finetune again with reconstruction and adversarial losses. Only the final layers of the decoder are trained.

The authors argue that this multi-stage approach is more efficient than training the full model on high-res images, and it prevents the adversarial loss from affecting the encoder which helps with generalization.

The authors validate their claims by training models and evaluating reconstruction quality (Table 2) and generation quality (Tables 3 and 4). They also provide some example images for qualitative evaluation (Fig. 7 and in the supplementary material).
weaknesses: A more detailed summary of the DC-AE architecture would help. Fig. 4 provides a high-level overview, but I'm not sure how literal or complete it is, and the details of the "Encoder/Decoder Stages" are not provided. To be fair, the authors do provide code, but architecture and training details should be in the paper (probably in the appendix).

The conjecture in Section 3.3 that DC-AE helps because otherwise the "diffusion model needs to simultaneously learn denoising and token compression when using a patch size > 1" is pretty hand-wavy. For example, I can always set patch_size=1 and add add a space-to-depth layer to the encoder and claim that now the encoder alone is responsible for token compression. That said, I don't have a better explanation, but I think it's ok for papers to accept that the contribution is a better architecture justified with empirical evaluation.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 4
questions: My overall rating is due to the somewhat limited set of modifications and the limited experiments, which should be a topic for rebuttal. In addition, I also have the following questions:

1. Did you adopt all the same loss functions and GAN architectures as in the stable diffusion paper?
2. How often does the proposed architecture output artifacts?
3. Are the images in all figures random-samples or selected? This should be made clear in the text.
4. Did you consider other operations for space-to-channel other than averaging? Same question for channel duplication in the upsampling process.
soundness: 4
strengths: 1. The paper is well-written and easy to follow.
2. The baselines using stable diffusion are well-known and recognized by the community.
3. The qualitative examples seem compelling.
4. The advantages with using less compute could be of significant impact.
summary: The paper presents a new type of autoencoder for high resolution generative diffusion models. Autoencoders are used in generative diffusion models primarily to reduce the computation necessary for training and evaluation. However, there is some loss of representation capacity in mapping the image to the latent domain. The loss of representation capacity becomes more severe with higher downsampling ratios, which limits compute gains beyond those provided by 8X downsampling, as well as high-resolution generation.

The present paper aims to expand the possible range of downsampling ratios to as high as 64X. The paper achieves this first by analyzing the challenges of higher downsampling ratios, finding that staged architectures actually have worse performance than architectures with simple space-to-channel operations. Following this, the paper proposes a new architecture with residual space-to-channel and channel-to-space connections that effectively skip the information across network blocks, improving optimization High-resolution performance after adopting the residual space-to-channel operations still left some to be desired, so the paper further proposes a decoupled high-resolution adaptation procedure that finetunes the head of the decoder while finetuning the rest of the model. A second stage tunes the middle layers.

In numerical experiments the trained autoencoders exhibit better reconstruction accuracy than their counterparts from stable diffusion, which are commonly used in the community. The full diffusion models show promising FID numbers vs. stable diffusion with more efficient training in both quantitative and qualitative examples.
weaknesses: 1. The authors mostly use stable diffusion for all experiments, although several datasets are considered.
2. A more detailed quantitative analysis of why standard staged training has difficulties with gradient propagation is not presented.
3. The proposed modifications are somewhat incremental - this seems close to simply removing some skip connection convolutions + a reshape operation in the standard stable diffusion autoencoder. The extra fine-tuning procedure is more detailed, but is only necessary due to the autoencoder changes.
4. Autoencoder comparisons in Table 2 are only done at high downsampling ratios - not sure how the proposed architecture performs with shallow downsamples.
5. Only the FID metric (presumably with the Inception V3 backbone) is used to evaluate the models.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - What is the role of the averaging compared to simple space to channel? How would the model behave without the averaging? Would the channel dimension grow too fast?
- How does the use of EfficientViT blocks affect the quality (reconstruction/generation) of the method? Does it lead to better behavior when trained with large subsampling factors compared to the baseline?
- “..., we assume the number of sampling steps is 1.”: Does this also hold for inference throughput and latency? If so, I would suggest reporting the throughput/latency to sample a full image.
- What is the resolution of the images in Figure 7? It would generally be helpful to indicate the resolution for the visual examples in the caption.
soundness: 3
strengths: - The proposed architectural modifications are simple and sound, and the corresponding multi-step training recipe is relatively simple as well.
- The speed-ups are quite substantial.
- The method is tested on a broad range of data sets (including text-to-image modeling), and in combination with different latent diffusion models. So the results are quite comprehensive.
summary: The paper proposes a new VAE architecture for latent diffusion generative modeling along with a multi-step training recipe. More specifically, the proposed design combines space-to-channels with averaging to build deep autoencoders with a high spatial subsampling factor (up to 64x). In order to train these models effectively at high resolution, they are first trained at low resolution without GAN loss, then partially tuned at higher resolution, and then refined with a GAN loss. The more aggressive subsampling leads to lower spatial resolution/sequence length in the latent space than VAEs from the literature and hence leads to substantial speed-ups of about 4x to 16x while maintaining generation quality, in combination with different latent diffusion models.
weaknesses: - The role of the space-to-channels-based autoencoder design and the multi-step training recipe are entangled. How does the baseline autoencoder (with more aggressive subsampling) perform, when it is trained with the multi-step recipe?
- The paper seems to lack important experimental details, including optimizer, learning rate, schedule etc. These are important in particular given the multi-stage training with different losses (including a GAN loss). Also, what losses does the method use exactly? It looks like the low resolution stage uses a perceptual loss (Figure 5). 

\
Minor:
- In the side-by-side comparisons it would be useful to have the original image as additional reference.
- Some prior works discuss the relevance of space-to-channel in image enhancement/generation methods, probably worth mentioning (see e.g. Shi et al. "Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network." CVPR 2016)

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: -	Could you please clarify the experimental settings and AE architectures in Section 3.1? By Line 161 ‘… convert the latent to higher spatial compression…’, do you mean that the f8 to f64 AEs have the same output latent size initially, yet f32 and f64 AEs further squeeze the space dimension into channel dimension?
-	In Table 5, the proposed method shows a lower CLIP score with the Pixart diffusion transformer, is it a typo?
soundness: 3
strengths: -	Novel autoencoder design for high compression ratios. The authors introduce additional non-parametric residual connections in the downsampling and upsampling layers, which facilitate the optimization of autoencoders in high spatial compression ratio scenarios.
-	The authors carefully design a three-stage training recipe to mitigate the high-resolution generalization gaps by fine-tuning the trained AE on high-resolution data with minimal additional costs.
-	The proposed DC-AEs show a substantial reduction in training and inference costs of latent diffusion models on various datasets.
summary: The submission presents Deep Compression Autoencoder (DC-AE), a new series of autoencoders that enable highly efficient high-resolution image synthesis with latent diffusion models. The key architecture design is the adding residual connection with space-channel inter-play in both downsampling and upsampling layers, which alleviates the optimization difficulty significantly. They further propose decoupled high-resolution adaptation, a three-phase training strategy that mitigates the generalization penalty of high spatial-compression autoencoders. By improving the autoencoder's spatial compression ratio up to 128 while maintaining reconstruction quality, DC-AE provides significant speedups for both training and inference of latent diffusion models without sacrificing image generation accuracy. The authors demonstrate the effectiveness of DC-AE on various datasets and diffusion transformer architectures.
weaknesses: -	Although the VAE shows little FID loss as it progresses to a higher spatial compression ratio (DC-AE-f32 to DC-AE-f64), it does not translate to similar results in the corresponding latent diffusion models as shown in Table 3. The underlying reason is probably that a diffusion model with higher capacity is required when dealing with high compression ratio latents.
-	The current submission lacks either theoretical analysis or illustrative toy examples on the proposed design choice of space-channel residual connection, e.g., comparisons on the loss landscapes.

Meta-review: This paper explores image auto-encoders with a high spatial compression rate in the context of image generation. The proposed model results in efficiency gains while maintaining downstream performance. The manuscript was reviewed by four knowledgeable reviewers who acknowledged that the paper was well written and easy to follow (jBHL, 9DQC), the proposed solution was relatively simple, sound and effective (rPFh, pFgo), offering significant speedups (pFgo) and reduction of resources required to train diffusion models (jBHL), highlighting the potential of the proposed approach (9DQC).

The reviewers's main concerns were:
1. Missing ablations (jBHL) and analyses on the generalization across resolutions (rPFh)
2. Missing architecture details (rPFh)
3. Proposed modifications somewhat incremental (9DQC), and no theoretical analyses or toy examples highlighting the potential of the proposed design (iQFp)
4. Only FID metric used for evaluation (9DQC)

During rebuttal and discussion, the authors emphasized novelty and technical contribution, added comparisons with more auto-encoders and ablations, as well as metrics (PR, IS, CMMD), shared additional details on the method and experimental setup, and scaled some experiments. After discussion, the reviewers unanimously recommend acceptance. The MR agrees with the reviewers' assessment and recommends to accept.

Decision: Accept (Poster)