MC-DiT: Contextual Enhancement via Clean-to-Clean Reconstruction for Masked Diffusion Models

Q1: Visual quality comparison.

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

Q2: Details about training time, inference time, and memory usage.

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

Besides, 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.

Q3: Two extra EMA branches for addressing model collapse.

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

Impact of classifier-free guidance.

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

Q1: Training overhead of extra branches.

With 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). In a word, the additional overhead of MC-DiT is relatively small (7.6$\%$ parameters and $8\%$ FLOPs), but the FID performance improvement is significant.

Q2: Improvements of MC-DiT-XL/2-G over MDT-XL/2-G.

In $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.

Q3: Measure of Mutual information in Fig. 1.

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

Q4: Will the code be released?

We have provided the source codes in the supplementary material for review. The overall source codes and pretrained models will be released upon acceptance.

Q1: Generalization to other domains or datasets.

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

Q2: Difference between MC-DiT and other masked DiTs.

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

Q3: Computational resources and acceleration.

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

Q1: Complexity.

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

1. 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.
2. 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. |Strategy|Iterations|FID| |:---:|:---:|:---:| |MaskDiT-XL/2 w unmasked tuning|1300K|12.15| |MC-DiT-XL/2 w unmasked tuning|1300K|7.92| |MC-DiT-XL/2 w/o unmasked tuning|1300K|8.33|

Q2: Training and inference speed.

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

Q3: More detailed comparative analysis

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

Q4: Specific context information.

The 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. ###Q5: Ablation on hyperparameters. Following 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$.

Q1: Motivation for MC-DiT

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

Q2: Denotations on Figure 2.

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

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

Q3: Classier-free guidance for ImageNet-512x512 generation.

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

Q1:Claim in Line 38

We 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. Proposition 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.

Q2:Equations(5) and (6) and Claim in Line 163.

The details of the KL divergence in the Proposition 2 are as follows

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

Q3:Training of the 2 branches of DiT

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. $\theta_{ema}=\alpha \times \theta_{ema}+(1-\alpha)\times \theta_{dec}$ where $\alpha$ denotes the weight coefficient. The two decoders are updated only using the EMA method without using gradient updates.

Q4:Much part of $\mathcal{I}(x_0^1;x_0^2)$

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

Q5:Training loss in Figure 3

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

Q6:Different numbers of total iterations in Figures 3a and b.

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

Q7:Metrics in Figure 3b.

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

Q8:Inference speed in limitations

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

Q9:Evaluation of speed of convergence.

We 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. The 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.

Q10:Speed of convergence shown in Figure 3(a).

Although 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. Besides, 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.

Thank you for the clarifications and my questions have been sufficiently answered. I have updated my recommendation for the paper now.

Q1: Writing error in line 107.

We 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]$.

Q2: Visualization results.

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

Q3: Parameter tuning required for testing.

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

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.

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.