Matryoshka Diffusion Models

We thank the reviewer for acknowledging the motivation and the results of our paper. We also appreciate your constructive feedbacks. We answer your questions below, and we will provide more details with the updated draft.

The paper lacks an analysis on "comparison with literature"

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.

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.

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.

Some important implementation details are missing: learning rate, batch size, model configurations, etc.

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.

Missing quantitative evaluation on text-to-video generation compared with existing baselines.

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.

No video files included for illustrating text-to-video generation results.

We will include video results in the supplementary material.

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

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

computational budgets of all the experiments

We will update the draft with the exact computational budgets for the experiments. But roughly speaking, our ImageNet experiments are all using 8xA100 GPUs, and Text2Image experiments are using 32xA100 GPUs. All experiments run between 1-2 weeks, which correspond to <500K training iterations.

Are there any other differences compared to f-DM

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.

Figure 5 CFG weight

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.

For how long has the (video) model been trained and was there joint image/video training (or image pretraining) used?

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.

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?

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.

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

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

We thank the reviewer for acknowledging our contributions and results, and for the comprehensive review and constructive questions. We answer your questions below, and we will provide more details with the updated draft.

The method is not end-to-end

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.

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.

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.

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

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.

FID scores on ImageNet are ~3x times higher than the current SotA

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.

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

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).

We’d like to thank the reviewer for acknowledging our contributions and results. We answer your questions below, and we will provide more details with the updated draft.

The variables in the provided pseudocode for the NestedUNet architecture are not clear

We will update to pseudocode to clarify this.

In Table 1 we see that there is a noticeable gap between MDM and the baselines.

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.

Video results

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.

Sampling speed

MDM’s sampling speed is comparable to that of a standard UNet, and we will provide measurements in the updated draft.

We thank the reviewer for the constructive feedbacks. We answer some of your questions below, and we will provide more details with the updated draft.

It is anticipated that there will be an increase in computational load

Note that the computational cost of Nested Unet is almost identical to a standard Unet, as shown in Figure 3.

And the performance does not seem to surpass that of LDM, which employs classifier-free guidance

In Figure 4, we show the comparison of MDM to our own implementation of LDM. The LDM baseline uses the same Unet architecture as our 64x64 Unet, and we see that MDM outperforms LDM in this controlled setting. Note also that we didn’t apply CFG to either MDM or LDM, and we show in Figure 8 that CFG is applicable to MDM, similar to that shown in LDM.

Also, since multi-ratio and resolution training is already being conducted in stable-diffusion XL, the proposed method with multi-resolution training is not much novel.

We agree that multi resolution training in SDXL is related to MDM, but we respectfully point out that they work in fundamentally different ways. In SDXL, multi-resolution training can be considered as form a scheduled data augmentation, where the same model is applied to data of different resolutions, and during inference, the model is only capable of generating images of the highest resolution. MDM on the other hand, consumes images of multiple resolutions as part of the model architecture as well as the inference pipeline. It allows one to train images of multiple resolutions at the same time, and also generate images of multiple resolutions in a single inference run (see Figure 1 (a) top left corner as an example).

Also, it would be nice if the performance difference with usual UNet and NestedUNet is provided

We will provide additional experiments training a single NestedUnet and standard Unet in the updated draft — our observation is that they have near identical performance as two stand alone architectures.