Matryoshka Diffusion Models
Authors: Jiatao Gu
Shuangfei Zhai
Yizhe Zhang
Joshua M. Susskind
Navdeep Jaitly
Venue: ICLR 2024 poster (2024)
Keywords: diffusion models
text-to-image generation
nesting architecture
progressive training
Primary area: generative models
TL;DR: We propose a single-stage diffusion model for high resolution inputs

Abstract: Diffusion models are the de-facto approach for generating high-quality images and videos, but learning high-dimensional models remains a formidable task due to computational and optimization challenges. Existing methods often resort to training cascaded models in pixel space, or using a downsampled latent space of a separately trained auto-encoder. In this paper, we introduce Matryoshka Diffusion (MDM), an end-to-end framework for high-resolution image and video synthesis. We propose a diffusion process that denoises inputs at multiple resolutions jointly and uses a NestedUNet architecture where features and parameters for small-scale inputs are nested within those of large scales. In addition, MDM enables a progressive training schedule from lower to higher resolutions, which leads to significant improvements in optimization for high-resolution generation. We demonstrate the effectiveness of our approach on various benchmarks, including class-conditioned image generation, high-resolution text-to-image, and text-to-video applications. Remarkably, we can train a single pixel-space model at resolutions of up to 1024x1024 pixels, demonstrating strong zero-shot generalization using the CC12M dataset, which contains only 12 million images. Code and pre-trained checkpoints are released at https://github.com/apple/ml-mdm.

Published as a conference paper at ICLR 2024
MATRYOSHKA DIFFUSION MODELS
Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Josh Susskind & Navdeep Jaitly Apple {jgu32, szhai, yizzhang, jauigsssdsiiiiiiiiiillllllleeee.
Figure 1: (<+) Images generated by MDM 642, 1282, 256 5122 and 1024 resolutions using the prompt "a Stormtrooper Matryoshka doll, super details, extreme realistic, 8k": (+- 1 and 16 frames of 642 video generated by method using the prompt "pouring milk into black coffee"; All other samples are at 10242 given various prompts. Images were resized for ease of visualization.
ABSTRACT
Diffusion models are the de-facto approach for generating high-quality images and videos but learning high-dimensional models remains formidable task due to computational and optimization challenges. Existing methods often resort to training cascaded models in pixel space, or using downsampled latent space of separately trained auto-encoder. In this paper, we introduce Matryoshka Diffusion (MDM), a novel framework for high-resolution image and video synthe- sis. We propose diffusion process that denoises inputs multiple resolutions jointly and uses a NestedUNet architecture where features and parameters for small scale inputs are nested within those of the large scales. In addition, MDM enables progressive training schedule from lower to higher resolutions which leads to significant improvements in optimization for high-resolution generation. We demonstrate the effectiveness of our approach on various benchmarks, in- cluding class-conditioned image generation, high-resolution text-to-image, and text-to-video applications. Remarkably, we can train single pixel-space model at resolutions of up to 1024 x 1024 pixels, demonstrating strong zero shot generaliza- tion using the CC120 2M dataset, which contains only 12 million images. Code and pre-trained checkpoints are released at https: //g thub. com/apple, hmmpee/mm e/ml mdm.
1 INTRODUCTION
Diffusion models et al., 2015; Ho et al., 2020; Nichol & Dhariwal, 2021; Song et al., 2020) have become increasingly popular tools for generative applications, such image (Dhariwal

---

Published as a conference paper at ICLR 2024
& Nichol, 2021; Rombach et al., 2022; Ramesh al., 2022; Saharia et al., 2022), video (Ho et al., 2022c;a), 3D (Poole et al., 2022; Gu et al., 2023; Liu et al., 2023b; Chen et al., 2023), audio (Liu et al., 2023a). and text (Li et al., 2022; Zhang al., 2023) generation. However scaling them to high- resolution still presents significant challenges as the model must re-encode the entire high-resolution input for each step (Kadkhodaie et al., 2022). Tackling these challenges necessitates the use of deep architectures with attention blocks which makes optimization harder and uses more resources.
Recent works (Jabri et al., 2022; Hoogeboom al., 2023) have focused on efficient network archi- tectures for high-resolution images. However, none of the existing methods have shown competitive results beyond 512 x 512. and their quality still falls behind the main-stream cascaded/latent based methods. For example, DALL-E (Ramesh et al.. 2022), IMAGEN (Saharia et al., 2022) and eDiff- I (Balaji et al., 2022) save computation by learning a low-resolution model together with multiple super-resolution diffusion models, where each component is trained separately. On the other hand, latent diffusion methods (LDMs) (Rombach et al., 2022: Peebles Xie, 2022; Xue al., 2023) only learn low-resolution diffusion models, while they rely on a separately trained high-resolution autoencoder (Oord et al., 2017; Esser et al., 2021). In both cases, the multi-stage pipeline complicates training & inference, often requiring careful tuning of hyperparameters.
In this paper. we present Matryoshka Diffusion Models (MDM), a novel family of diffusion models for high-resolution synthesis. Our main insight is to include the low-resolution diffusion process as part of the high-resolution generation, taking similar inspiration from multi-scale learning in GANs (Karras et al., 2017; Chan al., 2021; Kang et al., 2023). We accomplish this by performing a joint diffusion process over multiple resolution using Nested UNet architecture ( (see Fig. 2 and Fig.3) 3). Our key finding is that MDM, together with the Nested UNets architecture, enables 1) multi- resolution loss that greatly improves the speed of convergence of high-resolution input denoising and 2) an efficient progressive training schedule, that starts by training low-resolution diffusion model and gradually adds high-resolution inputs and outputs following schedule. Empirically, we found that the multi-resolution loss together with progressive training allows one to find an excellent balance between the training cost and the model's quality.
We evaluate MDM on class conditional image generation, and text conditioned image and video generation. MDM allows to train high-resolution models without resorting to cascaded or latent diffusion. Ablation studies show that both multi-resolution loss and progressive training greatly boost training efficiency and quality. In addition, MDM yield high performance exx-to--maae generative models with up to 1024 resolution, trained on the reasonably small CC12M dataset. Lastly, MDM generalize gracefully video generation, suggesting generality of our approach.
2 DIFFUSION MODELS
Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) are latent variable models given a pre-defined posterior distribution (named the forward diffusion process), and trained with de- noising objective. More specifically, given data point e R% and fixed signal-noise schedule fas, oih= we define a sequence of latent variables {za (2a)=,,,..,. that satisfies:
q( x) = N(2 a, I) and g( = (Z) 1),
(1)
where 2o œ. 0// = az/a., oj.- = of - -aj. < I. By default, the signal-to-noise ratio (SNR, af /f decreases monotonically with t. The model then learns to reverse the process with a backward model p0(=z1122) which can be re-written as denoising objective:
Co = F lx) [w> [ B) =012]
where (z),).)) 1) is a neural network (often a variant of a UNet model (Ronneberger et al., 2015)) that maps a noisy input 2; to its clean version æ. conditioned on the time step t; w, R is loss weighting factor determined by heuristics. In practice, one can reparameterize 2g with noise- or v-prediction (Salimans & Ho, 2022) for improved performance. Unlike other generative models like GANs (Goodfellow et al., 2014), diffusion models require repeatedly applying deep neural network æg the ambient space as enough computation with global interaction is critical for denoising (Kadkhodaie al., 2022). This makes it challenging to design efficient diffusion models directly for high-resolution generation, especially for complex tasks like text-to-image synthesis. As common solutions, existing methods have focused on learning hierarchical generation:
2

---

Published as a conference paper at ICLR 2024
642
#7
2
256
Modil
(r2.
x
1024'
θ
Figure 2: An illustration Matryoshka Diffusion. #t a/ and a are noisy images three different resolutions, which are fod into the denoising network together, and predict targets independently.
Cascaded diffusion (Ho et al., 2022b; Ramesh et al., 2022; Saharia et al., 2022; Ho et al., 2022a; Pernias et al., 2023) utilize a cascaded approach where a first diffusion model is used to generate data at lower resolution, and then second diffusion model is used generate super- resolution version of the initial generation, taking the first stage generation as conditioning. Cascaded models can be chained multiple times until they reach the final resolution. Ho al. (2022a); Singer et al. (2022) uses a similar approach for video synthesis as well models are cascaded from low spatio-temporal resolution high spatio-temporal resolution. However, since each model trained separately, the generation quality can be bottlenecked by the exposure bias (Bengio et al.. 2015) from imperfect predictions and several models need to be trained corresponding different resolutions.
Latent diffuusin (LDM, Rombach al., 2022) and its follow-ups (Peebles & Xie, 2022; Xue et al., 2023; Podell al., 2023), the other hand, handle high-resolution image generation by performing diffusion in the lower resolution latent space of a pre-trained auto-encoder, which is typically trained with adversarial objectives (Esser et al., 2021). This not only increases the complexity of learning, but bounds the generation quality due to the lossy compression process.
End-to-end models Recently, several approaches have been proposed (Hoogeboom et al., 2023; Jabri et al., 2022; Chen, 2023) to train end-to-end models directly on high-resolution space. Without relying separate models, these methods focus efficient network design well as shifted noise schedule to adapt high-resolution spaces. Nevertheless, without fully considering the innate structure of hierarchical generation, their results lag behind cascaded and latent models.
MATRYOSHKA DIFFUSION MODELS
In this section, we present Matryoshka Diffusion Models (MDM), a new class of diffusion models that is trained in high-resolution space, while exploiting the hierarchical structure of data formation. MDM first generalizes standard diffusion models in the extended space ($ I), for which specialized nested architectures (S 3.2) and training procedures (Appendix B) are proposed.
3.1 DIFFUSION MODELS IN EXTENDED SPACE
Unlike cascaded or latent methods, MDM learns single diffusion process with hierarchical structure by introducing a multi-resolution diffusion process in an extended space. An illustration is shown in Fig. 2. Given a data point E R.,, we define time-dependent latent 2 ..... 2 R Nn...N Similar to Eq. (1), for each ž,,
R:
q(> =
o
(2)
where D N + N, is a deterministic "down-sample" operator depending on the data. Here, D' is coarse / lossy-compressee version of For instance, (*(.) can be avgpool(.) for generating low-resolution images.
By default, assume compression in progressive manner such that N1 N2 NR N and (x)) = æ. Also, {a,, oo] are the resolution-specife noise schedule. In this paper, we follow Gu et al. (2022) and shift the noise schedule based the input resolutions. MDM then learns
3

---

Published as a conference paper at ICLR 2024
condicioing flent, voc!
Supervaion
Figure 3: An illustration of the NestedUNet architecture used in Matryoshka Diffusion. We follow the design of Podell et al. (2023) by allocating more computation in the low resolution feature maps (by using more attention layers for example), where in the figure we the width of a block to denote the parameter counts. Here the black arrows indicate connections inherited from UNet, and red arrows indicate additional connections introduced by Nested UNet.
the backward process po(zt- (z,) with R neural denoisers (z)) Each variable z depends all resolutions {=] ..!" time step During inference, MDM generates resolutions in parallel. There is no dependency between z7.
Modeling diffusion in the extended space has clear merits: (1) since what care during inference is the full-resolution output z;" all other intermediate resolutions are treated as additional hidden variables 2;, enriching the complexity of the modeled distribution;(2) the mustiiresolution depen- dency opens up opportunities to share weights and computations across 2f, enabling to re-allocate computation more efficient manner for both training and inference efficiency.
3.2 NESTEDUNET ARCHITECTURE
Similar to typical diffusion models, we implement MDM the flavor of UNet (Ronneherger al., 2015; Nichol & Dhariwal, 2021): skip-connections used in parallel with computation block to preserve fine-grained input information, where the block consists multi-level convolution and self-attention layers. In MDM, under the progressive compression assumption, it is natural that the computation for x is also beneficial for 2; This leads us to propose NestedUNet, an architecture that groups the latents of all resolutions {z() in one denoising function nested structure, where low resolution latents will be fed progressively along with standard down-sampling. Such multi-scale computation sharing greatly eases the learning for high-resolution generation. A pseudo code for NestedUNet compared with standard UNet is present as follows.
def NestedUNet (z:tist (Tensor), h:Tensor=None, o: [Tensor)=
# 2: list of inputs with increasing resolutions
# h: output hidden states from previous resolution
# f_merge, f f_sip, f_up, f_down: neural layers
21-1] if is None else f_merge (zl-1),
if len (2) 1: move to next resolution
x f f_skip (x, f_up (NestedUNet (z [:-1], f_down (x), ))))
else:
inner UNet at lowest resolution
f_skip (x, f_up (f_midddd (f_doww (x))
append return x
return results of all resolutions
Aside from the simplicity aspect relative to other hierarchcal approaches, NestedUNet also allows to allocate the computation the most efficient manner. As shown Fig. 3, our early exploration found that MDM achieved much better scalibility when allocating most of the parameters & computation the lowest resolution. Similar findings have also been shown in Hoogeboom al. (2023).
3.3 LEARNING
We train MDM using the normal denoising objective jointly multiple resolutions, as follows:
R
La = E2 E. Esing(zila)
W [100( - D -D'(w))1))
(3)
4

---

Progressive Training While MDM can be trained end-to-end directly following Eq. (3) which has already shown better convergence than naive baselines, we found simple progressive training technique, similarly proposed in GAN literature (Karras et al., 2017; Gu et al., 2021), greatly speeds up the training of high-resolution models w.r.t. wall clock time. More precisely, we divide up the training into R phases, where we progressively add higher resolution into the training objective in Eq. (3), This is equivalent to learning a sequence of MDMs on [z],... =]] until r reaching the final resolution. Thanks to the proposed architecture, we can achieve the above trivially if progressive Published as a conference paper at ICLR 2024
where w/ is the resolution-specifee weighting, and by default we set /wf NR/N,.
growing the networks (Karras et al., 2017). This training scheme avoids the costly high-resolution training from the beginning, and speeds up the overall convergence.
4 EXPERIMENTS
MDM is a versatile technique applicableto any problem where input dimensionality can be progres- sively compressed. We consider two applications beyond class-conditional image generation that demonstrate the effectiveness of our approach - text-to-image and text-to-video generation.
4.1 EXPERIMENTAL SETTINGS
Datasets In this paper, we only focus on datasets that are publicly available and easily reproducible. For image generation, we performed dlasscoodittonee generation on ImageNet (Deng et al., 2009) at 256 x 256, and performed general purpose text-to- to-tte--mmae generation using Conceptual 12M (CC12M, Changpinyo et al., 2021) at both 256 x 256 and 1024 x 1024 resolutions. As additional evidence of generality, we show results on text-to-video generation using WebVid-1OM (Bain et al., 2021) at 16 x 256 x 256. We list the dataset and preprocessing details in Appendix F.
The choice of relying extensively on CC12M for text-to-image generative models the paper significant departure from prior works (Saharia et al., 2022; Ramesh et al., 2022) that rely on exceedingly large and sometimes inaccessible datasets, and so we address this choice here. We find that CC12M is sufficient for building high-quality text-to-image models with strong zero-shot capabilities in a relatively short training time (see details in Appendix D.2). This allows for a much more consistent comparison of methods for the community because the dataset is freely available and training time is feasible. We submit here, that CC12M is much more amenable as a common training and evaluation baseline for the community working on this problem.
Evaluation In line with prior works, evaluate image generation models using Fréchet Inception Distance (FID, Heusel et al., 2017) (ImageNet, CC12M) and CLIP scores (Radford et al., 2021) (CC12M). To examine their zero-shot capabilities, also report the FID/CLIP scores using COCO (Lin et al., 2014) validation set togenerate images with the CC12M trained models. We also provide additional qualitative samples for image and video synthesis in supplementary materials.
Implementation details We implement MDMs based on the proposed NestedUNet architecture, with the innermost UNet resolution set to 64 x 64. Similar to Podell et al. (2023), we shift the bulk of self-attention layers the lower-level (16 X 16) features, resulting in total 450M parameters for the inner UNet. As described in $ 3.2, the high-resolution part of the model can be easily attached on top of previous level of the NestedUNet, with minimal increase in the parameter count. For text-to-image and text-to-video models, we use the frozen FLAN-T5 XL (Chung et al., 2022) as our text encoder due to its moderate size and performance for language encoding. Additionally, we apply two learnable self-attention layers over the text representation to enhance text-image alignment.
For image generation tasks, we experiment with MDMs of (642, 256²). {642, 128?, 256²) for 256 x 256, and {642, 2562, 10243), [642, 128 256 5122, 10242] 3 for 1024 x 1024, respectively. For video generation, MDM is nested by the same image 64 x 64 UNet with additional attention layers for learning temporal dynamics. The overall resolution is (642, 16 x 642, 16 x 2562). We use bi-linear interpolation for spatial (.)., and first-frame indexing for temporal Dr..). Unless specified, we apply progressive and mixed-resolution training for all MDMs. We use 8 A100 GPUs for ImageNet, and 32 110 GPUs for CC12M and WebVid-10M, respectively. See Appendices A and B for more implementation
and training details.
5

---

Published as a conference paper at ICLR 2024
MDM
MDM
MDM tho PT:
MDM
Stmple
3
Bimple
Concarteed DM
Canpaded DM
FD 10k
Latent DM
RD 1bx
9
Latert DM
CLIF SX
8
MDM (ne PT)
88
Simp DM
Latert DM
150
202
253
332
300
400
450
100
300
300
350
150
150
200
250
300
350
Faral oma
Traiiing Baralora 0C)
Training Earaitonn (K.
(a) FID (L) of ImageNet 256 x 256. (b) FID (1) on CC12M 256 x 256. (e) CLIP (1) on CC12M 256 x 256.
Figure 4: Comparison against baselines during training. FID (4) (a, b) and CLIP(1) (c) scores of samples generated without CFG during training of different class conditional models of ImageNet 256 x 256 (a) and CC12M 256 x 256 (b, c). As can be seen, MDM models that were first trained lower resolution (200K steps for ImageNet, and 390K for CC12M here) converge much faster.
Baseline models Aside from the comparisons with existing state-of-the-art approaches, we also report detailed analysis on MDMs against three baseline models under controlled setup:
1. Simple DM: A standard UNet architecture directly applied to high resolution inputs; We also consider the Nested UNet architecture. but ignoring the low resolution losses; Both cases are essentially identical to recent end-to-end diffusion models like Hoogeboom et al. (2023).
Cascaded DM: we follow the implementation details of Saharia et al. (2022) and train CDM that is directly comparable with MDM where the upsampler has an identical configuration to our NestedUNet. We also apply noise augmentation to the low resolution conditioning image, and sweep over the optimal noise level during inference.
Latent DM: we utilize the latent codes derived from the auto-encoders from Rombach et al. (nd22).and subsequently train diffusion models that match the dimensions of the MDM UNet.
Comparison with literature on Im- ageNet (FID-50K), and COCO (FID-30K). indicates samples generated with CFG. Note existing text-to-image models are mostly trained on much bigger datasets than CC12M.
Table I:
4.2
MAIN RESULTS
Comparison with baseline approaches Our compar- isons to baselines are shown in Fig. 4. On ImageNet 256 256, we select a standard UNet our simple DM baseline. For the Cascaded DM baseline, we pretrain 64x64 diffusion model for 200K iterations, and apply an upsampler UNet also in the same size. We apply standard noise augmentation and sweep for the opti- mal noise level during inference time (which we have found to be critical). For LDM experiments, we use pretrained autoencoders from Rombach et al. (2022)
Models
FID
ImageNet 256 256
ADM (Nichol Dhariwal, 2021)
10.94
CDM (Ho et al., 2022b)
4.88
LDM-4 (Rombach et al., 2022)
10.56
L.M-4 (Rombach al., 2022)
3.60
Ours (cfg=1) Ours (cfg=1.5)*
8.18
which downsamples the imput resolution and we use the same architecture for these experiments our 64x64 low resolution models. For MDM variants, we use a NestedUNet of the same size as the baseline UNet. We experiment with two variants, one trained directly with the multi resolution loss Eq. (3) (denoted no PT), and another one resuming from the 64x64 diffusion model (ie, progressive training). CC12M 256x256 follows similar setting, except that we use single loss Neste- dUNet as our simple DM architecture. We monitor the FID curve on ImageNet, and the FID and CLIP curves CC12M.
3.51
MS-COCO 256 256
LDM-8 (Rombach et al., 2022)
23.31
LDM-8* (Rombach al., 2022)
12.63
Dalle-2* 22 (Ramesh et al., 2022)
10.39
IMAGEN (Saharia et al., 2021)
7.27
Ours (cfg=l)
18.35
Ours (cfg=1.35)*
13.43
Comparing simple DM to MDM, we see that MDM clearly has faster convergence, and reaches better performance in the end. This suggests that the multi resolution diffusion process together with the multi resolution loss effectively improves the models convergence, with negligible added complexities. When following the progressive training schedule, see that MDM's performance and convergence speed further improves. As a direct comparison, we see that the Cascaded DM baseline significantly underperforms MDM, while both starting from the same 64x64 model. Note

---

Published as a conference paper at ICLR 2024
Figure 5: Random samples from our class-conditions MDM trained on ImageNet 256 256.
that this is remarkable because Cascaded DM has more combined parameters than MDM (because MDM has extensive parameter sharing across resolutions), and uses twice as many inference steps. We hypothesize that the inferior performance of Cascaded DM is largely due to the fact that our 64x64 is not aggressively trained, which causes a large gap between training and inference wrt the conditioning inputs. Lastly, compared to LDM, MDM also shows better performance. Although this is a less direct control as LDM is indeed more efficient due to its small input size, but MDM features a simpler training and inference pipeline.
Comparison with literature In Table MDM is compared to existing approaches in literature, where we report FID-50K for ImageNet 256x256 and zero shot FID-30K on MSCOCO. On Ima- geNet, for which our architecture and hyperparameters competitive FID of 3.51 with CFG. Our FID results comparable to the literature, although MDM is trained on significantly less data than the baselines like Imagen and Dalle-2.
optimized, MDM is able achieve
Qualitative Results We show random samples from the trained MDMs on for image generation (ImageNet 256 x 256, Fig. 5), text-to-image (CC12M, 1024x 1024 Fig. 6) and text-to-video (WebVid- 10M. Fig. 7). Despite training on relatively small datasets, MDMs show strong zero-shot capabilities of generating high-resolution images and videos. Note that we use the same training pipelines for all three tasks, indicating its versatile abilities of handling various data types.
4.3 ABLATION STUDIES
Effects of progressive training We experiment with the progressive training schedule, where we vary the number of iterations that the low-resolution model is trained before continuing on the target resolution (Fig. 8a). We see that more low resolution training clearly benefits that of the high-resolution FID curves. Note that training on low resolution inputs much more efficient w.r.t. both memory and time complexity, progressive training provides a straightforward option for finding the best computational trade-offs during training.
Effects of nested levels Next, we compare the performance of using different number of nested resolutions with experiments on CC12M. The result is shown in Fig. 8b. We see that increasing from two resolution levels to three consistently improves the model's 's convergence. It's also worth noting that increasing the number of nesting levels brings only negligible costs.
CLIP-FID trade-off Lastly, we show in Fig. 8c the pereto curve of CLIP-FID on the zero-shot evaluation of COCO, achieved by varying the classifier free guidance (CFG) weight. MDM is similarly amendable to CFG as other diffusion model variants. As comparison, we overlap the same plot reported by Imagen (Figure A.11). We see that Imagen in general demonstrates smaller FID, which attribute it to higher diversity MDM demonstrates strong CLIP score, whereas we have found in practice that such high CLIP scores correlate very well with the visual quality of the generated images.
result of training on large dataset. However,
7

---

Published as a conference paper at ICLR 2024
Atuiliinnal Dinma ou lengly Cladi
ICLR R
Aruh nd paddredd wodig prevgarr
Cieeais ploto Wall book whi biking weg.ooplllitee.oommooooooo diseug ug.
MDM
Peckeelie Bandl Creor
publo of #1 #: maurgis
WrRch wornya spbmne toruniied grflt, denncp (cu pak
yyuung bidge calisaiie yalln tox Rocand pa
pondooottinnn ow colloiiifee 1900 Gieears, weldo
Figure 6: Samples from the model trained on CC12M at 1024 with progressive training.
5
RELATED WORK
In addition to diffusion methods covered in $ 2, multiscale models have been widely used image generation and representation learning (Kusupati al., 2022). A well-known Generative Adversarial Network (GAN) is the LAPGAN model (Denton et al., 2015) which generates lower-resolution images that subsequently fed into higher-resolution models. Pyramidal Diffusion (Ryu & Ye, 2022), applies a similar strategy with denoising diffusion models. Autoregressive models have also been applied for generation from early works for images (Van Den Oord al., 2016: Oord al., 2016) and videos (Kalchbrenner et al., 2017; Weissenborn et al., 2020), to more recent text-to-iman to-image models (Gafni et al., 2022; Yu et al., 2022) and text to video models (Wu et al., 2021; Singer al., 2022). While earlier works often operate in pixel space, recent works, such as Parti (Yu et al., 2022) and MakeAScene (Gafni et al., 2022) use autoencoders to preprocess images into discrete latent features which can be modeled autoregressively using large sequence-to-sequence models based on transformers. f-DM (Gu et al., 2022) proposed a generalized framework enabling progressive signal transformation across multiple scales, and derived corresponding de-noising scheduler to transit from multiple resolution stages. This scheduler is employed in our work. Similarly, IHDM (Rissanen al., 2023) does coarse-to-fine generation implicitly increase the resolution.
8

---

Published as a conference paper at ICLR 2024
The Bund, Shanghai, with the ship moving on the river, oil painting
A happy giant panda playing guitar in the forest,day,
shark swimming in the ocean
Figure 7: Samples from model trained on WebVid- at 16 256² with progressive training. Videos are subsampled for ease of visualiation.
200
20
NDM
WEM UEM maning BOOK
Imagan
MEM 2008
Bimple DV
CUP
MDM (5 200
MCB (5--120-250)
mirple DM
Traiiin solonnn (K)
Training Reratioa (0)
CLP soore
(a) FTD (1) on ImageNet 256 2h). CL.IP (t) on CC12M 256 256. (c) Trade-off on COCO 256 256.
Figure 8: (a) Increasing the number of steps of low resolution training in the progressive training improves
results. (b) Larger number of nesting levels CLIP produces improvements in speed of convergence and
final score (c) FID vs CLIP trade-off seen by varying the weight CFG (using evaluation on COCO)
6 DISCUSSIONS AND FUTURE DIRECTIONS
In this paper showed that sharing representations across different resolutions can lead to faster training with high quality results, when lower resolutions are trained first. We believe this is because the model is able to exploit the correlations across different resolutions more effectively, both spatially and temporally. While we explored only a small set of architectures here, we expect more improvements can be achieved from a more detailed exploration of weight sharing architectures. and new ways of distributing parameters across different resolutions in the current architecture. Another unique aspect of our work is the use of an augmented space, where denoising is performed over multiple resolutions jointly. In this formulation resolution over time and space are treated in the same way, with the differences in correlation structure in time and space being learned by different parameters of the weight sharing model. A more general way of conceptualizing the joint optimization over multiple resolutions is to decouple the losses at different resolutions, by weighting them differently. It is conceivable that a smooth transition can be achieved from training lower to higher resolution. We also note that while we have compared our approach to LDM in the paper, these methods complementary. It is possible to build MDM on top of autoencoder codes. While we are not making the claim that the MDM based models are reaching the SOTA, we leave the evaluation of MDM large scale dataset and model sizes future work.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENT
We thank Miguel Angel Bautista, Jason Ramapuram, Alaaeldin El-Nouby, Laurent Dinh, Ruixiang Zhang, Yuyang Wang for their critical suggestions and valuable feedback this project. We thank Ronan Collobert, David Grangier and Awni Hanun for their invaluable support and contributions to dataset pipeline.
REFERENCES
Max Bain, Arsha Nagrani, Gül Varol, and Andrew Zisserman. Frozen in time: A joint video and image encoder for end-to-end retrieval. IEEE International Conference on Computer Vision, 2021.
Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Karsten Kreis, Mika Aittala, Timo Aila, Samuli Laine, Bryan Catanzaro, et al. ediffi: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01322 2022.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015.
Eric R Chan, Connor Lin, Matthew Chan, Koki Nagano, Boxiao Pan, Shalini De Mello, Orazio Gallo, Leonidas Guibas, Jonathan Tremblay, Sameh Khamis, et al. Efficient geometry-awar 3d generative adversarial networks. arXiv preprint 1rX1v:2112.07945, 2021.
Soravit Changpinyo, Piyush Sharma, Nan Ding, and Radu Soricut. Conceptual 12M: Pushing web-scale image-text pre-training to recognize long-tail visual concepts. In CVPR, 2021.
Hansheng Chen, Jiatao Gu, Anpei Chen, Wei Tian, Zhuowen Tu, Lingjie Liu, and Hao Su. Single- stage diffusion nerf: unified approach to 3d generation and reconstruction, 2023.
Ting Chen. On the importance of noise scheduling for diffusion models. arXiv preprint 10X1:22301.10772, 2023.
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay. William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mírac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoe V. Le, and Jason Wei. Scaling cnttuuctionfinetuune language models, 2022. URL https: https://arxii. http://aaxvrrgga/r/////eee
2200.220.... 11416.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A Large-scale Hierarchical Image Database. IEEE Conference on Computer Vision and Pattern Recognition, pp- 2009.
Emily Denton, Arthur Szlam, and Rob Fergus. Deep Generative Image Models using a Laplacian Pyramid of Adversarial Networks. NIPS, pp. 1-9, 2015.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8799 2021.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution im- age synthesis. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 12873-12883, 2021.
Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin. Devi Parikh, and Yaniv Taigman. Make- a-scene: Scene-based text-to-image generation with human priors. 2022. doi: 10.48550/ARXIV. 2203.13131, URL. httpp:://axiii http::ts/aaxiiv..orraba/////////
wwtp:///rrrrbb///2..33 13131.
Ian Goodfellow, Jean Pouget-Abadic, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In NeurlPS, 2014.
10

---

Published as a conference paper at ICLR 2024
Jiatao Gu, Lingjie Liu, Peng Wang, and Christian Theobalt. Stylenerf: A style-based 3d-aware generator for high-resolution image synthesis. arXiv preprint arXiv:2110.08983 2021.
Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Miguel Angel Bautista, and Josh Susskind. f-dm: A multi- stage diffusion model via progressive signal transformation. arXiv preprint arXiv:2210.04955, 2022.
Jiatao Gu, Alex Trevithick, Kai-En Lin, Josh Susskind, Christian Theobalt, Lingjie Liu, and Ravi Ramamoorthi, Nerfdiff: Single-iaag view synthesis with nerf-guided distillation from 3d-aware diffusion. arXiv preprint arXiv: 2023.
Yuwei Guo, Ceyuan Yang, Anyi Rao, Yaohui Wang, Yu Qiao, Dahua Lin, and Bo Dai. Animatediff: Animate your personalized text-to-image diffusion models without specific tuning, arXiv preprint arXiv:2307.04725 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02.:00 2022a.
Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusson models for high fidelity image generation. J. Mach. Learn. Res., 23:47-1, 2022b.
Jonathan Ho, Tim Salimans, Alexey A Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. In ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022c.
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. In International Conference Machine Learning, 2023. URL https://api semanticscholaa
Allan Jabri, David Fleet, and Ting Chen. Scalable adaptive computation for iterative generation. arXiv preprint arXiv:2212. 2022.
Zahra Kadkhodaie, Florentin Guth, Stéphane Mallat, and Eero Simoncelli. Learning multi-scale local conditional probability models of images. In The Eleventh International Conference Learning Representations, 2022.
Nal Kalchbrenner, Aäron van den Oord, Karen Simonyan, Ivo Danihelka, Oriol Vinyals, Alex Graves, and Koray Kavukcuoglu. Video pixel networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1771-1779. PMLR, 06-11 Aug 2017.
Minguk Kang, Jun-Yan Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung Park. Scaling up gans for text-to-image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10124-10134, 2023.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv: 2017.
Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ra- manujan, William Howard- -nwaar,Snyder, Kaifeng Chen, Sham Kakade, Prateek Jain, et al. Matryoshka representation learning. Advances in Neural Information Processing Systems, 35:30233-3022 2022.
11

---

Published as a conference paper at ICLR 2024
Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion- Im improves controllable text generation. Advances in Neural Information Processing Systems, 35:4328 4338,,4333400 2022.
Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft COCO: Common Objects in Context. European Conference on Computer Vision, pp. 740-755, 2014.
Haohe Liu, Zehua Chen, Yi Yuan, Xinhao Mei, Xubo Liu, Danilo Mandic, Wenwu Wang, and Mark D Plumbley. Audioldm: Text-to-audio generation with latent diffusion models. arXiv preprint arXiv:2301.12503, 2023a.
Ruoshi Liu, Rundi Wu, Basile Van Hoorick, Pavel Tokmakov, Sergey Zakharov, and Carl Vondrick. Zero-I-to-3; Zero-shot one image 3d object, 2023b.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171. PMLR, 2021.
Aaron van den Oord, Nal Kalchbrenner, Oriol Vinyals, Lasse Espeholt, Alex Graves, and Koray Kavukcuoglu. Conditional Image Generation with PixeICNN Decoders. Advances in Neural Information Processing Systems, pp. 4790-4798, 2016. ISSN 10495258.
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural Discrete Representation Learning. NIPS, 2017.
William Peebles and Saining Xic. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022.
Pablo Pernias, Dominic Rampas, and Marc Aubreville. Wuerstchen: Efficient pretraining of text-to- image models. arXiv preprint arXiv:3306.00637, 2023.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: improving latent diffusion models for high-resolution image synthesis. arXiv preprint 237.02230. 2023.
Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209. 14988, 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. arXiv preprint arXiv:2103.0003 2021.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.0m 06125, 2022.
Severi Rissanen, Markus Heinonen, and Arno Solin. Generative modelling with inverse heat dis- sipation. In The Eleventh International Conference on Learning Representations, 2023. URL https: h/tppr///ppeeevien
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and puttern recognition, pp. 10684-10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net Convolutional Networks for Biomed- ical Image Segmentation. International Conference on Medical Image Computing and Computer- Assisted Intervention, pp. 234-241, 2015.
Dohoon Ryu and Jong Chul Ye. Pyramidal denoising diffusion probabilistic models. arXiv preprint 018444 2022.
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. 1rXiv:2104.001 2021.
12

---

Published as a conference paper at ICLR 2024
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorcalistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 45:36479-36494 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005122.20
Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang. Oron Ashual, Oran Gafni, Devi Parikh, Sonal Gupta, and Yaniv Taigman. Make-a-video: Text-to-video generation without text-video data, 2022.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference Machine Learning, pp. 2256-2265. PMLR, 2015.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 13456, 2020.
Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International conference on machine learning, pp. 1747-1756. PMLR, 2016.
Dirk Weissenborn, Oscar Täckström, and Jakob Uszkoreit. Scaling autoregressive video models. 2020.
Chenfei Wu, Jian Liang, Lei Ji, Fan Yang, Yuejian Fang, Daxin Jiang, and Nan Duan. Nüwa: Visual synthesis pre- vre-trainigg for neural visual world creation, 2021.
Zeyue Xue, Guanglu Song, Qiushan Guo, Boxiao Liu, Zhuofan Zong, Yu Liu, and Ping Luo. Raphael: Text-to-image generation via large mixture of diffussin paths. arXiv preprint 2.t.:2353.119555 2023.
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang. Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, Ben Hutchinson, Wei Han, Zarana Parekh, Xin Li, Han Zhang, Jason Baldridge, and Yonghui Wu. Scaling autoregressive models for content-rich text-to-imen image generation. 2022.
Yizhe Zhang, Jiatao Gu, Zhuofeng Wu, Shuangfei Zhai, Josh Susskind, and Navdeep Jaitly, Plan- ner: Generating diversified paragraph via latent language diffusion model. arXiv preprint arXiv:2306.02531, 2023.
13

---

Published as a conference paper at ICLR 2024
APPENDIX
THANK YOU
Figure 9; Random samples from MDM trained on CC12M dataset at 256 256 and 1024 x 1024 resolutions. See detailed captions in the Appendix G.
14

---

Published as a conference paper at ICLR 2024
A ARCHITECTURES
First, we show the following as the core architecture for MDM for the lowest resolution of 64 x 64. Following (Podell al., 2023), we increase the number of self-attention layers for each resnet blocks for 16 x 16 computations. To improve the text-image correspondence, we found it useful to apply additional self-attention layers on top the language model features.
Base architecture (MDM-S64)
config:
resolutions=[64.33 16]
,12,76811 768]
num_res_ 1ww_rres_blocks=(222.22...
nwm.a...._lyyrs.prr..bocco..
num_heads=8,
schedule='cosine"
emb
emb_channels=1024.
num_in_attm__laerrs=2.
Then, we configure the models for 2562 and 1024 resolutions nested way as follows:
Nested architecture (MDM-5645256)
config:
resolutions==256, 128 128, 64]
resolutiion_chamell---ff,, 128, 256]
inner_config:
resolutions=:64.33 16]
resolutionn.chammels=[(3333 512 768]
num.res.blocks--22..222
L.m.attt_layyrr.pr..bookk-k
num_heads=8.
schedule='cosine"
num_res_blocks=[222.11
num.:tttalaaeespppr..bckkk......
schedule="cosiie-shhff4
emb mtb__chanees=1024.
num_lan.ttmmllayers:2,
in..feeturr.pogggceee..cmmeeel_p
Nested architecture (MDM-564S1288256)
Architecture config (NDM 64, 128, 256):
resolutionns=255, 128]
resolutiin__chamels--ff6 128]
inner_config:
resolutions-[128. 64]
resollttion__hannells:: (attaaannnels11222555592
inner_config
resolutions-(64.32.156
resolution..chamels=-1355 512 768]
num num_res_blockee jum_res_.blocss(2.222]]
nwm.att._layer..er.bloot-
num_heads=8,
schedule='cosiine
nun_res.blockk-12.111
awn.atth_Jrrers.pe..bloko......
schedule="ccssine-shhte
nun_res_blocks-(2.1
wum.attuulaayrr_.pr..blo-- 0]
sciedulee"'coiine-shhtte
emb mmb_channels=1024..
num_In_attn_laaerssz2,
1m_
mm.ffetuurrpppooccceelllaaaaeess
15

---

Published as a conference paper at ICLR 2024
Nested architecture (MDD--564S2566112222
config:
resol 32, 64]
inner_config:
resollutions=((56,122.. 128,
resollattoon.chameels--4444/ 128, 256]
inmer_config:
resolutions=:(4.32,111
resoluttion_channells--ff 512, 768]
num_res_blocks=[222.2]
num_heads=8,
schedule='cosine"
num.res.blocks=[2...11
nun_resbblocks:(222]]1
sciedule='coiiie-shhff...
mmb_chaamees=1024..
num_in_attt_laayers=:,
in..feturreppoeeeeeelccaaeeess_
Nested architecture (AMII00MMMAD0544522222266533133
config:
resolutiions=1102,
resolution_.haamels--2..
inner inmer_config:
resoluttions=(112,25))
resolluttool.ccammllss=.2..2.2 inmer_config:
resolutions=1256.1 128]
resoluction_chamelss-(44.. 128]
inner_config:
resolutions==128. 64]
resolutiinn.chamels=1222 256]
immer_config:
resolutionss:64,,2,, 16]
resolutiion_chammells==:22 512, 768]
num_res_blockks-222
num.atu.layers.per.bbe hwwp::sspppebboooee-e-ee block= 1,5]
num_heads=8.
schedule='cosiine
num_res_bloocss=[2.11
schedule="cocsine-shhhrt
nun_res_block=-[211
aww.atta_layers.pe..bookkk 0]]
schedule=""osiiie-shhfff"4
nun_resbblocks:[21]]
mun_attn_layerrsper.bboot-
schecdule=scosiieeehhfff$$'
nun.res.blocks--2211111
wun.aatt._aaprrr.pe..bbolk
schedulle=rcosiin---hhftte
emb_chhamels-1122..
num num_im_attn_laaerse
In addition, we also show the models for video generation experiments, where additional temporal
attention layer is performed across the temporal dimension connected with convolution-based
16

---

Published as a conference paper at ICLR 2024
sampling. An illustration of the architecture of video modeling is shown in Fig. 10. For ease of visualization, we use 4 frames instead of 16 which was used in our main experiments.
Nested architecture (MDM-S64T16) for video generation
config:
temporal_axis=11
resoolutionn_chmmmlll-55555,,, 256, 256, 256, 256]
inner_config:
resolatiions:[6433..1]]]
resollution.cchmmelss((6666 52,,,11,,6661666 518.
nus_res_bloocks([222.]]
num_heads=8,
schedule='cosiine"
nun_res.biocks:(22222...11
aww.tenpptal.attnnlyyyessseee..p
schelulle=rcoiinee-hhit4
emb_ mobb_chanees=1124..
nun_ln_attn_laaers=2,, num
Nested architecture (MDM-56641165256 for video generation
config:
resolutions==256. 128 128,64) 64]
resoluttion_chammel--(8881, 128, 256]
inner_config:
temporal_axis=Irue
cenporal_resolliioons:116 8, 4,2,1]
resoluttion.cchmmeels-13336 256,256,222 256,
inmer_config:
resollttions:(64432,.6]]
resoluttion_chanmell==:5555 512,
(21,)))222,666,,,,,6 768]
num_heads=8.
schedule='cosine
num.res.bblocks:2222222...1
sciedbler-"osiieeshhfff
numrrs_bloccss((22.]]]]
num.stt..Javerr.ppe...bkkk......
schedule=='ooiiie-shhttt emb_
chbbccaaneess1124.
num num_in_attn_layers=22, in.feetture_ppppeceel..camameesp
B TRAINING DETAILS
For all experiments, we share all the following training parameters except the batch size and training steps differ across different experiments.
default training config:
optimizer='adie adam
adam_betal=0.9 adam_beta2-0.99
adam adam_eps=1.e-8
Learning_rate=le-2
Beaminggraaee_wrmuppssteess:
weight_decay=o.0
17

---

Published as a conference paper at ICLR 2024
conditioning (toxt, lobel, ctc)
Figure 10:
illustration NestedUNeet architecture used Matryoshka Diffusion for video generation. We allocate more computation in the low resolution feature maps, and use additional temporal attention layers aggregate information across frames.
gradient_ - 0 ema_decay=0.. 9999
mixed_precision.traiiinee
For ImageNet experiments, the progressive training setting is set default without specifying:
progressive training config:
target_ressluttoons==444 226]
batch_size=(511,222 256]
craining_steps= 500K]
For text-to-image generation on CC12M, we test on both 256 x 256 and 1024 x 1024 resolutions, while each resolution two types of models with various nesting levels are tested. Note that, the number of progressive training stages is not necessarily the same the actual nested resolutions in the model. For convenience, we always directly initialize the training of 1024 x 1024 training with summarize all experiments into config:
trained model for 256 x 256. Therefore,
progressive training config:
Garget_resolutions-ff6.. ,256, 1824 (optioaaall batch_size=(2048, 1624, 768] training_steps= 500K, 100K]
Similarly, we list the training config for the video generation experiments as follows.
progressive training config:
target_resolutions=
16x64, 16x256] batch (iceh_s8ie= 512 128]
training_steps==500k. 500k, 300K]
C INFERENCE DETAILS
In Fig. 11, we demonstrate the typical sampling process of trained MDM. We start by sampling independent Gaussian noises for each resolution which gives us (= We then pass all the noisy images to the Nested UNet in parallel, which yields the denoised outputs {f
We then perform one step of denoising for each { with { following the procedure of a standard diffusion model, and this gives us (2f-1/r=1... inference steps 250 and v-prediction (Salimans & Ho, 2022) Similar to Saharia et al. (2022), we apply "dynamic thresholding" to avoid over satruation problem in the pixel predictions.
In practice we set the number of model parameterization.
18

---

Published as a conference paper at ICLR 2024
Noisy images across time
Predicted clean images across time
64x64
256x256
1024x1024
64x64
256x256
1024x1024
t=1000
t=856
1=712
t=568
t=424
t=280
t-136
1-1
Figure 11: An example the inference process of MDM for text-to-image generation 1024 1024 with three levels. text caption panda playing guitar garden." "
D ADDITIONAL ABLATIONS
D.I. UNET vs NESTED UNET
As mentioned in Sec. 4, UNet and Nested UNet are very similar in capacity and they yield near identical performances when treated standalone architectures. We verify this in Figure 12, where we train a standard diffusion with a UNet and Nested UNet on ImageNet 256x256. We see that they indeed yield tightly coupled FID curves. In addition, also measure their efficiency during both
19

---

Published as a conference paper at ICLR 2024
200
Nested UNet
UNet
90
80
FID 10K
60
50
40
30
20
0
50
100
150
200
250
300
350
400
450
Training iterations (K)
Comparison between standard UNet and Nested UNet ImageNet 256x256. As expected, they share similar performance, which confirms that Nested UNet standalone architecture (with the multi resolution loss and progressive) does not offer immediate advantage over UNet.
Figure 12:
Table Speed comparison UNet and Nested UNet on ImageNet 256x256. Speed is measured on 8xA 100 GPUs with batch size of 128 for training and 256 for inference.
Training (secfiter)
Inference (sec/step)
UNet
0.50
0.40
Neseted UNet
0.52
0.41
training and inference, as shown in Table 2. It's evident that they have near identical computation costs.
D.2 TRAINING COSTS
For all three datasets/tasks, namely ImageNet, CC12M and WebVid-10M, best results obtained on 4x8 A100 GPUs, for which we here report the total training costs. The results can be seen in Table 3.
D.3 WALL CLOCK TIME
To complete the comparison in Figure 4 (a), we also provide the time efficiency measurements for each the baseline runs (only ImageNet experiments are shown, and the comparison is similar on CC12M). This is shown in Table 4.
Table 3: Training cost our main experiments. INet: ImageNet: CC: CC12M: WebVid-IOM. - denotes progressive training stage, eg. CC 64 + 256 refers to training a 256x256 model from a 64x64 model on CC12M. All experiments use 4x8 A100 GPUs.
INet 64
INet -565
CC
CC
156
CC (C:256-> 1024
CC
Vid 16x64
Vid 16x64
16x256
Sample/GPU Iter()
32
100
375
500
500
100
500
300
See/lier
101
1.21
1.14
1.38
1.88
1.17
2.34
20

---

Published as a conference paper at ICLR 2024
Table Training speed MDM and baseline models in Figure 4 (a). MDM, simple DM and Cascaded DM have practically the same training speed (and the difference is mostly due to variances in the hardware), while LDM is considerably more efficient.
MDM
MDM (no PT)
Simple DM
Cascaded DM
Latent DM
Sec/lter
1.11
0.96
E BASELINE DETAILS
E.1 CASCADED DIFFUSION MODEL.
For our cascaded diffusion baseline models, we closely follow the guidelines from (Ho et al., 2022b) while making it directly comparable to models. In particular, our cascaded diffusion models consist of two resolutions, 64x64 and 256x246. Here the 64x64 resolution models share the same architecture and training hyper parameters as MDM. For the upsampler network from 64x64 to 256x256, we upsample the 64x64 conditioning image to 256x256 and concatenate it with the 256x256 noisy inputs. Noise augmentation applied by applying the same noise scheduels on the upsampled conditioning images, as suggested in (Sabaria et al.. 2022). All the cascaded diffusion models are trained with 1000 diffusion steps, same as the MDM models.
During inference, we sweep over the noise level used for the conditioning low resolution inputs in the range of (1, 100, 500, 700. 1000), similar to Saharia al. (2022). We found that a relatively high conditioning noise level (500, or 700) is needed for our cascaded models to perform well.
E.2 LATENT DISFUSION MODEL
For the LDM experiments we used pretrained encoders from https://githui: com/CompVis/ latent- di ffusion (Rombach al., 2022). The datasets were preprocessed using the autoen- coders, and the codes from the autoencoders were modeled by our baseline U-Net diffusion models. For generation, the codes were first generated from the diffusion models, and the decoder of the autoencoders were then used to convert the codes into the images, at the end of diffusion.
In order to follow similar spatial reduction to our MDM-S644256 model we reduced the 256x256 images to codes at 64x64 resolution for the Imagenet experiments, using the KL-F4 model from https:// http://mmeerrre// lab. com/
f4. zip and we then trained our MDM-S64 baseline model on these spatial codes. However, for the text-to- diffusion exper- iments on CC12M we found that the model performed better we used the downsampling model (KL-F8) from https: / /ommer--aaa 1ab. cwm/filmeeeeee htt/fiile//eee siliesllatene hels/aaeentdddt -ifttent-ifffuioon http/kkl8.... aal-f..2ip zip. However, since this reduced the resolution of the input to our UNet model, we modified the MDM-S64 model to not perform downsampling after the first ResNet block preserve similar computational footprint (and this modification also performed better). The training of the models was performed using the baseline models.
same set of hyperparameters
F DATASETS
contains 1.28M im- ages across 1000 classes. We directly merge all the training images with class-labels. All images are resized 2562 with center-crop. For all ImageNet experiments, we did not perform cross-attention. and fuse the label information together with the time embedding. We did not drop the labels for training both MDM and our baseline models. FID is computed on 50K sampled images against the entire training set images with randomly sampled class labels.
ImageNet (Deng et al., 2009,
CC12M (Changpinyo et al., 2021, https: thub coom/google goo.ggol reseercheeeeechheeee datasetss conceptual-12m is a dataset with about 12 million image-text pairs meant to be used for vision- and-language pre-training. As mentioned earlier, we choose CC12M as our main training set considering its moderate size for building high-quality text-to-image models with good zero-shot capabilities, and the whole dataset is freely available with less concerning issues like privacy. In this paper, we take all text-image pairs dataset set for text-to-image generation. More specifically,
21

---

Published as a conference paper at ICLR 2024
randomly sample 1/ 1/1000 of pairs the validation set where monitor the CLIP and FID scores during training, and use the remaining data for training. Each image by default is center-cropped and resized to desired resolutions depending the tasks. No additional filtering or cleaning is applied.
is large-scale dataset of short videos with textual descriptions sourced from stock footage sites. The videos are diverse and rich in their content. Following the preprocessing steps of Guo et al. (2023)', we extract each file into sequence frames, and randomly sample images every frames create 16 frame long clip from the original video. Horizontal flip is applied as additional data augmentation. As the initial exploration of applying MDM on videos, we only sample one clip for each video, and training MDM on the extracted video clips.
Web Vid-100 (Bain et al., 2021, https://maxberr
com/webvid-date
G ADDITIONAL EXAMPLES
We provide additional qualitative samples from the trained MDMs for ImageNet 256 x 256 (Figs. 13 to 15), text-to-image x56 and 1024x 1024 (Figs. and to 18), and text-to-video 16x256x256 (Fig. 19) tasks.
In particular, the prompts for Fig. 9 are given as follows:
fluff owl with knitted hat holding wooden board with "Thank You" written it (1024 1024),
batman and Joker making sushi together.
squirrel wearing crown stage,
an oil painting of Border Collie, an oil painting of rain at traditional Chinese town, a broken boat a peacel lake, lipstick put front of pumpkins, frog drinking coffee fancy digital Art.
a lonely dog watching sunset, painting of royal girl in a classic castle, realistic photo castle, origami style, paper art, fat cat drives UFO,
teddy bear wearing blue ribbon taking selfie in small boat in the center of lake, paper art, paper cut style, cute bear,
crowded subway, neon ambiance, abstract black oil, gear mecha, detailed acrylic, photorealistic, a groundhog wearing a straw hat stands on top of the table, an experienced chief making Frech soup in the style of golden light, a blue jay stops on the top of a helmet of Japanese samurai, background with sakura tree (1024 x 1024).
22

---

Published as a conference paper at ICLR 2024
Figure 13: Uncurated samples from MDM trained on ImageNet 256 x 256 with the guidance weight 2.5 for labels of "srhinoceros beetle", "Siberian husky", "cliff. drop, drop-off", "coral reef", "space shuttle". "hummingbird".
23

---

Published as a conference paper at ICLR 2024
Uncurated samples from MDM trained on ImageNet 256 x 256 with the guidance weight 2.5 for labels of "sulphur-crester cockatoo, Kakatoe galerita, Cacatua galerita", "Ilama". "loggerhead, loggerhead turtle, Caretta caretta" ", "hot pot, hotpot", "jack-o" llntern". "espresso".
Figure 14:
24

---

Published as a conference paper at ICLR 2024
Figure 15:
Random samples from the trained MDM on ImageNet 256 x 256 given random labels. The guidance weight is set 2.0.
25

---

Published as a conference paper at ICLR 2024
cinematic photo of hand holding a -lanten
ckk
toy bulldozer on a small table in a garden
a Chinese ink painting of trees near a peaceful lake
cinematic photo of an old man with long curly hair, blue eyes and wearing a tophat
Vbrant portrait painting of an Egyptian cat with robotic half face
Figure 16: Uncurated samples from the trained MDM on CC12M 256 x 256 given various prompts. The guidance weight is set 7.0.
26

---

Published as a conference paper at ICLR 2024
80mm resin detailed miniature of fluffy black devil rabbit, hell, red eyes, Product Introduction Photos
A glass bottle containing a sailboat floats through the galaxy
The edge of the Russian forest in summer at night
Figure 17: Random samples from the trained MDM on CC12M 1024 1024 given various prompts. The guidance weight is 7.0.
27

---

Published as a conference paper at ICLR 2024
Concept of wizard blowing the trumpet the top of mountain
medieval knight in the daylight with intricate armor details, holding blooming violet bouquet in his hand. sunshine casts shadows
Old photo of a young girl in 1890s in an apron works at a local bar
Figure 18: Random samples from the trained MDM on CC12M 1024 1024 given various prompts. The guidance weight is set 7.0.
28

---

Published as a conference paper at ICLR 2024
Rude
9 Burd
Burd
shuude
suurs
Burua
Busy freeway night
1991
88
48
T 269
Me 119S
Pouring milk into black coffee
chutte
shutte
shutte chttte
chutte shutte shutt
shutt
A teddy bear riding skateboard by the lake
Figure 19: Random samples from the trained MDM WebVid 16 256 256 given various prompts. The guidance weight is 7.0.
29


Reviews:
[{"review_id": "XjQJS4TICK", "reviewer": "Reviewer_7XqF", "summary": "The paper proposes a diffusion model capable of denoising multiple resolutions simultaneously. To enhance computational efficiency, they use a nested UNet structure. While it is possible to train the model all at once, the results show that progressively training it led to better convergence.", "questions": "It is curious about the computational load increases. Also, it would be nice if the performance difference with usual UNet and NestedUNet is provided.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The proposed model converges faster compared to traditional cascaded diffusion models. In the case of MS-COCO, the model demonstrates superior performance.", "weaknesses": "It is anticipated that there will be an increase in computational load. And the performance does not seem to surpass that of LDM, which employs classifier-free guidance. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yZ7pBh0BVu", "reviewer": "Reviewer_iXvr", "summary": "This paper introduces a new framework for high-resolution image synthesis using diffusion models. Due to computational limitations, diffusion models are often limited to cascaded approaches in pixel-space or operating in latent space. The proposed framework, Matryoshka Diffusion Models (MDMs), denoises images at various resolutions simultaneously. MDM is trained using the standard diffusion objective jointly at multiple resolutions with a progressive schedule where the higher resolutions are added into the objective later in training. The authors demonstrate that MDM has greater efficiency with comparable performance on image and video generation.", "questions": "- Generally the experiments section would benefit from including more insights on the results, especially for the scenarios where MDM is outperformed by the baselines.\n- 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.\n- It is helpful to understand how the proposed multi-resolution prediction affects sampling speed.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed framework is straightforward and easy to understand, drawing inspiration from existing GAN literature to address limitations of current diffusion models for high-resolution synthesis.\n- The authors evaluate MDM on multiple synthesis tasks and demonstrate comparable results. The qualitative results look impressive especially given the relatively small scale of training data.\n- Ablation studies quantify the effects of progressive training and the number of nested levels on the quality and alignment of the outputs.", "weaknesses": "- 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.\n- 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.\n- 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.\n- There are several typos (especially with spacing before citation parentheses on page 8).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yL4kAwbD1o", "reviewer": "Reviewer_Z34F", "summary": "This paper presents a multi-stage image/video generation model based on f-DM. The model is structured as a NestedUNet: a UNet with multiple inputs/outputs of increasing/decreasing resolutions. Compared to f-DM, it incorporates progressive growing and benchmarks the approach on multiple large-scale text-to-image and text-to-video datasets. The works positions itself as a new paradigm for high-resolution diffusion models and rivals latent DMs and cascaded DMs. It features faster training convergence in terms of the amount of iterations compared to the existing paradigms. The obtained results visually look quite good to me.", "questions": "- What are the computational budgets of all the experiments? (I can only see the amount of GPUs being used for the experiments — without a notice on for how long). I believe that the smaller amount of training could also justify inferior FID results on ImageNet compared to SotA.\n- 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?\n- 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 generate them?\n- Please, include the comparison with existing text-to-image generators. Judging by the maximum attainable CLIP score is misleading.\n- 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?\n- 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?\n- 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).\n\nSome typos:\n- \"crtical\" => \"criticial\" (page 2)\n- \"under performs\" => \"underperforms\" (page 6)\n- \"eg\" => \"e.g.\" (page 7)\n- page 8 — no space before the bracket \"(\" in multiple places.\n\n[1] Gu et al \"f-DM: A Multi-stage Diffusion Model via Progressive Signal Transformation\"", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- Visually, the results look very good. And it is especially remarkable given the little compute used to train the models.\n- A good advantage of the given method is that, compared to CDMs, it does not need require the previous stage to be well-trained to have meaningful training of the current stage with reliable scores. For CDMs, while it's possible to train all the stages simultaneously, one cannot generate images from scratch in the middle since the base stage has not been fully trained yet. Somehow, this does not happen for the given model.\n- The ablations and evaluation in general is quite solid.\n- The exposition is good, and the paper is written well.", "weaknesses": "- 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.\n- 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.\n- 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.\n- FID scores on ImageNet are ~3x times higher than the current SotA (e.g., MDT).\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PxTfS1G9nA", "reviewer": "Reviewer_cN8u", "summary": "This paper proposes Matryoshka Diffusion Models for high-resolution image and video synthesis. Unlike conventional approaches that use either cascade models or latent diffusion models with an additional autoencoder, Matryoshka Diffusion Models uses a diffusion process that denoises the multi-resolution input jointly, where such a process can be trained progressively and improves the optimization efficiency significantly. The paper shows the effectiveness of the method on popular image generation and video generation tasks and verifies the training efficiency of the method compared with existing diffusion model variants.", "questions": "- Why some points in Figure 4 (e.g., after 200K of Latent DM in Figure 4(a)) are missing?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is generally well-written and easy to follow. \n- The paper is well-motivated.\n- The paper conducts experiments with various datasets, including ImageNet, MSCOCO, and WebVid-10M.\n- The high-resolution image generation results are quite impressive.", "weaknesses": "- 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. \n- 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. \n- Some important implementation details are missing: learning rate, batch size, model configurations, etc. \n- Missing quantitative evaluation on text-to-video generation compared with existing baselines.\n- No video files included for illustrating text-to-video generation results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)