Paper: Würstchen: An Efficient Architecture for Large-Scale Text-to-Image Diffusion Models (ICLR-2024)

Reviewer comment (Evaluation):
I'm uncertain about the inference efficiency of this approach, as it appears to add an "extra" computation (Stage C) on top of the conventional LDM and SD (Stage B and Stage A). In particular, how could the proposed method achieve better inference time than SD-v2.1 in Figure 4? A detailed computational comparison would be beneficial for different components in the system (the text encoder, LDM(s), and image decoder) instead of just an overall process.

Author response:
The compute to process a sampling step scales quadratically with resolution. Stage C, operates on very low resolution, making sampling steps much more efficient. Stage B, is strongly guided by the latents of Stage C (see Appendix E) and thus needs a much lower number of diffusion steps (12, in our experiments) to produce high-fidelity results.
However, we acknowledge that this should be communicated better. To facilitate an easier understanding of this, we made the following changes to the paper:

First, we provide a more detailed plot on the inference-time, splitting up the time spent on the different inference stages in Fig. 4. Second, we overhauled Tab. 2, to show the properties of the individual Stages B and C, instead of only showing Stage C (more on this later).

Reviewer comment (Experiments):
I think the Baseline LDM (trained for 25,000 GPU-hours (same as Stage C)) needs to be trained for GPU hours of Stage B + Stage C, given that both stages contribute to the final latent representation of the proposed method. More importantly, a baseline with the same architecture of the upper part in Stage B, Figure 3 (i.e., a conventional LDM obtained by just removing Stage C and the below part of Stage B, Figure 3) seems necessary to show the benefit of the proposed approach.

Author response:
We also discussed this internally before submission, as both models are diffusion models. We eventually decided to train our baseline with Stage C’s budget, as Stage C is the image generator (see Appendix E)
Training a second baseline is not possible for us, as the 36K GPU-hours is outside our available budget and would require approx. 3 weeks to complete, not including evaluation.
As a compromise, we are in the process of fine-tuning Baseline LDM for 11K  additional gpu-hours to match the compute of Stage B+C training. We plan to add this model as a second baseline to the camera-ready version, as the results will unfortunately not be ready in the discussion period.

Reviewer comment (Presentation):
The parameter values in Table 2 might confuse readers due to inconsistencies in their presentation. For some models, like LDM, the table seems to consider all the parameters, including the text encoder. Yet, for other models such as the proposed method and SD, only the diffusion parameters are listed. I strongly suggest presenting the "total" parameters (because several components work together for a single text-to-image system) or, preferably, detailing both the "total" and diffusion parameters separately.

Author response:
Thank you for pointing this out. Based on your feedback, we added CLIP-Score as well as separating total and generative parameters. We list Stage B and C as requested.

Following your comment, we again updated our manuscript and changed Table 2 to differentiate between parameters of the generative model (approx. 1B+1B for Würstchen Stage B+StageC) and the total number of parameters (approx. 2.7B). We also added the total number of parameters of SDXL (3.4B).

Due to the double-blind policy, we cannot comment on open-source software you may have used to compute those numbers. However, an older development version of our model trained earlier this year used two different text encoders, due to Stage B and Stage C having overlapping development cycles, which resulted in an outdated version of Stage B being used to train a newer version of Stage C, with the key difference being the text encoders. 

However, our current model architecture, as presented in the paper, is designed to have only a single text encoder that is shared by both stages. This is also reflected in the parameter count (2.7B parameters vs. 3B.). 

Nevertheless, we are **very grateful** for your comment, as this discussion sparked an internal debate about the relevance of Stage B text-conditioning, which we suspected to have only a minor impact. To quantify the impact of text-conditioning on Stage B and C we conducted another ablation study, which has been added as Appendix J to the revised supplementary material.

This additional ablation study compares the image quality of our original model with two variants that only receive text-conditioning in Stage B or C, respectively. These variants are created by simply replacing the text embeddings produced by the text encoder model with zero-tensors of the same shape. The model was not adapted, fine-tuned or retrained in any other way, the entire experiment was inference-only.

We find that removing the text-conditioning from the trained Stage B has, in our assessment, a negligible impact on the generative quality of the model. Automated metrics register slight changes (a slight decrease of 49.4%/49.6% in PickScore, where 50% would be expected for identical sets, and, in contrast, a slight improvement in FID), which we think can well be attributed to random effects. In contrast, removing text-conditioning from Stage C results in a catastrophic loss of alignment with the text prompt, demonstrating its importance and reinforcing our understanding that Stage B acts as a latent super-resolution module for Stage C.

This indicates to us that future versions of Stage B do not need to be trained on text-conditioning in Stage B at all and that even in existing versions text-conditioning in Stage B can be dropped without significant changes in generative quality. While this insight is not directly relevant to your question regarding parameters, we found it interesting enough to bring it up anyway, since it highlights potential future improvements to our architecture, further boosting training efficiency. 

**Therefore, we, again, would like to thank you for the great diligence and high effort spent for the review. A discussion like this, in our mind, really shows the strengths of the peer review system.**

For the sake of reproducibility, we also provide the source code used to disable the Stage B conditioning. This will work likewise for the older version with dual text encoders that might have been used by the reviewer:

```
class TextEncoderWarapper:

   def __init__(self, to_wrap):
       self.to_wrap = to_wrap

   def __call__(self, *args, **kwargs):
       x = self.to_wrap(*args, **kwargs)
       x.last_hidden_state = x.last_hidden_state * 0
       x.pooler_output = x.pooler_output * 0
       return x

   @property
   def dtype(self):
       return self.to_wrap.dtype

pipeline = AutoPipelineForText2Image.from_pretrained(weight_path, torch_dtype=torch.float16).to(device)
if compile:
   pipeline.prior_prior = torch.compile(pipeline.prior_prior, mode="reduce-overhead", fullgraph=True)
   pipeline.decoder = torch.compile(pipeline.decoder, mode="reduce-overhead", fullgraph=True)

pipeline.set_progress_bar_config(leave=True)

# for disabling Stage B text-conditioning
pipeline.decoder_pipe.text_encoder = TextEncoderWarapper(pipeline.decoder_pipe.text_encoder)

# for disabling Stage C text-conditioning
#pipeline.prior_pipe.text_encoder = #TextEncoderWarapper(pipeline.prior_pipe.text_encoder)
```

Reviewer comment (Evaluation):
The popular MS-COCO benchmark has been conducted at the resolution of 256x256. Why did the authors change the resolution for IS in Table 2? In my experience, the resolution affects the metric scores. Furthermore, for some models (LDM, DALL-E, CogView), the IS results at 256x256 were reported. I also highly recommend including CLIP score.

Author response:
Concerning the CLIP-Score, we agree that this is a missing piece in our automated evaluation, we added it to Tab. 2 and more results can be found in Appendix B of the revised manuscript. We found only minor changes in CLIP score, especially between the SD models and Würstchen.
Concerning IS-Score, the generated images used to compute the IS are the same images used to compute the FID-score. We think this is a misunderstanding, as a matter of fact, we used the code of CogView in particular for computing IS to make reproducibility more straightforward: https://tinyurl.com/f8mn9f2a 
We write 299x299 since In line 36 of the linked, public source-file, images are resized to 299x299 pixels, hence ensuring all images are evaluated at the same resolution (which is, in fact, the fixed-size input resolution of the used InceptionV3 model). Because of this implementation quirk, we decided to report the “true” input resolution of the Inception model in the paper, to enhance clarity and reproducibility.

Reviewer comment (Theory):
I think the description “By conditioning Stage B on low-dimensional latent representations, we can effectively decode images from a 16x24x24 latent space to a resolution of 3x1024x1024, resulting in a total spatial compression of 42:1” in page 5 looks incorrect or overclaimed, because Stage B also takes a high-resolution latent map, 4x256x256, as input.

Author response:
We think the core of this misunderstanding lies in the question of how Stage B and C interact. Intuitively, Stage C conditioning Stage B suggests that Stage B is generating the image while Stage C is merely guiding this process, or both stages forming a complex generative cascade. However, Stage B should be viewed as a latent super resolution model (employing denoising for the decoding process), not as a generative model. Stage B samples 4x256x256 latents, which  are initialized to random Gaussian noise. The semantic compressor, on the other hand, whose image representation lies at a 16x24x24 space, is not noised during sampling of Stage B. Effectively, the compressed 16x24x24 latents are passed as a conditioning to Stage B to be upsampled to 4x256x256 latents, which are then decoded to 3x1024x1024 by Stage A. This results in an image in pixel space which is 42 times larger that the latent representation of the same image in Stage C. In Appendix E we provide further empirical evidence for this, by decoding the image directly from Stage C’s 16x24x24 latents using a simple convolutional network without diffusion, which, as we show, directly yields a lower-resolution version of the same image as when decoding the latents with Stage B and Stage A, demonstrating that Stage B indeed only acts as a superresolution upsampler for the encoded image.

Reviewer comment (Experiments):
The behavior of the proposed model seems less explored. The representative analysis with different classifier-free guidance scales to show the tradeoff between FID-CLIP score [SD, GLIDE, Imagen] is missing. Furthermore, it would be interesting to analyze the tradeoff between the number of sampling steps and generation quality.

Author response:
We agree that a more in-depth analysis of hyperparameters is valuable. We hence addressed this in two ways: We added Appendix I containing FID-CLIP plots for classifier-free-guidance and sampling steps. We also added Appendix H, which is a more detailed analysis on the generative qualities of Würstchen based on the results of Sec. 4.2.

Reviewer comment (Experiments):
Since the Semantic Compressor is one of the main novelties, I wonder if you tested other feature extractors (e.g., could also use CLIP or Dino) and how that would affect training and quality, or by simply training an autoencoder with strong compression rate instead of using a pretrained feature extractor.

Author response:
> *"Since the Semantic Compressor is one of the main novelties, I wonder if you tested other feature extractors (e.g., could also use CLIP or Dino) and how that would affect training and quality. [...] Did you try other model architectures for the Stage C model (e.g., transformer based models) instead of only the ConvNext blocks?"*

Interestingly, we found that lower capacity semantic compressors did not impact the quality of the reconstruction negatively during our early exploration. For this reason, we switched from EfficientNetV2 L (120M parameters) to EfficientNetV2 S (22M parameters) to improve training efficiency in the early stages of development.
Due to the considerable cost of training Würstchen (approx.80.000-100.000$ per full training based on current AWS pricing), we were unable to study the effects on semantic compressor backbones and capacity thoroughly. For the same reason, we did not experiment with the Stage C architecture.
We hope to explore these avenues in future works.

Reviewer comment (Evaluation):
As the paper targets robust latent visual representations, there should be a detailed analysis (e.g., the quality of the latent features / the distribution of the compression space). This can make its claim more convincing.

Author response:
We also agree that an in-depth analysis of the compressed latent space would be insightful. Yet, as we enforce a normalized latent space (mean of 0, standard deviation of 1) of Stage C analyzing the distributions of the latent representation directly will not yield meaningful results. 
However, we expanded our analysis on the generative qualities further by analyzing categories and difficulty groups of partiprompts individually (Appendix H), to provide a more nuanced picture of the generative qualities of our model.
In a similar vein, we added FID-CLIP-score plots to characterize the model's behavior at different classifier-free-guidance scales and number of sampling steps (Appendix I). We also added collages for randomly sampled images for both ablation studies.

Reviewer comment (Evaluation):
For automatic evaluation metrics in Section 4.1, only FID and Inception score are evaluated, and there are no metrics evaluating how well the generated images are aligned with the input text instructions, such as CLIPScore.

Author response:
Following the recommendation of the reviewer, we provide MS-COCO CLIP-Score in Table 2 and on two additional datasets in Appendix B. We agree that an automated metric should be added showing the alignment of prompt and image directly. Technically, our third metric, the Pick-Score, considers prompts in the automated evaluation. However, Pick-Score also mixes general aesthetic qualities with text-alignment, making it insufficient for this purpose.

Reviewer comment (Evaluation):
The paper does not elaborate on the possible limitations or potential failure cases of the proposed method. Could the authors clarify this aspect?

Author response:
Following this comment, we performed an in-depth investigation of our human evaluation experiment on the partiprompts dataset to identify failure patterns, both regarding the image category and the challenge class (see Appendix H of the revised supplementary material). 
We indeed observe that Würstchen is underperforming compared to Stable Diffusion 2.1 in some aspects, specifically in tasks involving prompts that ask for writing/text, symbols and prompts requesting specific quantities of objects. In other words, tasks where the composition of patterns and objects in relation to each other is particularly important. 
We attribute this to the fact that Stage C operated on a very high compression, making fine-grained composition more challenging compared to Stable Diffusion 2.1. Stage B could correct this but is shown to primarily act as a super-resolution model, which we demonstrate in a small study in Appendix E.
We also briefly elaborate in Section 4.2 and Appendix I on the fact that the heavy filtering of the data (see Appendix G) results in a characteristic style for the generated images, which may not be desirable for a general image generator.