Title: MC-DiT: Contextual Enhancement via Clean-to-Clean Reconstruction for Masked Diffusion Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: Diffusion

Transformer

Masked Autoencoder
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.


Peer reviews:
comments: 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.

STRENGTHS:
- The paper provides a perspective on the limitations of noisy-to-noisy masked reconstruction, supported by theoretical insight and empirical analysis.
- The method is overall reasonable.
- 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?
- 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:
See the weakness part.
Besides, why is the IS of MC-DiT-XL / 2 much higher than other competitors?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

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

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.
2. 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.
2. Lack of experiment details about training time, inference time and memory usage.

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:
They provide accurate limitations in the conclusion section.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1.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.\n\nWEAKNESSES:\n1. 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).\n\nQUESTIONS:\nThe 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?\n\nLIMITATIONS:\nThe authors acknowledge that the training and inference speeds need improvement.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1.\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.\n\nWEAKNESSES:\n1.\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.\n\nQUESTIONS:\nPlease refer to the Weakness Section.\n\nLIMITATIONS:\nThe authors have mentioned the limitations in this paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-26"}

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

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.

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:
The authors mentioned that the training speed and inference speed still need to be improved and a possible future mitigation on the issue.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-28

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1.\tSufficient theoretical analysis\n2.\tThe overall writing of the paper is logically clear\n\nWEAKNESSES:\n1.\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.\n\nQUESTIONS:\n1.\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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nQuestions listed above.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-29"}