Paper: Matryoshka Diffusion Models (ICLR-2024)

Reviewer comment (Evaluation):
The paper lacks an analysis on "comparison with literature". The paper simply states the result is comparable to other baselines, but the results show a clear gap (e.g., FID 3.60 (LDM) while 6.62 (MDM) on ImageNet 256x256). In this respect, the authors should provide an extensive analysis and reasons why the performance is worse than the baselines, not just saying the proposed method shows comparable performance.

Author response:
Our design for the ImageNet experiments are mostly focused on the control experiments, and we have directly borrowed the same architecture and hyper parameters from our text2image experiments. Also, due to computational demands, we did not train our models with a lot of iterations. This means that we have not optimized our FID results on ImageNet, and Table 1 is a rather a pessimistic estimate of our model’s performance. We are working on improving the results, and will update Table 1. 

W1 -> R3: 1. **ImageNet Results and comparison to literature**. We have updated Table 1 with much improved FID results on ImageNet. As stated before, we did not optimize our hyper parameters specifically for ImageNet, and also was not training the models long enough. During the rebuttal period, we have tried to increase the training batch size, and also fixed our previous way of performing CFG (we were not properly dropping labels during training in the previous version). After these simple changes, we see drastic improves in the FID numbers which now match the baseline results in Table 1. We are optimistic that there is yet more room to improve with some more hyper parameter tuning. We have also added comparison to the Imagen paper's FID-CLIP plot in Figure 8 (c). More discussions are also included in the paper wrt these updates, see Sec 4.2, 4.3.

Reviewer comment (Evaluation):
To verify the "faster convergence", I think x-axis in Figure 4 should be wall-clock time rather than training iterations. Otherwise, I think the authors should provide time/iteration for each baseline used for the evaluation.

Author response:
Thanks for the suggestion, we will provide the training cost details in the updated draft. Roughly speaking, the cost per iteration is near identical between single diffusion, MDM and CDM for the high resolution training, all are more costly that LDM.

W2 -> R3: 2. **Wallclock time for Figure 4**. We have included the time measurement for models in Figure 4, see Table 4. Essentially, MDM, simple diffusion and Cascaded diffusion have near identical training costs due to the similarity of Nested UNet, UNet, and the upsampling UNet. LDM is considerably faster than the other three due to its lower input dimensionality. Combining Figure 4 and Table 4, we can see that MDM provides a better computational efficiency than simple diffusion and Cascaded diffusion in our controlled setting.

Reviewer comment (Reproducibility):
Some important implementation details are missing: learning rate, batch size, model configurations, etc.

Author response:
We will update the draft for the implementation details. Our design for all the experiments have largely been using the same set of hyper parameters regarding lr, architecture etc, and we only tune the batch size to fit the GPU memory.

Reviewer comment (Experiments):
Missing quantitative evaluation on text-to-video generation compared with existing baselines.

Author response:
Due to resource constraints, the experimental results on the text-to-video tasks are preliminary, and our goal is to demonstrate qualitatively that the same idea behind MDM also applies to video generation. There is great interest to us to further improve the video generation’s quality and compare it with state of the art methods, which we leave as future work.

W4 -> R3: 3. **Video results**. We have included a discussion at the very last paragraph of the paper to make clear that the positioning of our paper is not to compare with state of the art video generation models. We leave such a full exploration as future work.

Reviewer comment (Presentation):
No video files included for illustrating text-to-video generation results.

Author response:
We will include video results in the supplementary material.

W5 -> R3: 4. Video samples are included in the supplementary material, please download the zip ball to view it.

Reviewer comment (Presentation):
Why some points in Figure 4 (e.g., after 200K of Latent DM in Figure 4(a)) are missing?

Author response:
The LDM experiments were not finished at the point of submission, but we will upload a complete comparison for LDM in Figure 4.

W6 -> R3: 5. We have completed the baseline plots in Figure 4.

Reviewer comment (Experiments):
The method is not end-to-end (or at least does not perform well when trained in the end-to-end manner), despite what the paper claims. If I am not mistaken, at the end it still trains stage-by-stage similarly to CDMs — and without such stage-by-stage training it produces considerably worse results.

Author response:
First, we apologize for misunderstandings on the end-to-end claim in our presentation. We agree that our progressive training mode does not exactly fit into the common perception of end-to-end training, and we are happy to remove such claims in the paper. However, we’d like to highlight that even without progressive training, MDM still offers significant advantage over simple diffusion, as well as our own implementations of CDMs and LDM (see Figure 4). In addition, the inference mode of MDM follows a single inference pass in the same latent space, with or without progressive training, which also differentiates itself from multi-stage approaches like CDM, LDM and fDM.

Reviewer comment (Evaluation):
The paper does not report training costs rigorously for all the experiments, and it's impossible to compare between methods without knowing their training cost.

Author response:
We agree, thanks for pointing it out. Roughly speaking, the cost per iteration is near identical between single diffusion, MDM and CDM for the high resolution training, all are more costly that LDM. We will provide the actual training cost in the updated draft.

Reviewer comment (Experiments):
The comparison to CDMs does not seem fair, since the paper compares to under-trained CDMs. If one trains CDMs within a limited computational budget, then more focus should be put on the base stage, since the final stages converge much faster.

Author response:
We agree that our own CDM baseline is not optimally tuned, however we do believe that it’s still a fair comparison and reflects the advantage of MDM over CDM. As correctly stated by the reviewer, CDM is sensitive to quality of the low res model, which requires one to both train the low resolution model to near optimal as well as carefully tune the noise augmentation level for the low resolution inputs. The comparison in Figure 4 shows that, starting from the same low resolution model, MDM clearly shows an advantage which demonstrates the robustness of our progressive training scheme over cascaded diffusion.

Reviewer comment (Evaluation):
FID scores on ImageNet are ~3x times higher than the current SotA (e.g., MDT).

Author response:
Our design for the ImageNet experiments are mostly focused on the control experiments, and we have directly borrowed the same architecture and hyper parameters from our text2image experiments. Also, due to computational demands, we did not train our models with a lot of iterations. This means that we have not optimized our FID results on ImageNet, and Table 1 is a rather a pessimistic estimate of our model’s performance. We are working on improving the results, and will update Table 1.

Reviewer comment (Evaluation):
The paper makes a claim about good results on a small text-to-image dataset (CC12M), but does not compare to existing large-scale text-to-image generators. This makes it impossible to evaluate this claim — e.g., Figure 8c should contain the results of existing text-to-image generators to make the existing model comparable. Otherwise, such a claim is not grounded. Judging by the maximum CLIPScore on a CLIP/FID trade-off chart is wrong since in the Imagen's paper, one can notice that even their bad models can attain very high CLIP scores under a strong enough guidance.

Author response:
We are working on updating the draft to include a comparison with other text-to-image baselines. However, we respectfully disagree that high CLIP score is not indicative of a model’s perceptual quality — in our experience, CLIP score correlates much better with the sample’s perceptual quality than FID. The reviewer refers to the Imagen paper’s results and indicate that bad models are capable of achieving high CLIP scores, however we could not find such evidence (eg, in Figure 4(a) ). Can you elaborate which Figure/Table supports this claim? (the closest thing we could find is Figure A.11, however these plots are from the same model size and the FID differences are also small).

Reviewer comment (Novelty):
Are there any other differences compared to f-DM [1] apart from the progressive training idea and larger-scale experiments? It seems that f-DM uses the same NestedUNet idea, but the f-DM authors just do not call it a "NestedUNet". Do you use the same noise schedule as f-DM?

Author response:
MDM is related to f-DM in terms of parameter sharing across different resolutions, however with one fundamental difference. In f-DM, different resolutions are chained up sequentially, whereas the network handles one resolution at a time. This means that for f-DM the network architecture is identical to that of a standard diffusion model when denoising the highest resolution. In MDM, all resolutions are processed in parallel by the NestedUNet, both during training and inference. This sequential vs parallel difference is subtle, but makes a big difference wrt their empirical performance.

Reviewer comment (Evaluation):
To be honest, I do not quite understand why the FID on ImageNet is so high. Samples in Figure 5 looks very good to me (given that they are random samples). What CFG weight was used to generate them?

Author response:
Figure 5 is generated with  a CFG weight of 2. We have inspected results from different weights within the range [1.1, 5], they mostly visually appealing which indicates that CFG is applicable to MDM similarly to a standard diffusion model.

Reviewer comment (Experiments):
Please, include the comparison with existing text-to-image generators. Judging by the maximum attainable CLIP score is misleading.

Author response:
3. **CLIP vs FID**. Thank you for acknowledging our points and we are very glad that we could reach an agreement on this subtle matter :) We also agree with the reviewers suggestion and have updated Figure 8 to include the results from the Imagen paper as well. We have also edited the discussions in Sec 4.3 to reflect this update.

Reviewer comment (Reproducibility):
I find the results on video generation to be quite good. For how long has the model been trained and was there joint image/video training (or image pretraining) used?

Author response:
The video model was trained following the same protocol as the image domain, whereas we first train an image generator for ~400K iterations. Then the video model is trained with the (temporal) NestedUnet, whereas there is a video loss and image loss (similar to image losses in different resolutions). The training time of the video model is less than two weeks on 32xA100 GPUs.

Reviewer comment (Theory):
Why do you think your model does not suffer from the "train/test gap" problem of CDMs and its late stages can generate meaningful images even when the low-resolution stage has not been trained yet?

Author response:
The key difference between progressively training MDM and CDM is that for MDM, all parameters are jointly optimized, whereas CDM freezes the low resolution generator. In other words, the role of progressive training in MDM can be viewed as a better initialization of the weights. During inference time, MDM generates low and high resolution images in parallel, where we rely on a good inference schedule to deal with the training/test gap issue.

Reviewer comment (Experiments):
I have a suspicion that the video generator can struggle in generating videos with moving scenes because the base low-resolution generator produces just a single frame. Could you please provide the video results for moving scenes? And in general include some mp4/gif videos in the submission (I have not found any in the supplementary).

Author response:
The reviewers hypothesis is correct and we will provide additional examples in the supplementary material.

Reviewer comment (Writing):
The variables in the provided pseudocode for the NestedUNet architecture are not clear. A quick description to clarify the inputs to the function would be beneficial.

Author response:
We will update to pseudocode to clarify this.

Reviewer comment (Evaluation):
In Table 1 we see that there is a noticeable gap between MDM and the baselines. It would be helpful if the authors provided insight as to why they think this may be the case.

Author response:
We have updated Table 1 with much improved FID results on ImageNet. As stated before, we did not optimize our hyper parameters specifically for ImageNet, and also was not training the models long enough. We have tried to increase the training batch size, and also fixed our previous way of performing CFG (we were not properly dropping labels in the previous version). After these simple changes, we see drastic improves in the FID numbers which now match the baseline results in Table 1. We are optimistic that there are yet more room to improve with some more hyper parameter tuning.

Reviewer comment (Evaluation):
The authors highlight video generation as a contribution of MDM, but there is no discussion of the experiments or results (aside from implementation details and a few subsampled frames). It is difficult to get a sense of MDM's performance for this task.

Author response:
We have provided video samples in the supplementary material; we have also added discussions (the last paragraph of the paper) to make clear that we are not claiming that our video results are comparable to the SOTA, but rather to leave such an exploration as future work.

Reviewer comment (Presentation):
While there is some differentiation already denoted through the colors, it would be helpful to explicitly label the novel pathways introduced by the NestedUNet architecture in Figure 3 to clearly distinguish from skip connections from the original UNet.

Author response:
We updated the pseudo code, as well as the illustration Figure 3 for better clarity. We also provided more details on the inference procedure in Figure 11 and Sec C in Appendix.

Reviewer comment (Evaluation):
It is helpful to understand how the proposed multi-resolution prediction affects sampling speed.

Author response:
We have provided a benchmarking result of Nested UNet vs standard UNet wrt their inference speed in Table 2. This shows that Nested UNet (and hence MDM) has the same inference speed as a standard diffusion model with UNet.