MC-DiT: Contextual Enhancement via Clean-to-Clean Reconstruction for Masked Diffusion Models
Authors: Guanghao Zheng
Yuchen Liu
Wenrui Dai
Chenglin Li
Junni Zou
Hongkai Xiong
Venue: NeurIPS 2024 poster (2024)
Keywords: Diffusion
Transformer
Masked Autoencoder
Primary area: machine_vision

Abstract: Diffusion Transformer (DiT) is emerging as a cutting-edge trend in the landscape of generative diffusion models for image generation. Recently, masked-reconstruction strategies have been considered to improve the efficiency and semantic consistency in training DiT but suffer from deficiency in contextual information extraction. In this paper, we provide a new insight to reveal that noisy-to-noisy masked-reconstruction harms sufficient utilization of contextual information. We further demonstrate the insight with theoretical analysis and empirical study on the mutual information between unmasked and masked patches. Guided by such insight, we propose a novel training paradigm named MC-DiT for fully learning contextual information via diffusion denoising at different noise variances with clean-to-clean mask-reconstruction. Moreover, to avoid model collapse, we design two complementary branches of DiT decoders for enhancing the use of noisy patches and mitigating excessive reliance on clean patches in reconstruction. Extensive experimental results on 256$\times$256 and 512$\times$512 image generation on the ImageNet dataset demonstrate that the proposed MC-DiT achieves state-of-the-art performance in unconditional and conditional image generation with enhanced convergence speed.

MC-DiT: Contextual Enhancement via Clean-to-Clean Reconstruction for Masked Diffusion Models
Guanghao Zheng, Yuchen Liu, Wenrui Dai ; Chenglin Li, Junni Zou, Hongkai Xiong School of Electronic Information and Electrical Engineering Shanghai Jiao Tong University
Abstract
Diffusion Transformer (DiT) is emerging cutting-edge trend in the landscape of generative diffusion models for image generation. Recently, masked-reconstrrucion strategies have been considered to improve the efficiency and semantic consistency in training DiT but suffer from deficiency in contextual information extraction. In this paper, we provide new insight to reveal that noisy-to-noisy masked- reconstruction harms sufficient utilization of contextual information. We further demonstrate the insight with theoretical analysis and empirical study on the mutual information between unmasked and masked patches. Guided by such insight, we propose a novel training paradigm named MC-DiT for fully learning contextual information via diffusion denoising at different noise variances with clean-to- -oea--n mask-reconttruutton. Moreover, to avoid model collapse, we design two comple- mentary branches of DiT decoders for enhancing the use of noisy patches and mitigating excessive reliance on clean patches in reconstruction. Extensive experi- mental results on 256 x 226 and 512x510 image generation the ImageNet dataset demonstrate that the proposed MC-DiT achieves state-of-the-an performance unconditional and conditional image generation with enhanced convergence speed.
Introduction
have emerged front-runners in the latest advancements of image-level generative models. and surpass previous state-of-the-art generative models [10. [4]. DPMs corrupt an input image into noise obeying the normal distribution by gradually injecting Gaussian noise and recover the image from the noise with step-by-step denoising via a pretrained network
Diffusion Probabilistic Models (DPMs) 19
29
42 43
43 U-Net 38 is popular early works [37 351 to predict noise from
disrupted images for image generation. Recently, Diffusion Transformer (DiT) 34 has been widely
considered for DPMs 29
42
43
conditional image generation
4
video generation
[6 22 30,,
and 3D generation
15
25
36 due to its excellent scalability and superior performance.
Different from vision transformers (ViTs) [1]. DiT is trained predict Gaussian noise from disrupted images at different noise levels. To achieve large-scale training, DiT suffers from slow convergence and heavy computational burden in the training process [49]. Moreover, due to its goal of noise prediction, DiT causes semantic inconsistency in generated images, since it struggles learn contextual information in different regions of images for noise prediction [3) To solve is introduced into the denoising- based training schedule for DiT trained to predict masked noisy patches from the given unmasked noisy patches. MDT 13 pioneers to propose the asymmetrical
these problems, Inspired by masked autoencoder (MAE) [8]
DIT B3
49 48
diffusion transformer that integrates mask-reconse mass-reconsttruction with denoising, which employs encoders to extract features from unmasked noisy patches and a lightweight decoder to reconstruct masked patches via extracted features. MaskDiT 48 accelerates the training process by masking at most 50%
"Correspondence Wenrui Dai <daiwennii@sju..dd....m
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

noisy image patches. SD-DiT 49 introduces self-supervised discriminative objective for knowledge distillation to reduce the fuzzy relation between the maak-reconstruuctinn and denoising. Despite superior performance over vanilla DiT, they are deficient in exploiting contextual information by neglecting different noise scales in different steps of diffusion process.
In this paper, we revisit maso-recontttuutiinn in training DiT and reveal that reconstructing masked noisy patches from unmasked noisy patches harms contextual information extraction. This issue is exaggerated under large noise variances, since both unmasked and masked noisy patches are similar to standard Gaussian noise and contain little contextual information. We demonstrate this phenomenon in Figure aTa by evaluating the mutual information between unmasked output patches and masked patches at different noise variances for different methods. With the growth of
noise variance, mutual information in noisy image patches generated by MDT [B3] and MaskDiT
48 decreases sharply, while mutual information in vanilla noisy images decreases slowly. This fact suggests that the information in masked patches rarely comes from unmasked patches, and thereby, the contextual information is not sufficiently exploited. We further demonstrate this empirical finding with theoretical analysis on mutual information and the mask graph [46), as elaborated in Propositions and 3
To address this problem, we propose MC-DiT to reconstruct clean unmasked patches from clean masked patches rather than resort to noisy patches. Benefiting from the merit that -eenen reconstruction is not influenced by the noise, the proposed MC-DiT is able to learn contextual information via the diffusion denoising process all noise scales. Furthermore, to avoid model collapse caused by over-emphasizing clean patches but neglecting noisy patches, we design two complementary branches to enforce the model focusing more on denoising. In summary, our contributions are listed as below.
We provide a new insight that noisy-to-noisy mask-reconstruution is insufficient in extracting contextual information. We demonstrate the insight on mask-reconstruution with theoretical analysis and empirical study on mutual information between unmasked and masked patches. We propose MC-DiT, novel training paradigm with new mask- recssstconttt objective, to fully exploit contextual information with clean-to-clean reconstruction. We further design two complementary branches of DiT decoders to avoid model collapse and focus on denoising. We evaluate the proposed MC-DiT in x56 256 and 512x 512 image generation on the ImageNet dataset and achieve state-of-the-ee ahrt FID score for DiT backbones of various scales.
2
Related Work
have attracted increasing attention due to their superior image generation ability compared with preceding generative models 14 477. Denoising diffusion probabilistic model (DDPM) [19 significantly advances generative models, particularly in tasks such as text-guided image synthesis. In specific, DDPM is realized as a Markov chain (32) that contains forward process and reverse process. In the forward process, clean images are disrupted by Gaussian noise step by step, and in the reverse process, the images are generated from the Gaussian noise with step-by-step denoising. The commonly used U-Net model (38) is trained to predict the Gaussian noise from noisy images. Score matching method 49 is introduced into diffusion models to design the forward and reverse process with elegant Stochastic Differential Equation (SDE) [44. EDM [2I] analyzes the design space of diffusion models and disentangles the effects of parametrization, sampling, and training. address the time- consuming iterative issue inherent in DPMs, several methods apply fast sampling strategy 22 28 or
Diffusion Probabilistic Models. Diffusion Probabilistic Models (DPMs) 19 42 43
latent diffusion training strategy [37].
Transformers in Diffusion Models. In DPMs, the most commonly used architecture for noise prediction is U-Net (8 which is symmetric encoder-decoder framework. Recently, transformers provide new perspective to noise prediction due to their excellent multi-modall fusion ability and remarkable scaling properties. U-ViT 3 integrates time embedding, image patches, and conditional patches into overall tokens and applies residual connection into transformers for consistency in generation. DiT
34 claims that transformers applied DPMs realize superior performance and achieve the scaling law. Therefore, various works adopt and improve DiT into 2D image generation
R4 (55, video generation 30 [6] 22 and 3D generation [5] 25 36). this paper, we take the DiT as our backbone and change the input from noisy patches to clean patches.

---

23-
Verille Clear image
Mask Graph Clean Patches
Vanilaa Holey image
Generated mage
MDT Generated Image
(XXummssiee) mymskee))
Our MC-DIT Generated Image
15-
h h(x'))ng((*))
1.0
Mask Graph of Noisy Putches
0.5
x
xj
0.0
0.
05
1.0
2.0
2.5 5.0
Noise Variance
h
(a)
(b)
Figure 1: (a) Mutual information of different methods between generated masked patches and unmasked patches. We generate masked noisy patches from unmasked noisy patches and calculate mutual information ((X (unmmaseed, X Xmmakke) under various noise scales. The 'vanilla clean images' and 'vanilla noisy images' denote the real clcan/noisy images, which are the upper bound of the mutual information. The other lines are computed with the generated images by three strategies. (b) Mask graph [46] different reconstruction targets. The left yellow ellipse denotes unmasked patches and the right green denotes masked patches. The black arrow denotes the positive pairs to pull in.
Masked Training with Transformers. reoosreccotiitutti has been broadly applied into con- volutional networks and transformers []]. MAE [8 takes the mask-reconstruutton to pretrain transformers and achieves stunning contextual modeling capabilities and zero-shot performance. Inspired by this method, various strategies for masked training have been introduced into transformers. For example, ConvMAE [12] leverages masked convolution to prevent information leakage. FreMiM 45 converts images into frequency domain and applies masked training for frequency information reconstruction. MultiMAE [ utilizes masked training into multi- multi-modality inputs for cross-modal fusion and generation. SdAE [8] improves masked autoencoder via self distillation. In summary, the mask-reconstruutton training objective transfers the information from unmasked patches to masked patches, and thereby, enhances contextual semantic modeling ability.
Proposed Method
3.1 Preliminaries
Masked AutoEncoder (MAE) is significant unsupervised pretraining paradigm in computer vision, which reconstructs masked patches from unmasked patches. Given an image x, MAE first patchifies into N patches denoted by á E R.Vxe, where the channel dimension. A random binary mask m {0. 1 is applied to obtain masked patches
Masked AutoEncoders 18]
x 1-1m] N.xe and unmasked patches Ø2 = z [m]]] E N2xx N1 and N2 are the number of masked and unmasked patches. An encoder-decoder framework h g is then applied. The encoder f maps the unmasked patches 1 into latent space _ = fx11), while the decoder reconstructs the pixel value of masked patches x/2 g(21). The MAE is trained ensure the reconstruction ability by minimizing the Mean Square Error (MSE) loss LMAE = E,
- x2ll?. U-MAE [46] provides theoretical understanding of MAE and establishes connection between MAE and contrastive learning 33
Proposition 1 ([46) The lower bound of the MAE loss is
MMA 2 -Ez,,,,((z)) (2(2)) E + const Lasym E + const,
where Casym denotes the asymmetric aligmment loss 46). is the fitting error, and hg o ff is the pseudo autoencoder that satisfies E, (12())) x2 2
U-MAE [66 combines Proposition T with the mask graph TigureT[i) (upper) to represent the contrastive objective in MAE. Specifically, Proposition I demonstrates that the MAE loss is equal
3

---

to the contrastive loss Laym. which calculates the similarity of h(x1) and (g())). If and 22 are neighboring patches, Laym is minimized when positive pairs (i.e., 21 and T2) are pulled closer. Proposition D is consistent with the mask graph in Figure (b), where unmasked and masked patches are considered as contrastive pairs. Thus, we employ mask graph as an effective tool to analyze the contrastive objective of MAE.
Diffusion Probabilistic Models (DPMs) emulate the diffusion process of physical atoms to convert the standard Gaussian distribution into the target distribution via differential equations. In the forward process, clean data Iy (Zata(Zo) corrupted into Gaussian noise TT NO, o maz step by step via stochastic differential equation (SDE):
Diffusion Probabilistic Models 19
42 43
dz = fzz, t)dt + g(t)dw,
(2)
where f is the drift coefficient, is the diffusion coefficient, w is a standard Wiener process, and / is the time value from 0 to T. The reverse process generates target samples using the inverse SDE:
dx [f(,t) g(t))((())))))) log (x))ddt + g())deū
where wi is a reverse-time Wiener process. Following EDM [2]] to set f(x, 1) = 0 and g(() = v 21, the forward and reverse SDEs are reformulated as dr = v 2tdw and dr = -tV, a log (r()ddtt where s(x, v log (x)) is the score function. To solve the reverse-time SDE, denoising network Da(x, f) is trained to minimize the score matching loss:
E BxgrrPatcc (oo Do + n,t) -
(4)
As a result, the score function is estimated by s(x, t) (Do xt/ During training, at step t, noisy images zo + n are sent to the denoising network tz to predict the clean images z0- However, directly optimizing this objective leads to poor contextual information [3] To solve this problem, the mask-reconstruution is applied into denoising network
48 49
3.2 Contextual Information in Noisy Patches Reconstruction
We first review the mask-reconstruucion between noisy patches and point out that applying noisy patches reconstruction task into the training process of DiT is ineffective and leads to insufficient contextual information utilization. With the mutual information and mask graph, we propose two propositions demonstrate this claim, where the first is for the mutual information of input-output patches and the second is for the contrastive objective of input unmasked ummasked-masked masked patches.
Mutual information between input and output patches. Besides the mutual information between masked and unmasked output patches in Figure [a). we consider the mutual information between input noisy patches and output (noisy and clean) patches. MaskDiT (48) and MDT 13 reconstruct masked noisy output patches from input unmasked noisy input patches, whereas SD-DIT recovers unmasked clean output patches from masked noisy input patches, as elaborated below.
MaskDiT & MDT. Noisy patches I; at step are obtained by injecting noise (0, (21) into clean patches Tn- Masked and unmasked noisy patches are generated from 2;, by =! = xz [m]]
and wi 2 [1 m using a random binary mask m. x² is reconstructed from x by minimizing the MAE loss MMakk-MAt = Ez, Ex!. =0] (x))) x²?1/ for the encoder f and decoder g.
The ability exploit contextual information is measured by mutual information Ixx
SD-DiT. SD-DiT extracts latent features of x and concatenates them with masked noisy patches
xt? to predict clean patches z/. The MAE loss LSD-MAB Ez, Ex; _aa is. (l: x²²) z/l??
The ability to exploit contextual information measured by mutual information I( xx))
The contextual information in both two formulations is transferred from the noisy patches to noisy patches. Subsequently, we analyze the contextual information utilization ability of mask- reconstruction via calculating mutual information of masked and unmasked patches.
2 Although the reconstruction targets of SD-DiT are clean patches, it is equivalent to distinguish the of and noise Therefore, the contextual information is used for better prediction of n.

---

Proposition 2
Given masked and unmasked clean patches zg and 2ó and their noisy versions z/
and zi, the mutual information = I( x and [( a satisfy that
= I( p(==)) p=
K =
(5)
xx) ~ [(zo: = Ep(= Bp (x))
KL 1p(=))))))
)p(e))BB(ct) (m)))
[KL (p
(6)
Proposition 2 suggests Suggesss that the mutual information Z( = and ==));=)) are lower than the ideal
mutual information [(x) z9). With the growth of t. the KL divergence terms in (5) and (6) increase due larger noise perturbation [( z5) becomes larger and the ability to extract contextual information is degraded.
zg and z Thus, the gap between I I( z;). [[=];=²] x²) and
Mask graph. We further analyze the rass-recontttructinn via contrastive objectives in mask graphs. In U-MAE 46), the maso-reconttruution can be transformed into a constrastive learning objective and there exists a bipartite mask graph to elaborate this transformation [46]. The mask graph is consistent with Lasym Proposition [ Note that MaskDiT, SD-DiT, and MDT share the same mask graph, since their inputs are all unmasked noisy patches and noisy masked patches. Figure Tbb) illustrates the mask graph for clean image reconstruction reconstruction in MaskDiT, SD-DiT, and MDT (bottom). In Proposition 3 prove that contrastive objective between noisy patches could interfere learning contextual information.
MAE (top) and that for noisy patch
Proposition 3
The asymmetric loss of noisy patch reconstruction and the asymmetric loss of clean patch reconstruction satisfy that
Laemm-NN = EEp(j .a7) [x(xt)"(,,z7])
(āhg 0ug.
Oh Our
Rd Laeym + IE - h(()))
+ E (2))
where Lusym is defined in ① and represents contextual information. The two noise- wrise-weighted items represent contrastive objective between h((z) and [0hg/Oxi) (hy (x7) and (0h/axt)) weighted by the Gaussian noise n.
Proposition 3 implies that the Gaussian noise introduces extra terms in 7 and could affect the optimization process of Lasym. Noisy patch reconstruction undermines the contrastive objective of contextual information, since larger Gaussian noise leads to more severe perturbation on Losym-
In summary, we leverage mutual information and contrastive asymmetric loss to demonstrate that the noisy patches mask-reconstruction is sub-optimal to learn real contextual information and larger noise could have more serious impact on context information extraction. This is consistent with the results in Figure aa). To solve this problem, in Section [3.3] propose MC-DiT to effectively reconstruct masked clean patches from unmasked clean patches.
3.3 Contextual Enhancement with Masked Clean Patches
As demonstrated in Propositions 2 and that reconstructing masked noisy patches from unmasked noisy patches is insufficient for contextual information extraction, we propose a novel method named MC-DiT to enhance contextual information extraction for DiT from the perspective of leveraging masked clean patches to reconstruct unmasked clean patches. Figure 2 depicts the proposed framework for MC-DiT. The clean images are disrupted by Gaussian noise where t is the time step. Then the noisy images are patchified and masked by a random binary mask m. The unmasked noisy patches z are fed into the DiT encoder for feature extraction. For contextual information extraction, the masked clean patches zi are concatenated with extracted feature concat (21, .g),, where 21 is the feature of a/. After that, the feature is sent to DiT decoder to reconstruct clean unmasked patches xh) which is consistent with previous masked diffusion [88
NO, t² I),
49
The training objective of unmasked clean patch reconstruction is:
Lelean Exorpasta (0 (1) (at.aarrlllpppeeee (Do((20 + +n) @ (1 11-m,,20 xo O m, t) - xo) o (1 - m)ll. (8)
By applying masked clean patches =f in (8. the information =3 is transferred to unmasked clean output patches žt, which is constrained to equal zf. It is not disrupted by noise n, since there is no
5

---

Decoder
Food
Clean Patches Target
noisy
shared
patches
DiT Decoder
DIT
Encoder
Decoder
Noisy Patches Target
mask
DIT Encoder
Gaussion
clean potches
shared
noise
musk and insert
patchify
Decader
(a) The pre-training paradigm MC-DIT
(b) Finetuning
Figure 2: Framework of the proposed MC-DiT. (a) Pre-training. MC-DiT introduces unmasked clean patches and learns sufficient contextual information by reconstructing unmasked clean patches from masked clean patches. Two complementary EMA branches are developed avoid model collapse. (b) Finetuning. MC-DiT is trained with unmasked patches for denoising.
noise in zj and zg. Morcover, as mentioned in Section 3.2 the mutual information for our MC-DiT is [xz]i z5), since leverage masked clean patches to recover unmasked clean patches. According to Proposition 2] [( =σ) is much higher than I(rb; =²) and I(r]; x²) under different noise n and time steps t. Thus, the mutual information learned by MC-DiT would not decrease for different
noise at different time steps and sufficient contextual information could be learned for reconstruction.
In addition, explore the potential benefits of z by enforcing the masked output patches 2² to match u?. As discussed in Section 3.1, the denoising network Do predicts clean images from noisy images and can be viewed as distinguishing the clean images and noise. Therefore, predicting x? from z is to recognize the noise n and add to z² To this end, we employ reversed constraint on masked output patches z², as illustrated in Figure 2 a). The training objective is
Lnoise EEvorpaaoa No
+ +n) C (1 m), To m,,) (2o n)) m||2. (9)
3.4 Addressing Model Collapse
Although MC-DiT learn sufficient contextual information in theory, there exists model collapse problem in practice. The model learns a shortcut way that it only leverages masked clean patches to reconstruct unmasked clean patches and neglect the unmasked noisy patches. We further address the model collapse problem caused by only using masked clean patches to reconstruct unmasked clean patches for training. We introduce two extra EMA (Exponential Moving Average) branches of DiT decodery"]to to balance the mask-reconstruction and denoising objective. As shown in Figure 2 given the noisy input to DiT encoder, we introduce two branches to achieve only mask-reccnstruuttoon and denoising alongside the original DiT decoder. The upper branch of DiT decoder D reconstructs from the whole noisy patches and captures denoising information, while the bottom branch processes clean patches and captures contextual information Dep- Constraints on the two branches are employed the loss function to balance the DiT decoder.
Loon = Paate E hat
(10)
Lron Patta B P
(11)
'The parameters of EMA decoders are initialized with those in the DiT decoders and are updated in the EMA fashion according to the parameters in DiT decoders: Orma x Pema + (1 a) x lave. where a denotes the weight coefficient. The two decoders updated only using the EMA method without using gradient updates.

---

Table 1: Comparison with state-of-the-art approaches for Image class conditional image
generation. Bold font represents the best result. 'G'' means using classifier-free guidance.
Methods
FID
sFID +
Prec.
Rec.
BiGAN-degy 5
6,95
7.36
171.40
0.87
0.28
StyleGAN-XI.
2.30
4.02
265.12
0.78
0.53
MaskGIT 6
6.18
182.10
0.80
0.51
CDM 20
4.88
158.71
ADM O
10.94
6.02
100.98
0.69
0.63
ADM-U
7.49
5.13
127.49
0.72
0.63
LDM-8
15.51
79.03
0.65
0.63
LDM-4 37
10.56
209.52
0.84
0.35
U-ViT-H/2 2
6.58
DiT-XL/2 34
9.62
6,85
121.50
0.67
0.67
MDT-XL//2
6.23
5.23
143.02
0.71
0.65
MaskDiT-XL/2 48
5.69
10.34
177.99
0.74
0.60
SD-DIT-XL/2
7.21
5.17
144.68
0.72
0.61
MC-DiT-XL/2
4.14
6.96
309.69
0.83
0.62
ADM-G [9]
4.59
5.25
186.70
0.82
0.52
ADM-U-G 9
3.94
6.14
215.84
0.83
0.53
LDM-8-G
7.76
103.49
0.71
0.62
LDM-4-G 37
3.60
247.67
0.87
0..8
U-ViT-H/2-112
2.29
5.68
263.88
0.82
0.57
DiT-XL/2-G [34
2.27
4.60
278.24
0.83
0.57
MDT-XL/2-GT [3
1.79
4.57
283.01
0,81
0.61
MaskDiT-XL/2-G
48
2.28
5.67
276.56
0.80
0.61
MC-DIT-XL/2-G
1.78
4.87
290.17
0.81
0.62
Here, Do and Di EMA DiT decoders. Therefore, the overall loss is:
Celean + A Cnoiie + A2hcont + A3Ccon2.
(12)
where Å1 A2 and A are hyper-parameters.
Unmasking Finetuning. Similar to MaskDiT 48
although MC-DIT captures contextual
information during masking training, directly applying the pretrained model in inference leads to unsatisfactory performance, which is caused by the training-nnferent discrepancy. The clean patches used in training are not provided in the inference time. Thus, after training our MC-DiT, we finetune it the unmasked scenarios for better performance, shown in Figure 2 (b). It is worth noting that MC-DiT only needs a few iteration in the finetuning to generated semantic coherence images.
Experiments
4.1 Implementation Details
Model Settings. We follow the same architecture in MaskDiT pretrained variational autoencoder (VAE) from Stable Diffuion 37 space, and then train our MC-DiT to reconstruct clean patches from noisy patches under the EDM (1 framework to approximate score function in the diffusion process. The pretrained VAE maps 256x256x3 input images to x2x32x4 latent features and 512x512x: x3 images to 64x64x4 latent features. Similar to SD-DiT, we apply DiT-S, DiT-B, and DiT-XL as our backbones.
488
Specifically, we first apply a map the images into latent
Training Details. Similar to previous works
48 we train MC-DiT on ImageNet [39] with resolutions 256x 256xx and 512x512x3 respectively. Most training settings are the same as MaskDiT 48 We train MC-DiT for 400K to 1M iterations using the Adam optimizer with learning rate 0.0001 and weight decay. By default, use 50% mask ratio and batch size 1024. X and A [2 are set to 0.1 and 0.05 for more denoising reconstruction. The EMA coefficient is set to 0.999 for smoothness and data augmentation employed.
Evaluation Metrics. Following DiT 34), we leverage Fréchet Inception Distance (FID) to measure the quality of generated images. For fair comparison, also ADM's Tensorflow evaluation

---

Table 2: Comparison with state-of-the- approaches for ImageNet-5122x512 class conditional image
generation. The bold font represents the best performance.
Methods
FID
sFID +
IS
Prec.
Rec.
BiGAN-deep [5]
8.43
8.13
177.90
0.88
0.29
StyleGAN-XI -LL 41
2,41
4.06
267.75
0.77
0.52
ADM [
23.24
10.19
58.06
0.73
0.60
ADM-U 9
9.96
5.62
121.78
0.75
0.64
DiT-XL
34
12.03
7.12
105.25
0.75
0.64
MaskDiT-XL//2
48
10.79
13.41
145.08
0.74
0.56
MC-DiT-XL/2
9.30
6.28
179.58
0.76
0.53
9DM-G[9]
7.72
6.57
172.71
0.87
0.42
ADM-U-G(
3.85
5,86
221.72
0.84
0.53
DiT-XL /2G [4]
3.04
5.02
240.82
0.84
0.54
MuskDiT-XLL2-G
48
2.50
5.10
256.27
0.83
0.56
MC-DiT-XL/2-G
2.03
4.87
272.19
0.84
0.56
Table 3: Comparison with state-of-the-arr approaches ImageNet-256x 256 class conditional image generation at different scales and iterations. 'S'' 'B'', '-XL' means 'small', 'base" and largest model size, respectively and '/2' denotes the patch size of 2 for all input patches.
Methods DIT-S/2 S
Training Iterations
FID-50K +
400K
68.40
MDT-S/2 B
400K
53,46
SD-DiT-S //
449
400K
48.39
MC-DiT-S/2
400K
41.67
D/T-B/2 [B4]
400K
43.47
MDT-B// [S]
400K
34.33
SD-DiT-B /2 49
400K
28.62
MC-DiT-B/2
400K
18.88
DIT-XL/2 B4
7000K
9.62
MaskDiT-XL/22
48
1300K
12.15
MDT-XL/2 B
1300K
9.60
SD-DIT-XL/2
1300K
9.01
MC-DiT-XL/2
1300K
7.92
suite [9] to compute FID-50K (FID for short), sFID 31, sion/Recall (24 as secondary metrics. More vivid images have lower FID and sFID, while their IS and Precision/Recall are higher.
Inception Score (IS) 40 and Preci-
4.2 Experimental Results
We evaluate vanilla training (i.e., LDM 37
341) and masked training (i.e., proposed MC-Dit, MaskDiT[48), MDT 13, and SD-DIT 49 using backbones of different scales for 256 and 512x512 image generation on ImageNet.
ADM 9. and DiT
Results ImageNee-2566 255.. Table shows that MC-DiT-XL/2 achieves the smallest FID score and the highest IS score. Compared with non-masked diffusion models, MC-DiT-XL./ decrease the FID score from 9.62 to 4.14. Compared with masked diffusion models, the FID score decreases from 5.69 to 4. 14. With classifier-fres guidance (-G), our MC-DiT-XI./2-C achieves the best FID score 1.78, and the highest IS score, which significantly outperforms previous works.
Results on ImageNet-512 x 512. Table shows that MC-DiT-XL/2 achieves a FID of 9.30 and outperforms MaskDiT 48 and DiT 34 The IS score of MC-DiT is also the highest, indicating the effectiveness of our method. With classifier free guidance (-G), our MC-DiT-XL/2-G achieves the best FID score of 2.03, indicating the effectiveness of MC-DiT.
(a) reports the mutual information between unmasked and masked output patches with different noise, which can be viewed as the metric of contextual informa- tion consistency. Our MC-DiT decreases slowly during the noise variance becomes larger, which indicates more sufficient contextual reconstruction regardless noise.
Contextual Enhancement. Figure I

---

Trakking Loss
725
150
FID-50K
12
75-
25-
200 (K)
Training loss
(b) FID
Figure 3: Training loss and FID for DIT-B/2, MaskDit-B/2, and MC-DiT-B/2 during training.
MDT
MaskDIT
MC-DIT
MDT
MuskDiT
MC-DIT
MDT
MuskDiT
MC-DIT
Figure 4:
Comparison of 256 x 256 images generated by MDT, MaskDiT and MC-DiT, Various details are strange in images generated by MDT and MaskDiT.
Backbones different scales. Table evaluates FID-50K at different scales and training iterations with various backbones. Notice that MaskDiT only reports the performance of 'XL' scale. Under fixed number of training iterations, MC-DiT outperforms vanilla DiT (34,, MDT [3], and SD-DiT [49] in FID by large margin, i.e., 6.72, 9.74, and 1.09 FID reduction for DiT-S, DiT-B, and DiT-XL backbones. This fair comparison fully demonstrates the effectiveness of our method.
Convergence speed. In order to evaluate the convergence speed of various methods, we compare the training loss curve Figure a). We report the MSE loss (Eq. (8) on clean patches for fairness. We train MaskDiT [48] and DiT 34 for 300K iterations due to the substantial time and GPU resource
overhead and use the training curve our MC-DiT trained for evaluations, which trained for 400K iterations. The training loss of MC-DiT decreases faster than DiT [34] and MaskDiT [48] Figure ()) measures FID-50K at each step after unmasked tuning and shows that MC-DiT achieves the lowest FID-50K score.
Generated image comparison. Figure 4 visualizes the 256 256 images generated by MDT ik] MaskDiT[48] [ and our MC-DiT. Our generated images are more realistic and have more consistent textual structure than MaskDiT and MDT. For example, images of 'hammer' generated by MaskDiT and MDT have incomplete structure, while our MC-DiT generates images with more complete structures, validating the superior contextual information extraction ability our MC-DiT.
4.3 Ablation Studies
For computation efficiency, we adopt 'BB' in the models for fair comparison, All the models are trained for 400K itreations with batch size 256 and mask ratio 50%.

---

Table 4: Ablation study of hyperparameters.
Table 5: Ablation study of the EMA branches.
FID
A
FID
FID
Branches
FID
0
43.23
0.0
38.83
0.0
37.77
Main Branch 22.10
0.01
40.99
0.01
36.15
0.05
35,20
W Noisy Branch 19.26
0.1
35.20
0.1
35.20
0.1
35.46
Clean Branch
18.88
1.0
38.97
1.0
37.54
1.0
36.35
Table 6: Comparison with different targets.
Table 7: Ablation study of unmasking tuning.
Reconstruction Target
FID-50K +
Strategy MaskDiT-XL/2 UT
Iterations
FID
All the clean patches
34.53
1300K
12.15
Only clean patches
25.88
MC-DiT-XL/2 w UT
1300K
7.92
Clean patches + Noisy patches
22.10
MC-DiT-XL/2 w/o UT
1300K
8.33
Main branch target. We evaluate the reconstruction targets of main branch by considering three cases, iee, clean patch reconstruction noisy patch reconstruction, all clean patch reconstruction and only clean patch reconstruction. 'All clean patches' means all the patches are constrained by clean reconstruction loss. 'Only clean patches' means only unmasked patches are constrained by clean reconstruction loss, and 'Clean patches Noisy patches' means masked patches are further constrained by noisy reconstruction loss. Table shows that our model performs the best, which validates the effectiveness of noisy reconstruction loss.
Effectiveness of two extra EMA branches. Table the influence of EMA branches. FID decreases obviously by 2.84 using the noise branch, indicating the necessity of noisy branch to address model collapse. Further experiments can be found appendix.
Unmasked tuning. Unmasked tuning (UT) can reduce the training-inferent discrepancy, as demon- strated by MaskDiT and is adopted our MC-DiT. However, we can remove unmasked tuning to reduce the complexity at little loss on FID. Table 7 shows that FID will increase by only 0.41 1 for MC-DiT by removing unmasked tuning.
Hyperparameters. We separately evaluate four values for Å1. X2, and Äs- Table 4 shows best FID is obtained when Å = 0.1, X2 0.1, and A 0.05. Note that Aę larger than Xs since the denoising objective is more important than contextual information utilization.
2_
Figure 5: Ablation study of mask ratio.
Mask ratio. Figure visualizes the influence of the mask ratio in m. FID is smallest at the mask ratio of 50% and increases rapidly when the mask ratio is larger than 50%.
5 Conclusion
In this paper, summarize the previous works that combine mask-reconstruction with DiT training and claim that reconstructing masked noisy patches from unmasked noisy patches is insufficient for contextual information extraction. To validate this claim, we analyze the mutual information and contrastive objective theoretically and experimentally. Besides, we propose a new pretraining paradigm (dubbed MDD-DTT), which reconstructs unmasked clean patches from masked clean patches and guarantees the contextual information extraction. Moreover, to avoid model collapse, two extra EMA branches are applied MC-DiT to adjust the balance between the mask-reconttruution task and denoising objective. Extensive experiments demonstrate the robustness of our method and our MC-DiT achieves the state-of-the-an performance in image generation.
Despite excellent performance, the training speed and inference speed MC-DiT still needs to be improved. We will mitigate this issue in future work by transferring the information in the encoder into the decoder, which decreases the training difficulty.
Limitations.
Acknowledgment. This work was supported in part by the National Natural Science Foundation of China under Grant 22125109, Grant T2122024, Grant 62320106003, Grant 62371288, Grant 62431017, Grant 62401357, Grant 62401366, Grant 1023, Grant 61932022, Grant 1200106607
10

---

References
[1] Roman Bachmann, David Mizrahi, Andrei Atanov, and Amir Zamir. MultiMAE: Multi-modal multi-task masked autoencoders. In European Conference on Computer Vision, pages 348-367. Springer, 2022.
[2] Fan Bao, Chongxuan Li, Yue Cao, and Jun Zhu. All are worth words: vit backbone for score- score-based diffusion models. In NeurlPS 2022 Workshop on Score- Baree Methods, 2022.
[3] Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A VIT backhone for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 27669922679, IEEE, 2023.
[4] Georgios Batzolis, Jan Stanczuk, Carola-Bil Bibiane Schönlicb, and Christian Ermann, Conditional image generation with score-based diffusion models. arXiv preprint arXiv:2111.1300 2021.
[5] Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In The Seventh International Conference Learning Representations, 2019.
[6] Huiwen Chang, Han Zhang, Jiang, Ce Liu, and William Freeman. MaskGIT: Masked generative image Transformer. In Proceedings of the IEEE/CVF Conferencee on Computer Vision and Pattern Recognition, pages 11315-11325. IEEE, 2022.
[7] Ting Chen, Simon Komblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of 37th International Conference on Machine Learning, pages 1.597-1007. PMLR, 2020.
[8] Yabo Chen, Yuchen Liu, Dongsheng Jiang, Xiaopeng Zhang, Wenrui Dai, Hongkai Xiong, and Qi Tian. SdAE: Self-distillated masked autoencoder. In European Conference on Computer Vision, pages 108-124. Springer, 2022.
[9] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems 34, pages 8780-8794, 2021.
[10] Carl Doersch. Tutorial on variational autoencoders. arXiv preprint arXiv: 1606.6600000000 2016.
[11] Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transfomers for image recognition at scale. In The Ninth International Conference on Learning Representations, 2021.
[12] Peng Gao, Teli Ma, Hongsheng Li, Ziyi Lin, Jifeng Dai, and Yu Qiao. ConvMAE: Masked convolution meets masked autoencoders. arXiv preprint arXiv:2205.03892 2022.
[13] Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion Transformer is strong image synthesizer. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 23164-23173. IEEE, 2023.
[14] Ian Goodfellow, Jean Pouget-Abadie Mehdi Mirza, Bing Xu, David Warle-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11) 11:39 144, 2020.
[15] Pengsheng Guo, Hans Hao, Adam Caccavale, Zhongzheng Ren, Edward Zhang, Qi Shan, Aditya Sankar, Alexander Schwing, Alex Colburn, and Fangchang Ma. StableDreamer: Taming noisy score distillation sampling for text-to-31e 3D- arXiv preprint arXiv:2312.022
2023.
[16] Xun Guo, Mingwu Zheng, Liang Hou, Yuan Gao, Yufan Deng, Chongyang Ma, Weiming Hu, Zhengjun Zha, Haibin Huang, Pengfei Wan, el al. 12V-Adapter: A general image-to-vided adapter for video diffusion models. In SIGGRAPH 24: ACM SIGGRAPH 2024 Conference Papers, number 122 pages 1 -22 ACM, 2024.
[17] Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self-supervised deep learning with spectral contrastive loss. In Advances in Neural Information Processing Systems 34, pages 5000-5011, 2021.
[18] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16000-16009. IEEE, 2022.
[19] Jonathan Ho Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems 33, pages 6840-685 164516651,155 2020.
11

---

[20] Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research, 23(47):1-33, 2022.
[21] Tero Karras, Miika Aittala. Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In Advances Neural Information Processing Systems 35, pages 16565-26577 2022.
[22] Dan Kondratyuk, Lijun Yu, Xiuye Gu, Jose Lezama, Jonathan Huang, Grant Schindler, Rachel Hornung, Vighnesh Birodkar, Jimmy Yan. Ming- Chang Chiu. Krishna Somandepalli. Hassan Akbari, Yair Alon, Yong Cheng. Joshua V. Dillon, Agrim Gupta, Meera Hahn, Anja Hauth, David Hendon, Alonso Martinez, David Minnen, Mikhail Sirotenko, Kihyuk Sohn, Xuan Yang. Hartwig Adam, Ming-Hsuan Yang. Irfan Essa, Huisheng Wang. David Ross, Bryan Seyhold, and Lu Jiang. VideoPoct: A large language model for zero-shot video generation. In Proceedings of the s1 International Conference on Machine Learning, pages 25105-25124 124. PMLR, 2024.
[23] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[24] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. In Advances in Neural Information Processing Systems 32, pages 3927-3936, 2019.
[25] Chen-Hsuan Lin, Jun Gao, Luming Tang, Towaki Takikawa, Xiaohui Zeng, Xun Huang, Karsten Kreis, Sanja Fidler, Ming- Yu Liu, and Tsung- Yi Lin. Magic3D: High-resolution text-to-3D content creation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 300-309. IEEE, 2023.
[26] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep lcarning face attributes in the wild. In Proceedings of the IEEE International Conference on Computer Vision, pages 3730-3738. IEEE, 2015.
[27] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-Solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps. In Advances in Neural Information Processing Systems 35, pages 5775-5787, 2022.
[28] Cheng La, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-Solver++: Fast solver for guided sampling of diffusion prohabilistic models. arXiv preprint arXiv:221 2022.
[29] Jianjie Luo, Yehao Li. Yingwei Pan, Ting Yao, Jianlin Feng, Hongyang Chao, and Tao Mei. Semantic- conditional diffusion networks for image captioning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 23359-23368. IEEF, 2023.
[30] Xin Ma. Yaohui Wang, Gengyun Jia, Xinyuan Chen, Ziwei Liu, Yuan- Li, Cunjian Chen, and Yu Qiao. Latte: Latent diffusion transformer for video generation. arXiv preprint arXiv:2401.03048, 2024.
[31] Charlic Nash, Jacob Menick, Sander Dieleman, and Peter W Battaglia. Generating images with sparse representations. In Proceedings of the 38th International Conference on Machine Learning, pages 7958- 7968. PMLR, 2021.
[32] James R Norris, Markov Chains. Cambridge University Press, 1998.
[33] Aaron van den Oord, Yazhe and Oriol Vinyals. Representation learning with contrastive predictive coding, arXiv preprint 188ivv1807.03344 2018.
[34] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195 4205. IEEE, 2023.
[35] 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. In The Twelfth International Conference on Learning Representations, 2024.
[36] Ben Poole, Ajay Jain, Jonathan T Barron. and Ben Mildenhall. DreamFusion: Text-to-3 using 2D diffusion. In The Eleventh International Conference Learning Representations, 2023.
[37] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent difffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Partern Recognition, pages 10684-10695. IEEE, 2022.
[38] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional networks for biomedical image segmentation. In 18th International Conference on Medical Image Computing and Computer-Asiisted Cittreattoo-MIICAA 2015, pages 234-241. Springer, 2015.
12

---

[39] Olga Russakovsky. Jia Deng, Hao Su, Jonathan Krause. Sanjeev Satheesh, Sean Ma. Zhiheng Huang. Andrej Karpathy, Aditya Khosla, Michael Bemstein, Alexander C. Berg, and Li Fei-Fei. ImageNet large scale visual recognition challenge. International Journal of Computer Vision, 115:211-252, 2015.
[40] Tim Salimans, Ian Goodfellow. Wojciech Zaremba, Vicki Cheung. Alec Radford, and Xi Chen. Improved techniques for training GANs. Advances Neural Information Processing Systems 29, pages 2234-2242. 2016.
[41] Axel Sauer, Katja Schwarz, and Andreas Geiger. StyleGAN-XL: Scaling StyleGAN to large diverse datasets. In ACM SIGGRAPH 2022 Conference Proceedings, pages 1 1-10. ACM, 2022.
[42] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. In Advances in Neural Information Processing Systems 33, pages 12438-12448, 2020.
[43] Yang Song. Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Poole. Score-based generative modeling through stochastic differential equations. In The Ninth International Conference on Learning Representations, 2021.
[44] Nicolaas G Van Kampen. Stochastic differential equations. Physics Reports, 24()):111-28,, 1976.
[45] Wenxuan Wang, Jing Wang. Chen Chen, Jianbo Jiao, Yuanxiu Cai, Shanshan Song, and Jiangyun Li. Fremim: Fourier transform meets masked image modeling for medical image segmentation. I Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 7860-7870, 2024.
[46] Qi Zhang, Yifei Wang, and Yisen Wang. How mask matters: Towards theoretical understandings of masked autoencoders. In Advances in Neural Information Processing Systems 35, pages 27127-27133 2022.
[47] Xingjian Zhen, Rudrasis Chakraborty, Liu Yang, and Vikas Singh. Flow-based generative models for learning manifold to manifold mappings. In Proceedings of the 35th AAAI Conference Artificial Intelligence, pages 11042- 11052. AAAI, 2021.
[48] Hongkai Zheng, Weili Nie. Arash Vahdat, and Anima Anandkumar. Fast training of diffusion models with masked transformers. In Transactions on Machine Learning Research (TMLR), 2024.
[49] Rui Zhu, Yingwei Pan, Yehao Li, Ting Yao, Zhenglong Sun, Tao Mei, and Chang Wen Chen. SD-DIT: Unleashing the power of self supervised discrimination in diffusion Transformer. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 8435 8445, IEEE, 2024.
13

---

A Supplemental Material
A.1 Theoretical Proof
Proposition 2 Given masked and unmasked clean patches zig and xo and their noisy versions z! and oi, the mutual information ; x.). [(z] x.)), and [( xx) satisfy that
=(z): == I(ab: 25)) Ep(x)) P(ac (g))
K [
(13)
x( x I(zb; Ep(=3) mm))
x
Pp(c))pp(at (mb)
(14)
Proof. Given the time step 1,masked noisy patches z², masked clean patches z clean unmasked
patches vd and noisy unmasked patches z7. where zf = = + n, = =d + n and n (I).
derive the mutual information [( zf) according to the definition.
p(zgl=3)
[( Eptx) p( log
(15)
p(x())
p
p(abla))
Pltt P log
(16)
p(a))
p(atlag)
p(abl1n)
= EEp(x) E p(= Pp(u) (uf) log
(17)
p(a))
p(abl=n)
Ep(a ((2a)(())))) log p(x (x))
+ Eptab ag log plan x
(18)
p(xz)
p m
p log p(z =g)
= + E p(un)
pla (=))
p(
(19)
E ( =[aa) - Ep
K [KLpp
(20)
Thus, the mutual information between noisy masked patches and unmasked clean patches I( x?) is less than [==)==0) due the non-negativity of KI. divergence. Moreover, during / increases, the variance of the Gaussian noise n becomes larger. As a result, noisy masked patches x? are disrupted heavily from clean patches zū. The distribution p(= the derivation of KL divergency can be written as:
(z) is very dissimilar from p(zj/=]) Formally,
p(ab at) pp(
E P((=3) Ee(e)1u3) (p(abla)) (m3) log
p(x; mb)
p(z1)
B p(o log
p(= zg)
p(z))
p(= xb) p(n|26)
p(=)) + p(n)
E p paoglng log
p p (x))
p(a
[1+ p(nlarg) p(aglat)
p(n)
-Epca Ep(efl (3) ( log
(21)
p(xz)
We approximate p(af) R p())) + p(n). since p(x;) is Gaussion distribution with mean value 2f and variance (2, As increases, the KL divergence K Lp(x wt) [p(()(=?) increases and the mutual information I( x²) achieves the larger difference with I(r ag). Thus, the mutual information I( x²) is lower than I(a]; =z²)-
Similarly, the mutual information between noisy patches I(x]: x can be drived according to Eq. 20
I(xl; ~ I(z]; 2;) - E Ep(ag) Ep(e!la))
K KL (
(22)
I( = E Ep(e
E Bn(tt 1E (=())
[LL((>) (
(23)
Therefore, the proposition has been proved.
14

---

The asymmetric loss of noisy patch reconstruction and the asymmetric loss of clean patch reconstruction satisfy that:
Proposition 3,
Lasym- NN = -Bplut aaf) [h(x))
hg (=?)]
Thg
Oh T Jeryy
Lasym + IE -h(x))"
E hi (x))
(24)
where Lasym is defined in ① and represents contextual information. The two noise-weighted items represent contrastive objective between h((z)) - -0hg/Oxūj and (2]) [0h/ weighted by the Gaussian noise n.
Proof. According to 1, the asymmetric loss can be written as:
Lasgm- - NN = -E plat .²)) [h(zt)) (( hg(())l
(25)
}
[8hg" [ixz
E
hg (1 h(xt)²²
o(xi))
(26)
@hh 8x
EE
hxx))² 7 hg + h(xx)²
(27)
Oh
Obhg 7 axrg
-E
h((1)) +
n + o(x())
h (a)) + h(x²)²
"1
(28)
abg Dxy
Öh
E [b()t) (a (=:) + -h(x))²
n
+ E
- hg (2²)
azc
n
(29)
Shg
T
Oh T Ourg
Lasym +
-1(x)))"
+ n
E
hg (w0)
(30)
where we leverage first-order Taylor's formula in Eq. db and un. since xi = xb + n and x² = x + n. o denotes the the higher order infinitesimal.
26 and Eq. 28 to calculate h(x)) and (x3) at
A.2 Experiment Details
Diffusion Settings. We leverage EDM 21 as our diffusion training framework, which predicts clean image patches from noisy images. For fair comparison, we use the default parameters in EDM (see [21] for more details). During inference, we generate conditional images from Gaussian noise via EDM-samples (2I). Specifically, the time steps in the reverse process are set via f;
(Lmar + N (1= min ) where N = 40, 7, Imax = 80 and 4min 0.002. Besides, the scooddorder correction is applied and the generated images are the average of first-order and second-order results.
We follow the LDM 37 and adopt a pretrained VAB to firstly map the images into the latent spaces. The weight of pretrained VAB is from Stable Diffusion (37]. Then, we train the denoising models with these latent features. We leverage AdamW optimizer with learning rate 0.0001, batch size 256, and 50% mask ratio. As for unmasking fintuning, we slightly change some hyper-parameters with learning rate 0.00005, batch size 128, mask ratio 0%. Some details can be found in T
Training Details.
A.3 Supplementary Experiments
Generalization Experiments. We adopt the ImageNet dataset in the experiments for fair compari- son, since MaskDiT [88), SD-DIT[49) and MDT[[3] are all evaluated on the ImageNet (3ta In fact, our MC-DiT can be generalized to different domains or datasets for improved image generation
15

---

which leads to the balance on the denoising Table 8: Experimental details about MC-DiT.
MC-DiT-B/2
MC-DiT-XL/2
MC-DiT-XL/2
Resolution
256 x 256
256 x 256
512 512
Training Time
50h
586h
623h
Inference Time (50K images)
12h
8h
15.2h
GPUs
2x RTXX-300 GPUs
4x V100 GPUs
4x V100 GPUs
Batch Size
256 2
256 4
128 x
Memory Usage per GPU
17GB
20GB
27GB
Table 9: Performance comparison on Cifar10 and CelebA dataset of MaskDiT and MC-DiT
CifarI0
CelebA
Methods
FID
Methods
FID
MaskDiT-B/2
11.52
MaskDiT-B/2
7.14
MC-DiT-B/2
9.28
MC-DiT-B/2
5.36
Table Performance comparisons with different branches.
Branches
FID
Main Branch
22.10
w/ Noisy Branch
19.26
w/ Clean Branch
18.88
Main Branch (unmasked noisy patch only)
25.72
Main Branch (masked clean patch only)
23.69
Main Branch (unmasked noisy patch only) w/ Noisy Branch
19.84
Main Branch (unmasked noisy patch only) w/ Clean Branch
19.57
due to the fact that it can extract contextual information from arbitrary images. Table 9 compares the performance of MaskDiT[48] and MC-DiT on the CIFAR-10 10 (3 and CelebA [26 (that collected for face anti-spoofing) datasets. Due to time limit, we train both MaskDiT and MC-DiT for 200K iterations. Experimental results show that MC-DiT outperforms MaskDiT 48 on both datasets.
Convergence Speed. In Figure 3 we compare the training curve DiT 34 MaskDiT 48 and
MC-DiT and point out that the training loss of MC-DiT decreases faster than DiT 34 and MaskDiT [48] This is due to the primary focus of our analysis is the overall effectiveness of the model. The blue line can achieve lower loss, despite similar iterqation counts for flattening, highlighting the model's efficiency in reaching a more optimal solution. Besides, the loss reported in Figure 3(a) denotes the MSE loss Lelean- Thus, lower MSE loss means the generated clean patches more similar to the ground-truth Moreover, MC-DiT achieves lower MSE loss with the same iterations with DiT 34 and MaskDiT (48), indicating the performant model with higher convergence speed,
Main Branch Target. The modal collapse occurs when the main branch only considers clean-to- clean mask-reconstruction for masked clean patches but ignores the denoising of unmasked noisy patches. We propose two EMA branches to balance the two tasks for the main branch. We use the noisy EMA branch to realize noisy-to- to-clann mapping for denoising, and the clean EMA branch to realize clean-to-clean mapping for mask-reconstruucton (mask ratio is 0%). The two EMA branches constrain the output of the main branch (minimize the MSE loss between the outputs of the main branches and EMA branches) via three task and clean-to clean mask-reconstrruttonn task.
To verify this, we report in Table 10 the FID score of the main branch with noisy and clean patch inputted only. The result of the main branch with unmasked noisy patches only is higher than that of masked clean patches, indicating the modal collapse problem. With noisy and clean branches, the FID score of the main branches decline distinetly, validating the effectiveness of the EMA branches.
Ablation Study hyperparameters. Following MaskDiT 44), we select 0.01, 0. and 1.0 as scaling values of three hyperparameters and supplement various values for ablation study. Table [3 Table 14 and Table evaluate evaaluute various values of the three hyperparameters and we find that the best FID is still obtained when \ = 0.1, X2 0.1, and A 0.05.
16

---

Noise Varisnce-1LS
Naie Varioace-L1
Clean Inage
MaskDiT
Figure 6: Feature visualization of MaskDiT and MC-DiT at different noise variance. Better viewed by zoom in.
Table 11: speed denotes the number of iterations per second.
Parameters and training cost comparison between MDT, MaskDiT and MC-DiT. Training
Setting
256 x 256
Modal
Params
FLOPs
Mem
Speed
FID
MDT-XL/2
742M
28G
20G
1.22
6.23
MaskDiT-XL/P
730M
24G
18G
3.09
5.69
MC-DiT-XL/2
786M
26G
20G
1.45
4.14
512 x 512
Modal
Params
FLOPs
Mem
Speed
FID
MDT-XL/2
742M
64G
28G
0.83
MaskDiT-XL/7
730M
56G
24G
1.98
10.79
MC-DiT-XL/2
786M
60G
27G
1.05
9.30
Attention Map Visualization. Figure 6 visualizes the attention map of MaskDiT and MC-DiT at different noise variance with Cifarl0 dataset (23). larger noise variance denotes the noise with large scale. Our MC-DiT extracts proper shape for various noise scale, while the features extracted by MaskDiT are messy in the large noise scale. This further proves the motivation and effectiveness of paper that clean-to-clean mask reconstruction promotes learning sufficient contextual information.
Training cost comparison. We compare the training cost (parameters, FLOPs, memory used and training speed) on 4x V000 GPUs in Table The training speed of MC-DiT a little bit slower than other methods due to two EMA branches. However, the inference speed of MC-DiT is similar to MaskDiT, since two EMA branches are removed during inference. The additional overhead of MC-DiT is relatively small (7.6% parameters and 8% FL.OPs), but the FID performance improvement is significant.
EMA Branch with DiT encoder. In the main branch of MC-DiT, the unmasked noisy patches are fed into the DiT encoder, while all the noisy patches are directly inserted into the EMA DiT decoder to avoid modal collapse, shown in the Figure 2. The reasons are on the two folds: (1) efficient. Only apply DiT decoder for EMA branches leads to small extra parameters and fast inference speed, while EMA DiT encoder slows down the entire EMA branches. (2) effective. The DiT decoder is trained to extract masked clean images patches in the main branch. Thus, directly apply image patches as input of EMA DiT decoder does not lead to poor denoising results. As shown in Table 12 applying FMA DiT encoder introduces extra 669M parameters, while FID score only decreases 1.35. Thus, to balance the parameters and performance, select DiT decoder in the EMA branches.
17

---

Table 12: Performance and parameters comparisons with and without DiT encoder in EMA branches.
Branch
Params
FID
DiT Decoder
56M
18.88
DiT Decoder + DiT encoder
725M
17.53
Table 13: Ablation study X1
A
0
0.01
0.03
0.05
0.07 0.09
0.1
0.3
0.5
0.7
0.9
1.0
FID
43.23
40.99
39.23
38, 44
37.95 36.53
35.
20
35.98
36.74
74
36.91
37.52
38.97
Table 14: Ablation study on A2
A2
0.01
0.03
0.05
0.07
0.09 0.1
0.3
0.5
0.7
0.9
1.0
FID
38.83
36.15
36.02
36.46
35.99 36.07
35.20 20
35.34
36.18
37.26
35.98
37.54
Table 15: Ablation study on X3
A
0
0.01
0.03
0.05
0.07
0.09
0.1
0.3
0.5
0.7
0.9
1.0
FID
37.77
37.25
36.63
55,20
36.07
37.93
35 46
35.88
37.2 26
36.19
19
37.40
36.35
"tniled frog'
oxygen
"sulphur butterfly"
lampshade'
"bubble"
'pencil
"mailbox
Knace pad"
"eustard
"scorn squash"
apple
'school bus"
'overskin
'pale'
"mobile
'bonse
'pillow'
theater
Figure Visualization of 256 256 images generated by our MC-DiT.
A.4 Generated Samples
Figure 7 visualizes some images generated by MC-DIT with 256 256 resolutions. Figure visualizes some images generated by our MC-DiT with 512 x 512 resolutions.
18

---

NW
Figure 8: Visualization of 512 x 512 images generated by our MC-DiT
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We clearly clarify our claim in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss our limitation in the conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
20

---

Justification:
We clearly describe the theoretical results and proof in Sec 3.2 and supplemen-
tary.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
We clearly clarify the setting and parameters of experiments in Sec 4.2.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer: [No]
Justification: We have provide the core file of our code in the supplementary. And the code will be released upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https /nipss.cc/pulli//guu guwde//ddee/// Cot//ooodeeesiseeeeeee wwwpbbmmsssioooooooooooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We describe this in details in Sec 4.1.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: We focused primarily on the exploratory analysis and preliminary results. Addressing statistical significance and error bars will be priority in our future research to provide more comprehensive evaluation of our findings.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
22

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative emrr rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [No]
Justification:
We will report this upon acceptance.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We have adhered to all ethical guidelines and standards throughout our study,
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is no societal impact of the work performed.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
23

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification:
We briefly discuss this in conclusion.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We follow the source code of MaskDiT, which is credited properly.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The pretrained weight and source code will be released upon acceptance.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
25


Reviews:
[{"review_id": "XfuFpwgZjQ", "reviewer": "Reviewer_ZkfZ", "summary": "This paper critiques previous masked-reconstruction strategies in DiT training for their poor contextual information extraction, attributing this to noisy-to-noisy reconstruction. The authors theoretically and empirically validate that this approach limits mutual information between unmasked and masked patches. To address this, they propose a new training paradigm, MC-DiT, which uses clean-to-clean mask-reconstruction combined with diffusion denoising at varying noise levels. To prevent model collapse, they design two complementary DiT decoder branches to balance the reliance on noisy and clean patches. Model collapse would happen in this context due to excessive reliance on clean patches for reconstruction, leading to insufficient utilization of noisy patches and imbalanced training. Extensive experiments on the ImageNet dataset show that MC-DiT achieves state-of-the-art performance in both unconditional and conditional image generation, with faster convergence.", "questions": "- Where can I see: Line 38 \"Despite superior performance over vanilla DiT, they are deficient in exploiting contextual information by neglecting different noise scales in different steps of diffusion process.\" Is there a citation which discusses this is important, or this concluded from your Figure 1 and table of results Table 1?\n- It is not obvious from equation 5 and 6 that in Line 163: \"With the growth of $t$, the KL divergence terms in (5) and (6) increase due to larger noise perturbation on $x^1_0$ and $x^2_0$\" should be true. This can be understood intuitively, but there is no \"decay\" term with respect to $t$ in these equations to suggest that. Can this be formalised with respect to strength of the gaussian noise $n$. Also, I do realise that due to non negativity of KL divergence the two expectation terms subtract from it to make the mutual information smaller, but I so not see how this is true let's say between t and t+1.\n- Line 216: how are the 2 branches of DiT trained in the EMA fashion here (student teacher or are do they also collect gradients)?\n- Line 203: \"$\\mathcal{I}(x^1_0; x^2_0)$ is much higher\", the **much** part is not clear from Proposition 2.\n- Figure 3, is the training loss that is logged for all the models the same? i.e. $\\mathcal{L}_{\\text{clean}}$? or for your method is it the composite loss?\n- Related to Figure 3, when we talk about speed of convergence in terms of iterations it does not say anything about the wall clock time (or FLOPs or Memory) that an iteration takes. In this adapted method, we do x3 forward passes through the DiT decoder, therefore how do the wall clock times (or FLOPs or Memory) compare? From a practical standpoint, this should be clear. Hypothetically, do you also expect the other methods to make up the difference in performance if they were trained for a proportional time longer.\n- Figure 3 a and b, why is the plot only shown for different number of total iterations?\n- Figure 3b, are these metrics calculated before or after fine-tuning for your method?\n- Line 309 in Limitations. Why does the inference speed need to be improved? Is the model inferred differently and requires more steps?\n\n*Minor Typos*\n- Line 13: MDT mentioned before it is defined, although this clear from the citation\n- Equation 1: it was not clear $\\mathcal{L}_{\\text{asym}}$ was defined as the expectation term, this lead to some confusion in the discussion later in Proposition 3.\n- Proposition 2, Equation 6 should end with a full stop.\n- Equation 8 9, 10, 11: brackets are not matched, duplicate sencond closing brackets?", "limitations": "Questions listed above.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper motivates the need for their research in their introduction and is an interesting idea.\n- The paper adds to the mathematical discussion surrounding image generation using diffusion using well understood mutual information metric prevalent in other areas of computer vision.\n- Presents experimental evaluation, with section on reproducible details and supplementary materials.", "weaknesses": "- While reading the article, there are many questions that arise which effect the reading experience of the article.\n- The main weakness of the paper is that at many occasions claims are made which are intuitive, but they are attributed to be implied from an equation / proposition which do not (at least not immediately) show the claim to be true. Look at questions for more details.\n- Some experiment details are unclear (in questions).\n- Table 1 is a bit difficult to read here with the number of methods and it is not obvious how the horizontal lines are drawn, i.e. what makes them different from other quadrants. I think there is enough space for a column or two to add a bit more detail instead of adding them all to the name of the method.\n- Figure 3 (a) is used to showcase speed of convergence. However, I think the distinction between convergence and a convergence to lower loss should be made. All 3 lines more or less flatten at the same time, you could actually argue the red and orange line are flattening faster. I agree the blue line is lower, but that does not mean it has converged faster, only converged to a lower loss. This also leads to a second point, a lower loss here does not necessarily mean a more performant model. As you notice in your own experiments, you require fine-tuning to make the output desirable. Therefore, I disagree that the model converges faster on the whole.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Q1:Claim in Line 38\nWe provide both theoretical and empirical evidences in Proposition~2 and Figure 1(a) in the manuscript to support the claim. We consider the mutual information $\\mathcal{I}(x_0^1;x_0^2)$ between unmasked patches $x_0^1$ and masked patches $x_0^2$ as the contextual information. \nProposition 2 points out that previous methods only learns $\\mathcal{I}(x_0^1;x_t^2)$ and $\\mathcal{I}(x_t^1;x_t^2)$, which is insufficient for $\\mathcal{I}(x_0^1;x_0^2)$. Figure 1(a) demonstrates that the mutual information learned by previous methods (MaskDiT, DiT and MDT) all declines greatly during the noise scale increases, indicating that these methods struggle to learn contextual information under different noise scales. \n### Q2:Equations(5) and (6) and Claim in Line 163.\nThe details of the KL divergence in the Proposition 2 are as follows\n$$\nE_{p(x_0^2)}E_{p(x_t^2|x_0^2)}E_{p(x_0^1|x_0^2)}\\log \\left[\\frac{p(x_0^1|x_t^2)}{p(x_0^1|x_0^2)}\\right] \\\\\n=E_{p(x_0^2)}E_{p(x_t^2|x_0^2)}E_{p(x_0^1|x_0^2)}\\log \\left[\\frac{p(x_t^2|x_0^1)}{p(x_0^2|x_0^1)}\\times \\frac{p(x_t^2)}{p(x_0^2)}\\right] \\\\\n\\approx E_{p(x_0^2)}E_{p(x_t^2|x_0^2)}E_{p(x_0^1|x_0^2)} \\log \\left[\\frac{p(x_0^2|x_0^1)+p(n|x_0^1)}{p(x_0^2|x_0^1)}\\times \\frac{p(x_0^2)+p(n)}{p(x_0^2)}\\right] \\\\\n$$\nwhere $x_t^2=x_0^2+n$ with $n\\in\\mathcal{N}(0,t^2I)$. We approximate $p(x_t^2)\\approx p(x_0^2)+p(n)$, since $p(x_t^2)$ is a Gaussion distribution with mean value $x_0^2$ and variance $t^2$. As $t$ increases, the KL divergence in Proposition 2 increases and the mutual information $\\mathcal{I}(x_0^1;x_t^2)$ and $\\mathcal{I}(x_t^1;x_t^2)$ achieve the larger difference with $\\mathcal{I}(x_0^1;x_0^2)$.\n### Q3:Training of the 2 branches of DiT\nThe parameters of EMA decoders are initialized with those in the DiT decoders and are updated in the EMA fashion according to the parameters in DiT decoders. $\\theta_{ema}=\\alpha \\times \\theta_{ema}+(1-\\alpha)\\times \\theta_{dec}$\nwhere $\\alpha$ denotes the weight coefficient. The two decoders are updated only using the EMA method without using gradient updates.\n### Q4:Much part of $\\mathcal{I}(x_0^1;x_0^2)$\nWe demonstrate in Proposition 2 that the difference between $\\mathcal{I}(x_0^1;x_0^2)$ and $\\mathcal{I}(x_0^1;x_t^2)$ lies in the KL divergence under various noise scale and further verify the theoretical results in Figure 1(a). Figure 1(a) shows that, with the growth of noise, the difference between $\\mathcal{I}(x_0^1;x_0^2)$ (red line) and $\\mathcal{I}(x_0^1;x_t^2)$ (gray and yellow lines) becomes larger. During the increase of noise variance from 0.0 to 1.0, the mutual information of MaskDiT and MDT decline 90% (1.95 vs 0.21), while mutual information of vanilla noisy image only decline 18% (2.07 vs 2.50). This decline refers to the KL divergence in Proposition 2. \n### Q5:Training loss in Figure 3\nIn Figure 3, $L_{clean}$ in Equation (8) is adopted as the training loss for all the models. Figure 3 shows that our MC-DiT converges to a lower value than other methods in terms of $L_{clean}$ and demonstrates the effectiveness of clean-to-clean mask reconstruction.\n### Q6:Different numbers of total iterations in Figures 3a and b.\nTo compare training performance, we train MaskDiT and DiT for 300K iterations in Figures 3a and b due to the substantial time and GPU resource overhead. We directly use the training curve of our MC-DiT trained for evaluations. MC-DiT is trained for 400K iterations for a fair comparison with other methods. In fact, we can find from the training curves for the first 300K iterations that, MC-DiT can obviously decrease the training loss and FID score in comparison to MaskDiT and DiT. This result implies the improvements of our MC-DiT. \n### Q7:Metrics in Figure 3b.\nThe FID reported in Figure 3(b) is calculated after unmasked tuning for MaskDiT and MC-DiT. It is consistent to the results reported in Table 3 in the manuscript.\n### Q8:Inference speed in limitations\nWe intend to improve the inference speed of diffusion model that requires multiple steps to achieve image generation, since MC-DiT is based on the diffusion model and shares the same architecture with MaskDiT. In fact, MC-DiT does not infer differently or requires more steps and yields the same time consumption in the inference stage, since the EMA branches and noisy target are removed during test.\n### Q9:Evaluation of speed of convergence.\nWe follow Figure 4 in SD-DiT to evaluate the speed of convergence for each iteration in Figure 3. Besides, as for the wall clock time, we report FLOPs and Memory used by MC-DiT in each iteration in Table R-1(a) in one page pdf. MC-DiT achieves 3x forward passes, thus the FLOPs and Memory are calculated via the summation of 3x forward process. Training speed of MC-DiT is slower than MaskDiT, leading to longer convergence time. However, MC-DiT converges to a lower loss and achieves superior FID score than MaskDiT. Meanwhile, due to the lightweight DiT decoder, MC-DiT achieves small extra FLOPs and Memory in the training stage. \nThe loss curve of MaskDiT may not decrease due to the small decrease magnitude after 100K iterations. And the FID score of previous methods doesn't decline faster than MC-DiT according to the decrease magnitude in Figure 2(b) of the manuscripts.\n### Q10:Speed of convergence shown in Figure 3(a).\nAlthough convergence to a lower loss is not bound to imply a faster convergence in terms of iterations, the primary focus of our analysis is the overall effectiveness of the model. The blue line can achieve a lower loss, despite similar iteration counts for flattening, highlights the model's efficiency in reaching a more optimal solution. \nBesides, the loss reported in Figure 3(a) denotes the MSE loss $\\mathcal{L}_{clean}$. Thus, lower MSE loss means the generated clean patches are more similar to the ground-truth, indicating the performant model. Moreover, we report the results with and without unmasked tuning in Table. Our MC-DiT outperforms MaskDiT without unmasking tuning.", "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": "Z7zT24HopD", "reviewer": "Reviewer_YSTZ", "summary": "This paper proposes a training strategy for diffusion transformers that fully learns contextual information by introducing clean to clean mask reconstruction during training, and designs complementary DiT decoder branches as well as corresponding supervisory losses to avoid the problem of model collapse, giving theoretical and experimental validation.", "questions": "1.\tWhy not release the qualitative results as proof of the effectiveness of the strategy?\n2.\tHow much does parameter tuning in the final testing phase affect the degree of merit of the final result? How can it be shown that it is the training strategy that is at work and not the parameter tuning?", "limitations": "The methodology proposed in the paper needs to be fine-tuned to the task during the inference phase, otherwise good results may not always occur. Further, it can be found that similar strategies that utilize unknown information in the inference phase for training can have application limitations, which are not conducive to extending and applying the strategy to tasks that do not have access to potentially clear images and information.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1.\tSufficient theoretical analysis\n2.\tThe overall writing of the paper is logically clear", "weaknesses": "1.\tThere are errors in the description of parts of the paper, e.g., x1 in lines 107 and 109 of the introductory section of Masked AutoEncoders is described as masked and unmasked, respectively.\n2.\tVisualization of experimental results is indeed missing, and only quantitative experimental results exist in the body of the paper.\n3.\tUsing the training strategy in the paper, although it can improve the results, it is not possible to conclude the size of the contribution of the training strategy to the final experimental results, as parameter tuning is still required in the testing phase.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Q1: Writing error in line 107.\nWe thank the reviewer for point out the writing error. The unmasked patches $x_1$ and masked patches $x_2$ in line 107 are corrected to $x_1=x[m]$ and $x_2=x[1-m]$. \n### Q2: Visualization results.\nWe have provided visualization results of generated $256\\times 256$ images in Figure 5 in the supplementary materials. Figure 5 shows that images generated by MC-DiT achieve vivid details and diverse styles. In this rebuttal, we further provide more visualization results for both generated $256 \\times 256$ and $512\\times 512$ images in the Figure R-1 in the global rebuttal file. It demonstrates that images of the 'school bus' and 'custard apple' in Figure R-1(a) exhibit different styles, while the details appear very realistic. The same thing happened in $512\\times 512$ images in Figure R-1(b), various images (e.g. 'dog', 'fox' and 'penguin') exhibit rich details and realistic styles, validating the effectiveness of our MC-DiT. \n### Q3: Parameter tuning required for testing.\nWe would like to emphasize that it is the enhanced ability of contextual information extraction by the clean-to-clean patch reconstruction rather than parameter tuning contributes to the superior performance of our MC-DiT. Existing masked diffusion models like MaskDiT and SD-DiT require parameter (unmasked) tuning to decrease the training-inference discrepancy. Different from these models, MC-DiT leverages clean-to-clean patch reconstruction to enhance the contextual information extraction for generation beyond the unmasked tuning to decrease the training-inference discrepancy. Experimental results further demonstrate the effectiveness of the clean-to-clean patch reconstruction in MC-DiT.\n\n**Experimental evaluations.** To validate the effectiveness of the clean-to-clean patch reconstruction, we compare MaskDiT-XL/2 (with unmasked tuning) and MC-DiT-B/2 with and without unmasked tuning. All the models are trained for 1300K iterations. As reported in below Table, both MC-DiT-B/2 with and without unmasked tuning outperform MaskDiT-XL/2 by evident margins of 4.27\\% and 3.82\\% in FID. By contrast, masked tuning only leads to a reduction of 0.41 in FID. These results demonstrate that the performance gain of our MC-DiT mainly comes from the enhanced ability of contextual information extraction by clean-to-clean patch reconstruction rather than unmasked tuning.\n|Strategy|Iterations|FID|\n|:---:|:---:|:---:|\n|MaskDiT-XL/2 w unmasked tuning|1300K|12.15|\n|MC-DiT-XL/2 w unmasked tuning|1300K|7.92|\n|MC-DiT-XL/2 w/o unmasked tuning|1300K|8.33|", "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": "rGidQtU6K5", "reviewer": "Reviewer_RKJM", "summary": "In this work, the authors reveal the issues of Diffusion transformers of having semantic inconsistency as they fail to learn the contextual information. Based on their theoretical analysis, they proposed a novel training paradigm to fully learn contextual information with clean-to-clean mask reconstruction. The paper is well organised and written.", "questions": "1. The authors gave a thorough analysis on the issues of diffusion transformers in section 3.  However, the motivation for the proposed MC-DiT to solve the issues is not very clear. \n2. The steps mentioned in section 3.3 are not so clear and probably are not cohesive with Figure 2. For instance, it mentioned as ‘the unmasked noisy patches x_t^{1} are fed into the DiT encoder for extraction’, but it seems those unmasked noisy patches go to the DiT decoder (?). I might be better to put the denotations on Figure 2 as well to guide readers.\n3. In Table 1, the results on using classier-free guidance were reported for ImageNet-256x256 generation. However, when it comes to the ImageNet-512x512 generation, they are ignored and not reported. Any particular reason behind?", "limitations": "The authors mentioned that the training speed and inference speed still need to be improved and a possible future mitigation on the issue.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The authors have a comprehensive understanding of issues and the state-of-the-art. In terms of originality and quality, the work is technically sound in general. The analysis and written are clear in general.", "weaknesses": "Please see the list of questions for improvement and clarification on some of the aspects.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Q1: Motivation for MC-DiT\nIn Section 3 of the manuscript, sufficient analysis is provided to claim that reconstructing masked noisy patches from unmasked noisy patches is insufficient for contextual information extraction. In details, the information used in noisy-to-noisy patch reconstruction only lie in $\\mathcal{I}(x_0^1;x_t^2)$ or $\\mathcal{I}(x_t^1;x_t^2)$, which are less than $\\mathcal{I}(x_0^1;x_0^2)$. Thus, our motivation is to directly model $\\mathcal{I}(x_0^1;x_0^2)$ via the clean-to-clean patch reconstruction, which reduces the impact of the noise and learns sufficient contextual information. This is realized by inserting masked clean patches into the unmasked noisy features for unmasked clean patches reconstruction. The contextual information flows from the masked clean patches to the unmasked clean patches, which equals $\\mathcal{I}(x_0^1;x_0^2)$.\n### Q2: Denotations on Figure 2.\nThank you for the comment and we will put the denotations on Figure 2 in the revised version. Here, we would like to explain the reason that `$x_t^{1}$ are fed into the DiT decoder for extraction as below. \n\nIn the main branch, the unmasked noisy patches are fed into the DiT encoder, while all the noisy patches are directly inserted into the EMA DiT decoder to avoid modal collapse, as shown in the Figure 2 in the manuscript. The reasons are on the two folds: (1) efficient. Only apply DiT decoder for EMA branches leads to small extra parameters and fast inference speed, while EMA DiT encoder slows down the entire EMA branches. (2) effective. The DiT decoder is trained to extract masked clean images patches in the main branch. Thus, directly apply image patches as the input of EMA DiT decoder does not lead to poor denoising results. As shown in below Table , applying EMA DiT encoder introduces extra 669M parameters, while FID score only decreases 1.35. Thus, to balance the parameters and performance, we select DiT decoder in the EMA branches.\n|Branch|Params|FID|\n|:--:|:--:|:--:|\n|DiT Decoder|56M|18.88|\n|DiT Decoder+DiT Encoder|725M|17.53|\n\n### Q3: Classier-free guidance for ImageNet-512x512 generation.\nAccording to your suggestion, we provide the results for ImageNet-512x512 generation with classier-free guidance in Table R-1 (b) in the one-page pdf. Compared with MaskDiT, our MC-DiT is consistent to reduce FID by 0.47, demonstrating its effectiveness.", "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": "ZW0hkp70hm", "reviewer": "Reviewer_Ngpd", "summary": "The paper introduces a novel training paradigm for Diffusion Transformers (DiT) in the context of generative diffusion models for image generation. The authors propose MC-DiT, which focuses on enhancing contextual information extraction by reconstructing clean unmasked patches from clean masked patches, as opposed to the traditional noisy-to-noisy reconstruction. The method employs two complementary branches of DiT decoders to balance the use of noisy and clean patches, preventing model collapse.", "questions": "Please refer to the Weakness Section.", "limitations": "The authors have mentioned the limitations in this paper.", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1.\tThe paper presents a new insight into the use of clean-to-clean reconstruction for learning contextual information in masked diffusion models, which is a significant departure from traditional noisy-to-noisy reconstruction methods.\n2.\tThe authors provide a theoretical analysis of mutual information between unmasked and masked patches, demonstrating the limitations of existing methods and the benefits of their proposed approach.\n3.\tThe introduction of two complementary DiT decoder branches to prevent model collapse is a thoughtful addition that addresses a common issue in such models.\n4.\tThe paper reports state-of-the-art results in terms of FID scores and IS scores, indicating that the proposed MC-DiT is highly competitive with existing methods.", "weaknesses": "1.\tThe proposed MC-DiT model may be more complex than necessary, which could potentially hinder its adoption and implementation in practical applications.\n2.\tThe paper acknowledges that the training and inference speed of MC-DiT needs to be improved, which suggests that the current approach may have efficiency issues. The authors should provide specific comparisons to demonstrate that these efficiency sacrifices are worth the performance gains.\n3.\tThe paper could benefit from a more detailed comparative analysis with other state-of-the-art methods, including feature visualization, to better understand the advantages of MC-DiT.\n4.\tCan the author explain whether this specific context information is pixel-wise information or semantic information? And their role in the overall framework?\n5.\tAblation experiments can be further supplemented and improved. For example, the hyperparameters in Tab.5 can be further observed to have an impact. The current scaling still has some ambiguity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Q1: Complexity.\nOur MC-DiT has the same main branch and training objective as existing methods like MaskDiT, MDT, and SD-DiT. The additional complexity of MC-DiT lies on the extra two EMA branches and unmasked tuning. \n\n1) The extra two branches increases only 7.6\\% parameters and 8\\% FLOPs, as shown in Figure R-1(a) in the one-page pdf. Thus, MC-DiT is approximately similar as MaskDiT, MDT and SD-DiT in training cost. Moreover, during inference, the two extra EMA branches are dropped without causing extra complexity, and MC-DiT has the same architecture as DiT in the inference stage.  \n2) Unmasked tuning can reduce the training-inference discrepancy, as demonstrated by MaskDiT and is adopted in our MC-DiT. However, we can remove unmasked tuning to reduce the complexity at little loss on FID. Table below shows that FID will increase by only 0.41 for MC-DiT by removing unmasked tuning.\n|Strategy|Iterations|FID|\n|:---:|:---:|:---:|\n|MaskDiT-XL/2 w unmasked tuning|1300K|12.15|\n|MC-DiT-XL/2 w unmasked tuning|1300K|7.92|\n|MC-DiT-XL/2 w/o unmasked tuning|1300K|8.33|\n\n### Q2: Training and inference speed.\nFirst of all, we have to clarify that, in the limitations, we mean to improve the inference speed of diffusion models rather than the inference speed of MC-DiT compared to existing DiT methods. In fact, MC-DiT has the same inference speed as MaskDiT, since two extra branches are removed during inference. Regarding training speed, MC-DiT is slower than MaskDiT due to the two extra EMA branches. However, MC-DiT yields an evidently lower FID scores (i.e., 1.5 for $256\\times 256$ image generation and 1.4 for FID for $512\\times512$ image generation) than MaskDiT.\n### Q3: More detailed comparative analysis\nWe further provide visualization of the feature maps extracted by MC-DiT and MaskDiT at different noise scale on CIFAR-10 in Figure R-2(b) in the one-page pdf. A larger noise variance denotes the noise with large scale. Our MC-DiT extracts proper shape for various noise scale, while the features extracted by MaskDiT are messy in the large noise scale. This further proves the motivation and effectiveness of our paper that clean-to-clean mask reconstruction promotes learning sufficient contextual information.\n### Q4: Specific context information.\nThe context information is semantic information that denotes the relationship between current patches and other patches. For example, the patches of legs are correlated with patches of the body in a cat image.  In fact, MAE employs mask-reconstruction for contextual information extraction from the clean image. In this paper, we introduce mask-reconstruction into the denoising process to extract contextual information. Specifically, the masked clean patches contain contextual information about clean patches target. Using masked clean patches to reconstruct target clean patches helps the model to understand the shape and context.\n###Q5: Ablation on hyperparameters.\nFollowing MaskDiT, we select 0.01, 0.1, and 1.0 as the scaling values of three hyperparameters and supplement various values for ablation study. Below tables evaluate various values of the three hyperparameters and we find that the best FID is still obtained when $\\lambda_1=0.1$, $\\lambda_2=0.1$, and $\\lambda_3=0.05$. \n|$\\lambda_1$|0|0.01|0.03|0.05|0.07|0.09|0.1|0.3|0.5|0.7|0.9|\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n|FID|43.23|40.99|39.23|38.44|37.95|36.53|35.20|35.98|36.74|36.91|37.52|38.97|\n\n|$\\lambda_2$|0|0.01|0.03|0.05|0.07|0.09|0.1|0.3|0.5|0.7|0.9|\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n|FID|38.83|36.15|36.02|36.46|35.99|36.07|35.20|35.34|36.18|37.26|35.98|37.54|\n\n|$\\lambda_3$|0|0.01|0.03|0.05|0.07|0.09|0.1|0.3|0.5|0.7|0.9|\n|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|\n|FID|37.77|37.25|36.63|35.20|36.07|37.93|35.46|35.88|37.26|36.19|37.40|36.35|", "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": "n9PI8VtGXF", "reviewer": "Reviewer_cs4r", "summary": "This paper introduces MC-DiT, a novel training paradigm for Diffusion Transformers (DiT) in image generation. It addresses the limitations of current masked-reconstruction strategies, which fail to effectively extract **contextual information** due to noisy-to-noisy reconstruction. MC-DiT employs clean-to-clean reconstruction, allowing for better contextual information utilization during diffusion denoising. The authors also design dual decoder branches to prevent model collapse. Theoretical and empirical analyses validate their approach, and experiments on the ImageNet dataset show that MC-DiT achieves state-of-the-art performance in both unconditional and conditional image generation tasks.", "questions": "The proposed method may still require large computational resources due to the dual-branch decoder design and the clean-to-clean reconstruction process. How to accelerate it?", "limitations": "The authors acknowledge that the training and inference speeds need improvement.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1.The introduction of the MC-DiT paradigm, which utilizes clean-to-clean mask-reconstruction, represents a novel approach that addresses the limitations of existing methods in extracting contextual information.\n\n2.The authors provide a thorough theoretical and empirical analysis, particularly focusing on mutual information, which strengthens the validity of their claims.\n\n3. The proposed MC-DiT achieves superior results in both unconditional and conditional image generation tasks, as demonstrated by the state-of-the-art FID scores on the ImageNet dataset.", "weaknesses": "1. The paper primarily focuses on image generation using the ImageNet dataset. It remains to be seen how well the approach generalizes to other domains or datasets with different characteristics.\n\n2. The authors should clearly elaborate on the differences between MC-DiT and other masked diffusion transformers (such as MaskGiT,SD-DiT, and MaskDiT).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Q1: Generalization to other domains or datasets.\nWe adopt the ImageNet dataset in the experiments for a fair comparison, since MaskDiT, SD-DiT and MDT are all evaluated on the ImageNet dataset. In fact, our MC-DiT can be generalized to different domains or datasets for improved image generation due to the fact that it can extract contextual information from arbitrary images. Table below compares the performance of MaskDiT and MC-DiT on the CIFAR-10 and CelebA (that collected for face recognition) datasets. Due to time limit, we train both MaskDiT and MC-DiT for 200K iterations. Experimental results show that MC-DiT outperforms MaskDiT on both datasets.\n|Cifar10|FID|\n|:--:|:--:|\n|MaskDiT-B / 2|11.52|\n|MC-DiT-B / 2|9.28|\n\n|CelebA|FID|\n|:--:|:--:|\n|MaskDiT-B / 2|7.14|\n|MC-DiT-B / 2|5.36|\n### Q2: Difference between MC-DiT and other masked DiTs.\nIn summary, we propose clean-to-clean reconstruction for MC-DiT achieve enhanced contextual information extraction, while existing methods (MDT, SD-DiT, and MaskDiT) are limited in contextual information extraction using noisy-to-clean and noisy-to-noisy reconstruction. Specifically, in Section 3.2 in the manuscript, we first demonstrate the problem of limited ability in contextual information extraction for previous methods (MDT, SD-DiT, and MaskDiT), since they apply noisy patches reconstruction task into the training process of DiT. They could suffer from weak ability of contextual information extraction when the noise is large. Based on this finding, our MC-DiT leverages clean patches for mask-reconstruction task to model contextual information effectively. We insert masked clean patches into the unmasked noisy features for reconstruction to exploit clean contextual information for denoising unmasked patches at arbitrary noise scales. Experimental results demonstrate that our MC-DiT can extract sufficient contextual information and achieve superior performance.\n### Q3: Computational resources and acceleration.\nWe report the resource consumption of MC-DiT, MDT, and MaskDiT in Table R-1(a) in the one-page pdf. Compared with MaskDiT, our MC-DiT sligtly increases the parameters and FLOPs by 7.6\\% and 8\\% but yields obvious FID gains of 1.5 and 1.4 for $256\\times 256$ and $512\\times 512$ image generation. MC-DiT could be potentially accelerated by employing more layers for DiT encoder but less layers for DiT decoder to decrease the overhead of DiT decoder.\n|$256\\times256$|Params|FLOPs|Memory|Speed|FID|\n|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|\n|MDT-XL/2|742M|28G|20G|1.22|6.23|\n|MaskDiT-XL/2|730M|24G|18G|3.09|5.69|\n|MC-DiT-XL/2|786M|26G|20G|1.41|4.14|\n\n|$512\\times512$|Params|FLOPs|Memory|Speed|FID|\n|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|\n|MDT-XL/2|742M|64G|28G|0.83|-|\n|MaskDiT-XL/2|730M|56G|24G|1.98|10.79|\n|MC-DiT-XL/2|786M|60G|27G|1.05|9.30|", "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": "2Ozcl8hIri", "reviewer": "Reviewer_TpTX", "summary": "This paper observes that reconstructing masked noisy patches from unmasked noisy patches harms contextual information extraction during the training of DiT and then proposes a novel training paradigm named MC-DiT with clean-to-clean mask-reconstruction. Two EMA branches of DiT decoders are designed to avoid model collapse.", "questions": "1. Please clarify the reason why two extra EMA branches can address model collapse.\n2. It is suggested to provide visual comparisons compared with other SOTA methods.\n3. Investigating the impact of classifier-free guidance is recommended since it can improve the performance of many baselines such as ADM, DiT, and MaskDiT.", "limitations": "They provide accurate limitations in the conclusion section.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The manuscript adequately puts forth a number of propositions and commendably supports these with ample evidence and rigorous demonstrations, fostering a robust intellectual foundation for their arguments.\n2. The authors' perspective on applying a noisy-to-noisy mask reconstruction approach is convincingly articulated.", "weaknesses": "1. The presentation of generated images for visualization is rather limited in quantity, necessitating an expansion to adequately illustrate the diversity and quality of the results. It is suggested to present generated results with the resolution of $512\\times 512$. This paper only provides visual results in Figure 5 with the resolution of $256\\times 256$ and it also claims superiority on $512\\times 512$ image generation.\n2. Lack of experiment details about training time, inference time and memory usage.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Q1: Visual quality comparison.  \nIn Figures~R-1 and R-2(a) in the global rebuttal file, we further provide various $256\\times 256$ and $512\\times 512$ images generated by our MC-DiT and compare with SOTA methods MaskDiT and MDT. Our generated images are more realistic and have more consistent textual structure than MaskDiT and MDT. For example, images of the 'school bus' and 'custard apple' in Figure R-1(a) exhibit different styles, while the details appear very realistic. Moreover, images of 'hammer' generated by \n MaskDiT and MDT have incomplete structure, while our MC-DiT generates images with more complete structures. The same thing happened in $512\\times 512$ images in Figure R-1(b), various images (e.g. 'dog', 'fox' and 'penguin') exhibit rich details and realistic styles, validating the effectiveness of our MC-DiT.\n### Q2: Details about training time, inference time, and memory usage. \nBelow table elaborates the implementation details for training. We leverage two types of GPUs for training and summarize the training time, inference time and memory usage. \n|Setting|MC-DiT-B/2|MC-DiT-XL/2|MC-DiT-XL/2|\n|:-:|:-:|:-:|:-:|\n|Resolution|$256\\times256$|$256\\times256$|$512\\times512$|\n|Training Time|50h|586h|623h|\n|Inference Time (50K images)|12h|8h|15.2h|\n|GPUs|2$\\times$RTX-3090 GPUs|4$\\times$V100GPUs|4$\\times$V100GPUs|\n|Batch Size|256$\\times$2|256$\\times$4|128$\\times$4|\n|Memory Usage per GPU|17GB|20GB|27GB|\n\nBesides, we compare the training cost (parameters, FLOPs, memory used and training speed) on 4$\\times$V100 GPUs in below Tables, where training speed denotes the number of iterations per second. The training speed of MC-DiT is a little bit slower than other methods due to two EMA branches. However, the inference speed of MC-DiT is similar to MaskDiT, since two EMA branches are removed during inference. The additional overhead of MC-DiT is relatively small (7.6$\\%$ parameters and $8\\%$ FLOPs), but the FID performance improvement is significant. \n|$256\\times256$|Params|FLOPs|Memory|Speed|FID|\n|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|\n|MDT-XL/2|742M|28G|20G|1.22|6.23|\n|MaskDiT-XL/2|730M|24G|18G|3.09|5.69|\n|MC-DiT-XL/2|786M|26G|20G|1.41|4.14|\n\n|$512\\times512$|Params|FLOPs|Memory|Speed|FID|\n|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|\n|MDT-XL/2|742M|64G|28G|0.83|-|\n|MaskDiT-XL/2|730M|56G|24G|1.98|10.79|\n|MC-DiT-XL/2|786M|60G|27G|1.05|9.30|\n\n### Q3: Two extra EMA branches for addressing model collapse.\nThe modal collapse occurs when the main branch only considers clean-to-clean mask-reconstruction for masked clean patches but ignores the denoising of unmasked noisy patches. We propose two EMA branches to balance the two tasks for the main branch. We use the noisy EMA branch to realize noisy-to-clean mapping for denoising, and the clean EMA branch to realize clean-to-clean mapping for mask-reconstruction (mask ratio is 0\\%). The two EMA branches constrain the output of the main branch (minimize the MSE loss between the outputs of the main branches and EMA branches) via three hyper-parameters, which leads to the balance on the denoising task and clean-to clean mask-reconstruction task. \n\nTo verify this, we report in below Table the FID score of the main branch with noisy and clean patch inputted only. The result of the main branch with unmasked noisy patches only is higher than that of masked clean patches, indicating the modal collapse problem. With noisy and clean branches, the FID score of the main branches decline distinctly, validating the effectiveness of the EMA branches.\n|Branches|FID|\n|:-------:|:-------:|\n|Min Branch|22.10|\n|w Noisy Branch|19.26|\n|w Clean Branch|18.88|\n|Main Branch (unmasked noisy patch only)|25.72|\n|Main Branch (masked clean patch only)|23.69|\n|Main Branch (unmasked noisy patch only) w Noisy Branch|19.84|\n|Main Branch (unmasked noisy patch only) w Clean Branch|19.57|\n\n### Impact of classifier-free guidance.\nWe have reported the $256\\times256$ image generation results with classifier-free guidance (CFG) in Table 1 in the manuscript, indicating the superior performance of MC-DiT. Besides, we supplement the $512\\times512$ results in the below Table. CFG benefits our MC-DiT with a district margin (4.14 vs 1.78 for $256\\times 256$ and 9.30 vs 2.03 for $512\\times 512$).\n|Methods|FID|sFID|IS|Prec.|Rec.|\n|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|\n|DiT-XL / 2|12.03|7.12|105.25|0.75|0.64|\n|MaskDiT-XL / 2|10.79|13.41|145.08|0.74|0.56|\n|MC-DiT-XL / 2|9.30|6.28|179.58|0.76|0.53|\n|DiT-XL / 2-G|3.04|5.02|240.82|0.84|0.54|\n|MaskDiT-XL / 2-G|2.50|5.10|256.27|0.83|0.56|\n|MC-DiT-XL / 2-G|2.03|4.87|272.19|0.84|0.56|", "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": "DW3loLK2qa", "reviewer": "Reviewer_gqjf", "summary": "The paper introduces MC-DiT, a training paradigm for Diffusion Transformers (DiT) in the field of generative diffusion models for image generation.  By utilizing the proposed clean-to-clean mask-reconstruction approach, the model can better leverage contextual information at different noise variances.", "questions": "See the weakness part.\nBesides, why is the IS of MC-DiT-XL / 2 much higher than other competitors?", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper provides a perspective on the limitations of noisy-to-noisy masked reconstruction, supported by theoretical insight and empirical analysis.\n- The method is overall reasonable.\n- The performance seems good.", "weaknesses": "- Will the additional two branches of DiT decoders increase the training overhead compared with other baseline methods? How about the training cost of each iteration compared with baselines?\n- Comparing with MDT-XL / 2-, the improvements of MC-DiT-XL / 2-G seem to be marginal. \n- How is the natural information measured in Fig. 1?\n- Will the code be released?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Q1: Training overhead of extra branches.\nWith the additional two branches, the training cost of MC-DiT is a little bit higher than MaskDiT and MDT. As shown in below Table, the MC-DiT-XL/2 has more parameters, since the EMA branches introduce additional 56M parameters and these additional parameters only accounts for 7.6\\%. The FLOPs and training speed of  MC-DiT-XL/2 are lower than those of MaskDiT-XL/2. The MDT-XL/2 has higher FLOPs and lower training speed than MC-DiT-XL/2 due to the difference of mask ratio (50% in MC-DiT and 30% in MDT). \nIn a word, the additional overhead of MC-DiT is relatively small (7.6$\\%$ parameters and $8\\%$ FLOPs), but the FID performance improvement is significant. \n|$256\\times256$|Params|FLOPs|Memory|Speed|FID|\n|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|\n|MDT-XL/2|742M|28G|20G|1.22|6.23|\n|MaskDiT-XL/2|730M|24G|18G|3.09|5.69|\n|MC-DiT-XL/2|786M|26G|20G|1.41|4.14|\n\n|$512\\times512$|Params|FLOPs|Memory|Speed|FID|\n|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|\n|MDT-XL/2|742M|64G|28G|0.83|-|\n|MaskDiT-XL/2|730M|56G|24G|1.98|10.79|\n|MC-DiT-XL/2|786M|60G|27G|1.05|9.30|\n\n### Q2: Improvements of MC-DiT-XL/2-G over MDT-XL/2-G.\nIn $256\\times 256$ image generation, MDT-XL/2-G and MC-DiT-XL/2-G obtain FID scores of 1.79 and 1.78, respectively. However, MDT-XL/2-G is trained for 6500K iterations, while our MC-DiT-XL/2-G only requires 2500K iterations, less than 40\\% of MDT-XL/2-G. For a fair comparison, we train MDT-XL/2 and MDT-XL/2-G under similar 2500K training iterations for evaluations and report the results in below Table. MC-DiT-XL/2 and MC-DiT-XL/2-G are shown to outperform the retrained MDT-XL/2 and MDT-XL/2-G with evident margins. This results further demonstrates that our MC-DiT can achieves superior performance with much fewer training iterations. \n|Methods|Iterations|FID|sFID|IS|Prec.|Rec.|\n|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|:-------:|\n|MDT-XL / 2|2500K|7.41|4.95|121.22|0.72|0.64|\n|MaskDiT-XL / 2|2500K|5.69|10.34|177.99|0.74|0.60|\n|SD-DiT-XL / 2 |2400K|7.21|5.17|144.68|0.72|0.61|\n|MC-DiT-XL / 2|2500K|4.14|6.96|309.69|0.83|0.62|\n|MDT-XL / 2-G|2500K|2.15|4.52|249.27|0.82|0.58|\n|MaskDiT-XL / 2-G|2500K|2.28|5.67|276.56|0.80|0.61|\n|MC-DiT-XL / 2-G|2500K|1.78|4.87|290.17|0.81|0.62|\n\n### Q3: Measure of Mutual information in Fig. 1.\nIn Fig. 1, we measure the mutual information between unmasked and masked patches in different vanilla and generated images. Specifically, given a vanilla image, we patchify it via a tokenizer and apply the mask to obtain the masked and unmasked patches. Noise at different scales is add to these patches. For 'vanilla clean and noisy images' in Fig. 1, these unmasked patches and masked patches are then reshaped to one-dimensional tensors and the mutual information is calculated based on these tensors. For 'generated images' in Fig. 1, unmasked patches are sent to corresponding models for denoising and we use output patches for calculation.\n\n### Q4: Will the code be released?\nWe have provided the source codes in the supplementary material for review. The overall source codes and pretrained models will be released upon acceptance.", "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:
We thank all reviewers for their valuable comments. We appreciate the reviewers' recognition of our work, including **excellent motivation (ZkfZ, RKJM, Ngpd, cs4r and gqjf), reasonable method (gqjf), thorough theoretical analysis (cs4r, Ngpd, YSTZ), state-of-the-art results (Ngpd, gqjf, cs4r and ZkfZ), along with clear writing (YSTZ, RKJM)**. The reviewers also raised some issues, including training efficiency (TpTX,cs4r, RKJM and ZkfZ) and unmasked tuning (ZkfZ, YSTZ). According to their comments, we have addressed each concern raised by the reviewers in a point-by-point manner.

Meta-review: The paper introduces a clean-to-clean reconstruction strategy to train DiT, which contrast with the noisy-to-noisy reconstruction that is performed in sota models. This contributes to the learning of more contextual information that benefit the reconstruction quality. Next to that, the authors propose an architecture with a 2-branch (complementary) DiT decoder to address the risk of mode collapse during learning. The novelty of the approach is evident and also supported by a theoretical analysis that strengthens the empirical results. 
The experimental analysis is comprehensive, comparing performance results with existing methods (outperforming them on several tests) and performing ablations. The rebuttal has addressed the concerns of the reviewers, and the extra results and analyses provided in the author-reviewer discussions complement the already good results reported in the paper and appendices. In my opinion, they can be added to the supplementary material and referred at in the main paper.
Decision: Accept (poster)