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_cc_2024_Conference
Published: 2024-04-23 04:49:49

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.

Main paper: 
 introduction: 
 Diffusion Probabilistic Models (DPMs) [19, 29, 42, 43] have emerged as front-runners in the latest advancements of image-level generative models, and surpass previous state-of-the-art generative models [10, 14] . DPMs corrupt an input image into a 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 [42, 43] . U-Net [38] is popular in early works [37, 35] to predict noise from disrupted images for image generation. Recently, Diffusion Transformer (DiT) [34] has been widely considered for DPMs [19, 29, 42, 43] in conditional image generation [4, 37] , video generation [16, 22, 30] , and 3D generation [15, 25, 36] due to its excellent scalability and superior performance.Different from vision transformers (ViTs) [11] , DiT is trained to 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 to learn contextual information in different regions of images for noise prediction [13] . To solve these problems, mask-reconstruction is introduced into the denoising-based training schedule for DiT [13, 49, 48] . Inspired by masked autoencoder (MAE) [18] , DiT is trained to predict masked noisy patches from the given unmasked noisy patches. MDT [13] pioneers to propose the asymmetrical diffusion transformer that integrates mask-reconstruction 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% 2 Related Work Diffusion Probabilistic Models. Diffusion Probabilistic Models (DPMs) [19, 42, 43] have attracted increasing attention due to their superior image generation ability compared with preceding generative models [14, 47] . 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 [43] is introduced into diffusion models to design the forward and reverse process with elegant Stochastic Differential Equation (SDE) [44] . EDM [21] analyzes the design space of diffusion models and disentangles the effects of parametrization, sampling, and training. To address the timeconsuming iterative issue inherent in DPMs, several methods apply fast sampling strategy [27, 28] or latent diffusion training strategy [37] .Transformers in Diffusion Models. In DPMs, the most commonly used architecture for noise prediction is U-Net [38] , which is a symmetric encoder-decoder framework. Recently, transformers provide a new perspective to noise prediction due to their excellent multi-modality 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 in DPMs realize superior performance and achieve the scaling law. Therefore, various works adopt and improve DiT into 2D image generation [34, 35] , video generation [30, 16, 22] , and 3D generation [15, 25, 36] . In this paper, we take the DiT as our backbone and change the input from noisy patches to clean patches. 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 I(X unmasked ; X masked ) under various noise scales. The 'vanilla clean images' and 'vanilla noisy images' denote the real clean/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] in different reconstruction targets. The left yellow ellipse denotes unmasked patches and the right green one denotes masked patches. The black arrow denotes the positive pairs to pull in.Masked Training with Transformers. Mask-reconstruction has been broadly applied into convolutional networks and transformers [11] . MAE [18] takes the mask-reconstruction 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 [1] utilizes masked training into multi-modality inputs for cross-modal fusion and generation. SdAE [8] improves masked autoencoder via self distillation. In summary, the mask-reconstruction training objective transfers the information from unmasked patches to masked patches, and thereby, enhances contextual semantic modeling ability.3 Proposed Method 
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 a 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not 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 to address problems of privacy and fairness. • 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
preliminaries: 
 Masked AutoEncoders [18] . Masked AutoEncoder (MAE) is a significant unsupervised pretraining paradigm in computer vision, which reconstructs masked patches from unmasked patches. Given an image x, MAE first patchifies it into N patches denoted by x ∈ R N ×c , where c is the channel dimension. A random binary mask m ∈ {0, 1} N is applied to obtain masked patchesx 1 = x[1-m] ∈ R N1×c and unmasked patches x 2 = x[m] ∈ R N2×c. N 1 and N 2 are the number of masked and unmasked patches. An encoder-decoder framework h = g • f is then applied. The encoder f maps the unmasked patches x 1 into latent space z 1 = f (x 1 ), while the decoder reconstructs the pixel value of masked patches x ′ 2 = g(z 1 ). The MAE is trained to ensure the reconstruction ability by minimizing the Mean Square Error (MSE) loss [46] provides a theoretical understanding of MAE and establishes connection between MAE and contrastive learning [33, 17, 7] .L M AE = E x E x1,x2|x ∥g(f (x 1 )) -x 2 ∥ 2 . U-MAEProposition 1 ( [46] ) The lower bound of the MAE loss isEQUATIONwhere L asym denotes the asymmetric alignment loss in [46] , ε is the fitting error, andh g = g • f g is the pseudo autoencoder that satisfies E x ∥h g (x 2 ) -x 2 ∥ 2 ≤ ε.U-MAE [46] combines Proposition 1 with the mask graph in Figure 1 (b) (upper) to represent the contrastive objective in MAE. Specifically, Proposition 1 demonstrates that the MAE loss is equal to the contrastive loss L asym , which calculates the similarity of h(x 1 ) and h g (x 2 ). If x 1 and x 2 are neighboring patches, L asym is minimized when positive pairs (i.e., x 1 and x 2 ) are pulled closer. Proposition 1 is consistent with the mask graph in Figure 1 (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 [19, 42, 43] . 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 x 0 ∼ P data (x 0 ) is corrupted into Gaussian noise x T ∼ N (0, σ2 max I) step by step via stochastic differential equation (SDE):EQUATIONwhere f is the drift coefficient, g is the diffusion coefficient, w is a standard Wiener process, and t is the time value from 0 to T . The reverse process generates target samples using the inverse SDE:EQUATIONwhere w is a reverse-time Wiener process. Following EDM [21] to set f (x, t) = 0 and g(t) = √ 2t, the forward and reverse SDEs are reformulated as dx = √ 2tdw and dx = -t∇ x log p t (x)dt, where s(x, t) = ∇ x log p t (x) is the score function. To solve the reverse-time SDE, a denoising network D θ (x, t) is trained to minimize the score matching loss:EQUATION)As a result, the score function is estimated by s(x, t) = (D θ (x, t) -x)/t 2 . During training, at step t, noisy images x 0 + n are sent to the denoising network D θ (x, t) to predict the clean images x 0 . However, directly optimizing this objective leads to poor contextual information [13] . To solve this problem, the mask-reconstruction is applied into denoising network [13, 48, 49] .We first review the mask-reconstruction 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 to 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-masked patches.Mutual information between input and output patches. Besides the mutual information between masked and unmasked output patches in Figure 1 (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 [49] recovers unmasked clean output patches from masked noisy input patches, as elaborated below. = E xt E x 1 t ,x 2 t |xt ∥g(f (x 1 t ), x 2 t ) -x 1 0 ∥ 2 .The ability to exploit contextual information is measured by mutual information I(x 1 0 ; x 2 t ). The contextual information in both two formulations is transferred from the noisy patches to noisy patches. 2 Subsequently, we analyze the contextual information utilization ability of maskreconstruction via calculating mutual information of masked and unmasked patches.Proposition 2 Given masked and unmasked clean patches x 1 0 and x 2 0 and their noisy versions x 1 t and x 2 t , the mutual information I(x 1 t ; x 2 t ), I(x 1 0 ; x 2 t ), and I(x 1 0 ; x 2 0 ) satisfy thatEQUATIONEQUATIONProposition 2 suggests that the mutual information I(x 1 0 ; x 2 t ) and I(x 1 t ; x 2 t ) are lower than the ideal mutual information I(x 1 0 ; x 2 0 ). 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 . Thus, the gap between I(x 1 0 ; x 2 t ), I(x 1 t ; x 2 t ) and I(x 1 0 ; x 2 0 ) becomes larger and the ability to extract contextual information is degraded. Mask graph. We further analyze the mask-reconstruction via contrastive objectives in mask graphs. In U-MAE [46] , the mask-reconstruction 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 L asym in Proposition 1. 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 1(b) illustrates the mask graph for clean image reconstruction in MAE (top) and that for noisy patch reconstruction in MaskDiT, SD-DiT, and MDT (bottom). In Proposition 3, we prove that contrastive objective between noisy patches could interfere learning contextual information.The asymmetric loss of noisy patch reconstruction and the asymmetric loss of clean patch reconstruction satisfy thatEQUATIONwhere L asym is defined in (1) and represents contextual information. The two noise-weighted items represent contrastive objective between h(x 1 t ) and [∂h g /∂x 2 0 ] (h g (x 2 0 ) and [∂h/∂x 1 0 ]) weighted by the Gaussian noise n.Proposition 3 implies that the Gaussian noise introduces two extra terms in (7) and could affect the optimization process of L asym . Noisy patch reconstruction undermines the contrastive objective of contextual information, since larger Gaussian noise leads to more severe perturbation on L asym .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 1 (a). To solve this problem, in Section 3.3, we propose MC-DiT to effectively reconstruct masked clean patches from unmasked clean patches.As demonstrated in Propositions 2 and 3 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 (a) depicts the proposed framework for MC-DiT. The clean images are disrupted by Gaussian noise n ∼ N (0, t 2 I), where t is the time step. Then the noisy images are patchified and masked by a random binary mask m. The unmasked noisy patches x 1 t are fed into the DiT encoder for feature extraction. For contextual information extraction, the masked clean patches x 2 0 are concatenated with extracted feature z = concat(z 1 , x 2 0 ), where z 1 is the feature of x 1 t . After that, the feature z is sent to DiT decoder to reconstruct clean unmasked patches x 1 0 , which is consistent with previous masked diffusion ( [48] , [49] ). The training objective of unmasked clean patch reconstruction is:L clean = E x0∼p data E n∼N (0,t 2 I) ∥(D θ ((x 0 + n) ⊙ (1 -m), x 0 ⊙ m, t) -x 0 ) ⊙ (1 -m)∥ 2 . (8)By applying masked clean patches x 2 0 in (8), the information in x 2 0 is transferred to unmasked clean output patches x1 , which is constrained to equal x 1 0 . It is not disrupted by noise n, since there is no noise in x 2 0 and x 1 0 . Moreover, as mentioned in Section 3.2, the mutual information for our MC-DiT is I(x 1 0 ; x 2 0 ), since we leverage masked clean patches to recover unmasked clean patches. According to Proposition 2, I(x 1 0 ; x 2 0 ) is much higher than I(x 1 0 ; x 2 t ) and I(x 1 t ; x 2 t ) under different noise n and time steps t. Thus, the mutual information learned by our MC-DiT would not decrease for different noise at different time steps and sufficient contextual information could be learned for reconstruction.In addition, we explore the potential benefits of x 2 0 by enforcing the masked output patches x2 to match x 2 t . As discussed in Section 3.1, the denoising network D θ predicts clean images from noisy images and can be viewed as distinguishing the clean images and noise. Therefore, predicting x 2 t from x 2 0 is to recognize the noise n and add to x 2 0 . To this end, we employ reversed constraint on masked output patches x2 , as illustrated in Figure 2 (a). The training objective isL noise = E x0∼p data E n∼N (0,t 2 I) ∥(D θ ((x 0 + n) ⊙ (1 -m), x 0 ⊙ m, t) -(x 0 + n)) ⊙ m∥ 2 . (9)Although MC-DiT can 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 decoders3 to balance the mask-reconstruction and denoising objective. As shown in Figure 2(a) , given the noisy input to DiT encoder, we introduce two branches to achieve only mask-reconstruction 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 D φ . Constraints on the two branches are employed in the loss function to balance the DiT decoder. Here, D ϕ and D φ are two EMA DiT decoders. Therefore, the overall loss is:EQUATION)EQUATION)EQUATIONwhere λ 1 , λ 2 and λ 3 are hyper-parameters.Unmasking Finetuning. Similar to MaskDiT [48] , although our 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-inference discrepancy. The clean patches used in training are not provided in the inference time. Thus, after training our MC-DiT, we finetune it on the unmasked scenarios for better performance, as 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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer 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 as well as 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 it into the paper). 
implementation details: 
 Model Settings. We follow the same architecture in MaskDiT [48] . Specifically, we first apply a pretrained variational autoencoder (VAE) from Stable Diffuion [37] to map the images into latent space, and then train our MC-DiT to reconstruct clean patches from noisy patches under the EDM [21] framework to approximate score function in the diffusion process. The pretrained VAE maps 256×256×3 input images to 32×32×4 latent features and 512×512×3 images to 64×64×4 latent features. Similar to SD-DiT, we apply DiT-S, DiT-B, and DiT-XL as our backbones.Training Details. Similar to previous works [13, 48, 49] , we train MC-DiT on ImageNet [39] with resolutions 256×256×3 and 512×512×3, respectively. Most training settings are the same as MaskDiT [48] . We train MC-DiT for 400K to 1M iterations using the AdamW optimizer with learning rate 0.0001 and no weight decay. By default, we use 50% mask ratio and batch size 1024. λ 1 and λ 2 in ( 12) are set to 0.1 and 0.05 for more denoising reconstruction. The EMA coefficient is set to 0.999 for smoothness and no data augmentation is employed.Evaluation Metrics. Following DiT [34] , we leverage Fréchet Inception Distance (FID) to measure the quality of generated images. For fair comparison, we also use ADM's Tensorflow evaluation suite [9] to compute FID-50K (FID for short), sFID [31] , Inception Score (IS) [40] and Precision/Recall [24] as secondary metrics. More vivid images have lower FID and sFID, while their IS and Precision/Recall are higher.We evaluate vanilla training (i.e., LDM [37] , ADM [9] , and DiT [34] ) and masked training (i.e., proposed MC-Dit, MaskDiT [48] , MDT [13] , and SD-DiT [49] ) using backbones of different scales for 256×256 and 512×512 image generation on ImageNet.Results on ImageNet-256×256. Results on ImageNet-512×512. Table 2 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.Contextual Enhancement. Figure 1 (a) reports the mutual information between unmasked and masked output patches with different noise, which can be viewed as the metric of contextual information consistency. Our MC-DiT decreases slowly during the noise variance becomes larger, which indicates more sufficient contextual reconstruction regardless noise. Backbones at different scales. Table 3 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 [13] , and SD-DiT [49] in FID by a 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 in Figure 3(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 of our MC-DiT trained for evaluations, which is trained for 400K iterations. The training loss of MC-DiT decreases faster than DiT [34] and MaskDiT [48] . Figure 3 (b) 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 [13] ,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 of our MC-DiT.For computation efficiency, we adopt '-B' in all the models for fair comparison. All the models are trained for 400K itreations with batch size 256 and mask ratio 50%. Effectiveness of two extra EMA branches. Table 5 evaluates the influence of two 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 in appendix.Unmasked tuning. Unmasked tuning (UT) 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 7 shows that FID will increase by only 0.41 for MC-DiT by removing unmasked tuning. Hyperparameters. We separately evaluate four values for λ 1 , λ 2 , and λ 3 . Table 4 shows best FID is obtained when λ 1 = 0.1, λ 2 = 0.1, and λ 3 = 0.05. Note that λ 2 is larger than λ 3 since the denoising objective is more important than contextual information utilization.Mask ratio. Figure 5 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%. 
conclusion: 
 In Limitations. Despite excellent performance, the training speed and inference speed of 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. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [Yes]Justification: We briefly discuss this in conclusion.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a best faith effort.12. Licenses 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 use 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 a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. 
 
 Appendix: A.1 Theoretical Proof Proposition 2. Given masked and unmasked clean patches x 1 0 and x 2 0 and their noisy versions x 1 t and x 2 t , the mutual information I(x 1 t ; x 2 t ), I(x 1 0 ; x 2 t ), and I(x 1 0 ; x 2 0 ) satisfy thatEQUATIONEQUATIONProof. Given the time step t,masked noisy patches x 2 t , masked clean patches x 2 0 , clean unmasked patches x 1 0 and noisy unmasked patches x 1 t , where x 2 t = x 2 0 + n, x 1 t = x 1 0 + n and n ∼ N (0, t 2 I). we derive the mutual information I(x 1 0 ; x 2 t ) according to the definition.EQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONThus, the mutual information between noisy masked patches and unmasked clean patches I(x 1 0 ; x 2 t ) is less than I(x 1 0 ; x 2 0 ) due to the non-negativity of KL divergence. Moreover, during t increases, the variance of the Gaussian noise n becomes larger. As a result, noisy masked patches x 2 t are disrupted heavily from clean patches x 2 0 . The distribution p(x 1 0 |x 2 t ) is very dissimilar from p(x 1 0 |x 2 0 ). Formally, the derivation of KL divergency can be written as:EQUATIONWe approximate p(x 2 t ) ≈ p(x 2 0 ) + p(n), since p(x 2 t ) is a Gaussion distribution with mean value x 2 0 and variance t 2 . As t increases, the KL divergence KL(p(x 1 0 |x 2 0 )||p(x 1 0 |x 2 t )) increases and the mutual information I(x 1 0 ; x 2 t ) achieves the larger difference with I(x 1 0 ; x 2 0 ). Thus, the mutual information I(x 1 0 ; x 2 t ) is lower than I(x 1 0 ; x 2 0 ). Similarly, the mutual information between noisy patches I(x 1 t ; x 2 t ) can be drived according to Eq. 20:EQUATIONEQUATIONTherefore, the proposition has been proved.Proposition 3. The asymmetric loss of noisy patch reconstruction and the asymmetric loss of clean patch reconstruction satisfy that:EQUATIONwhere L asym is defined in (1) and represents contextual information. The two noise-weighted items represent contrastive objective between h(x 1 t ) -[∂h g /∂x 2 0 ] and h g (x 2 0 ) -[∂h/∂x 1 0 ] weighted by the Gaussian noise n.Proof. According to Eq.1, the asymmetric loss can be written as:EQUATIONEQUATIONEQUATION)where we leverage first-order Taylor's formula in Eq. 26 and Eq. 28 to calculate h(x 1 t ) and h g (x 2 t ) at x 1 0 and x 2 0 , since x 1 t = x 1 0 + n and x 2 t = x 2 0 + n. o denotes the the higher order infinitesimal.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-sampler [21] . Specifically, the time steps in the reverse process are set via t i = (t1 ρ max + i N -1 (t 1 ρ min -t 1 ρ max )) ρ, where N = 40, ρ = 7, t max = 80 and t min = 0.002. Besides, the second-order correction is applied and the generated images are the average of first-order and second-order results.Training Details. We follow the LDM [37] and adopt a pretrained VAE to firstly map the images into the latent spaces. The weight of pretrained VAE 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 Table 8 .Generalization Experiments. We adopt the ImageNet dataset in the experiments for a fair comparison, since MaskDiT [48] , SD-DiT [49] and MDT [13] are all evaluated on the ImageNet dataset [39] . 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 9 compares the performance of MaskDiT [48] and MC-DiT on the CIFAR-10 [23] and CelebA [26] (that collected for face anti-spoofing) datasets. Due to time limit, we train both MaskDiT [48] 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 of 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 a 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 L clean . Thus, lower MSE loss means the generated clean patches are 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-toclean 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.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 distinctly, validating the effectiveness of the EMA branches.Ablation Study of hyperparameters. Following MaskDiT [48] , we select 0.01, 0.1, and 1.0 as the scaling values of three hyperparameters and supplement various values for ablation study. Table 13 , Table 14 and Table 15 evaluate various values of the three hyperparameters and we find that the best FID is still obtained when λ 1 = 0.1, λ 2 = 0.1, and λ 3 = 0.05. [23] . 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.Training cost comparison. We compare the training cost (parameters, FLOPs, memory used and training speed) on 4×V100 GPUs in Table 11 . 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.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, as 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 the input of EMA DiT decoder does not lead to poor denoising results. As shown in Table 12 , 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. 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 be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.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 a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Justification: There is no societal impact of the work performed.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).Although the reconstruction targets of SD-DiT are clean patches, it is equivalent to distinguish the x 1 0 and noise n. Therefore, the contextual information is used for better prediction of n.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: θema = α × θema + (1 -α) × θ dec , where α denotes the weight coefficient. The two decoders are updated only using the EMA method without using gradient updates. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - While reading the article, there are many questions that arise which effect the reading experience of the article.
- 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.
- Some experiment details are unclear (in questions).
- 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.
- 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.
questions: 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?
- 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.
- Line 216: how are the 2 branches of DiT trained in the EMA fashion here (student teacher or are do they also collect gradients)?
- Line 203: "$\mathcal{I}(x^1_0; x^2_0)$ is much higher", the **much** part is not clear from Proposition 2.
- 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?
- 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.
- Figure 3 a and b, why is the plot only shown for different number of total iterations?
- Figure 3b, are these metrics calculated before or after fine-tuning for your method?
- Line 309 in Limitations. Why does the inference speed need to be improved? Is the model inferred differently and requires more steps?

*Minor Typos*
- Line 13: MDT mentioned before it is defined, although this clear from the citation
- 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.
- Proposition 2, Equation 6 should end with a full stop.
- Equation 8 9, 10, 11: brackets are not matched, duplicate sencond closing brackets?
limitations: limitations: Questions listed above.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1.	There 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.
2.	Visualization of experimental results is indeed missing, and only quantitative experimental results exist in the body of the paper.
3.	Using 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.
questions: questions: 1.	Why not release the qualitative results as proof of the effectiveness of the strategy?
2.	How 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: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: Please see the list of questions for improvement and clarification on some of the aspects.
questions: 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. 
2. 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.
3. 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: limitations: The authors mentioned that the training speed and inference speed still need to be improved and a possible future mitigation on the issue.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1.	The proposed MC-DiT model may be more complex than necessary, which could potentially hinder its adoption and implementation in practical applications.
2.	The 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.
3.	The 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.
4.	Can the author explain whether this specific context information is pixel-wise information or semantic information? And their role in the overall framework?
5.	Ablation 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.
questions: questions: Please refer to the Weakness Section.
limitations: limitations: The authors have mentioned the limitations in this paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. The authors should clearly elaborate on the differences between MC-DiT and other masked diffusion transformers (such as MaskGiT,SD-DiT, and MaskDiT).
questions: 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: limitations: The authors acknowledge that the training and inference speeds need improvement.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. Lack of experiment details about training time, inference time and memory usage.
questions: questions: 1. Please clarify the reason why two extra EMA branches can address model collapse.
2. It is suggested to provide visual comparisons compared with other SOTA methods.
3. 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: limitations: They provide accurate limitations in the conclusion section.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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?
- Comparing with MDT-XL / 2-, the improvements of MC-DiT-XL / 2-G seem to be marginal. 
- How is the natural information measured in Fig. 1?
- Will the code be released?
questions: questions: See the weakness part.
Besides, why is the IS of MC-DiT-XL / 2 much higher than other competitors?
limitations: limitations: N/A

Comments: comment: Thank you very much for your comprehensive evaluation and the time you dedicated to reevaluating our work! We will include the experiments and visualization provided in the rebuttal in the final version of our paper. Many thanks for your suggestions again!
Thank you for the author's response. I suggest including the visualization results and ablation experiments in the final version of the paper. I have raised my score accordingly.
Thank you very much for your careful and insightful feedback, which has greatly improved the paper. We greatly appreciate your insights and are pleased to hear that all of your concerns have been addressed. Many thanks for your comments again.
Thank you very much for recognizing the improvements made to the paper and re-assessing the rating. We sincerely appreciate your thoughtful evaluation and the time you dedicated to reevaluating our work! We will definitely include the results provided in the rebuttal in the final version of our paper.
Thank you for your rebuttal. Most of my concerns are addressed. I think the experiments shown in the original paper are not sufficient and enough and it would be better to include the results in the rebuttal. I will change the rating.
Thank you very much for the time and effort you have dedicated to reviewing our manuscript. We believe we have addressed all the concerns you raised in your review and expect your feedback sincerely. Thank you once again for your attention.
Thank you very much for the time and effort you have dedicated to reviewing our manuscript.  We believe we have addressed all the concerns you raised in your review and expect your feedback sincerely.
Thank you once again for your attention.
Thank you very much for acknowledging the improvements made to the paper and reconsidering the score! We greatly appreciate your thoughtful evaluation and the time you took to reassess our work.
Thank you for your thoughtful feedback and for taking the time to review our paper. We are glad that our rebuttal has addressed all of your concerns. Many thanks for your comments again.
Thanks for your reply. I do not have further concerns.
Thank you for the clarifications and my questions have been sufficiently answered. I have updated my recommendation for the paper now.
Thanks for the rebuttal, and I would like to keep my original rating.