DiMSUM: Diffusion Mamba - A Scalable and Unified Spatial-Frequency Method for Image Generation

1. Training longer looks like getting an overfitting curve, an analysis needs to be done. This behavior is contradict to Transformer-based diffusion DiT, MDT, MaskDiT, etc.

Thank you for pointing out, we confirm that our method does have overfitting. After converging to the best FID10K score of 5.49 at epoch 250, the model begins to exhibit an overfitting pattern. We also observe the same overfitting problem for transformer-based diffusion models like DiT and MDT. Figure 3 in the rebuttal file demonstrates that overfitting does occur in transformer-based models, albeit at a later stage, starting from epoch 800 onward, due to their slower convergence rate compared to our model.

We hypothesize that overfitting is not solely determined by model architecture but by various factors. For instance, the Improved DDPM [2] paper showed that with identical architecture in Appendix F, changing the noise scheduling from linear to cosine could lead to overfitting in same UNET architecture. Based on our empirical experiments above, we think that the overfitting effect is not only attributable to our architecture but rather a broader phenomenon in diffusion models. For MDT, overfitting emerges much later, around epoch 1100, due to its masking strategy, which could serves as a form of regularization to delay the onset of overfitting.

2. The paper claims the issue with quadratic computation of transformers-based diffusion but there is no comparison between the proposed method and the competitor for both 256 and 512 resolution.

DIMSUM-L/2 achieves 2.2 seconds latency compared to DIT-L/2 with 3.8 seconds (see table 2 in rebuttal file).

3. The scalability is important for this architecture, but it is missing, which would weaken the paper.

For this question, we provide our answers in the global response.

4. On ImageNet, the crucial benchmark, the reported performance of the proposed method is not clearly better than DiT, and of course, cannot be comparable with the more advanced method transformer-based model such as MDT. This indicates that the proposed method seems still under-discovered, not better than the existing methods in the performance of image generation.

To answer this question, we break it down into two parts: (1) not clearly outperforming DiT and (2) underperforming compared to MDT.

First, it is important to highlight that our method achieves competitive performance with DiT and DIFFUSSM-XL, requiring significantly fewer training epochs—specifically, 4x fewer than DiT and 2x fewer than DIFFUSSM-XL. Notably, our method also uses a smaller network size of 460M parameters, compared to 675M of DiT while still demonstrating strong generation capacity and faster training convergence. In the updated ImageNet table provided in the PDF, we show that our method can further surpass the performance of other diffusion baselines when trained for a similar training duration as DIFFUSSM-XL. Notably, our training iterations are still less than a third of those required by DiT and SiT, yielding the best FID score of 2.11.

Though our method does not outperforms the results of MDT yet, we believe that MDT is orthogonal to our proposed architecture, where a masking scheme is introduced to enhance further the contextual learning ability of diffusion models, including our method. It is an interesting topic to combine with our network for boosting model performance as mentioned in the limitation of global response.

5. Q1. How is performance if using the same sampling method as other competitors that are based on DDIM/DDPM, instead of ODE solver? This is to ensure a fair comparison.

To clarify, Euler solver and DDIM are the same, except that Euler operates on continuous time intervals and removes the additional stochastic noise in each denoising step. In the manuscript, we followed SiT, Zigma, LFM works on using adaptive solver like "dopri5" for evaluation. In table 5 of the attached PDF, we show that FID scores of adaptive ODE solver and Euler solver (or DDIM sampling) are similar with small numerical difference, reconfirming the fairness of the benchmarking with DiT.

6. Q2. How is the performance w.r.t more sampling steps? A curve is a good demonstration.

We plot the FID-10K scores of various NFE used for evaluation in Fig. of the attached pdf. This shows that increasing NFE beyond 250 leads to minimal or no improvement in the FID scores. This behavior is consistent with the observation of the flow-matching in Figure 7 of FM paper, which has been shown to require fewer NFEs compared to other SDE-based methods.

[1] Lipman, Yaron, et al. "Flow matching for generative modeling." ICLR 2023.

[2] Nichol, Alexander Quinn, and Prafulla Dhariwal. "Improved denoising diffusion probabilistic models." International conference on machine learning. PMLR, 2021.

This section was included to the global response above, however, we include it here for easy reference.

We'd like to clarify that when we refer to DiT-L/2 in the rebuttal, we specifically mean the DiT-L/2 model trained using the flow matching framework like SiT. This choice ensures a fair comparison of inference speed.

This comparison is more equitable because our DiMSUM model also employs the flow matching framework. Crucially, the sampling methods between traditional diffusion methods and flow matching methods differ significantly. By comparing DiMSUM to the flow matching-trained version of DiT (SiT), we ensure that both models use the same sampling approach during inference.

By aligning sampling methods, we can accurately evaluate inference speed differences, ensuring that any variations are due to architectural choices rather than differences in sampling or training processes.

We also sorry for the typo in table 2 of the rebuttal PDF, the parameters of DiMSUM-L/2 should be 460M, not 480M.

Thank you for your response! We provide answers below:

Q1. 512 latency

256 (latent size: $32 \times 32$)

512 (latent size: $64 \times 64$)

*Note: Previously, we measured our models's latency on two resolutions using the same device and GPU_ID (a single NVIDIA A100). However, for the result of our model upon 512x512, due to resource matter (occupied device), we used a different NVIDIA A100 (different device/GPU_ID), which could cause unfairness. We remesured using the same device/GPU_ID (same specs) for a fairer comparison.

We appreciate your suggestion and have conducted additional benchmarks for 512x512 resolution images, following the same methodology used for 256x256 images in Table 2 of the rebuttal file. Our benchmarking process involves first warming up the model with 10 forward passes, then generating a single image (batch size = 1) 100 times, and finally calculating the average latency from these 100 runs.

Note that the parameter change after changing image size is mainly due to the PatchEmbed layer of the architecture, which both models have.

Q2. MEM and GFLOPs

256 (latent size: $32 \times 32$)

512 (latent size: $64 \times 64$)

Given our time constraints, we focused our comparison on DiT-L/2, maintaining consistency with our latency comparison. The results, as shown in the table, reveal that DiMSUM-L/2's memory usage is slightly higher than its counterpart. This increase is expected, considering DiMSUM's slightly larger parameter count.

Regarding GFlops, we acknowledge that for 256x256 images, DiMSUM produces about 4% more GFlops than DiT. However, an interesting trend emerges when we examine 512x512 images: DiMSUM's GFlops scaling is actually slower than DiT's. Consequently, at this higher resolution, DiT's GFlops exceed DiMSUM's by approximately 7%. This observation again aligns with the question 1 of reviewer RLU7.

This observation aligns with the known quadratic complexity of transformers as sequence length increases. Our hybrid model mitigates this issue; the impact of attention blocks is reduced, while Mamba demonstrates its linear scaling complexity as the token count grows. This architectural choice allows DiMSUM to maintain efficiency at higher resolutions, offsetting the initial GFlops difference at 256 resolution.

Q3-Q4.

Thank you for asking the questions.

In table 2, we use dopri5 to sample image from both models. We hypothesis that our architecture converges to less ODE-curvature solution compared to DiT architecture [1], [2]. This means that we could use less NFE to achieve the better image quality compared to DiT architecture. It is noted that when using dopri5, each sampling process could require different NFEs depending on the initial noise.

Additionally, we acknowledge that our method requires more preprocessing before each MambaBlock, which may introduce latency to the inference speed. Furthermore, as shown in Table 2b, adding more Wavelet level processing costs only 0.01 GFLOPs. We also confirm that all preprocessing components are included in our measurements.

[1] Pooladian, Aram-Alexandre, et al. "Multisample flow matching: Straightening flows with minibatch couplings." arXiv preprint arXiv:2304.14772 (2023).

[2] Lee, Sangyun, Beomsu Kim, and Jong Chul Ye. "Minimizing trajectory curvature of ode-based generative models." International Conference on Machine Learning. PMLR, 2023.

Q5.

To our knowledge, our method achieves state-of-the-art performance on several datasets, including CelebA and LSUN-Church. It's worth noting that while MDT and MaskDIT focused their benchmarks solely on ImageNet, our evaluation is more comprehensive, comparing against both transformer-based models and Mamba-based diffusion models such as DIFFUSSM and ZigMa.

Regarding ImageNet, we acknowledge that our initial claim may cause confusions. We will revise our statement to more accurately reflect our findings: "Our method demonstrates superior results compared to DiT and DIFFUSSM, achieving faster training convergence and delivering high-quality outputs."

It's important to highlight that both MaskDiT and MDT report their results using significantly larger models (-XL/2, approximately 675M parameters), while our results are based on a smaller model (-L/2, 460M parameters). Despite this size difference, our model achieves competitive performance. In fact, DiMSUM-L/2-G attains a slightly better FID score (2.11) than MaskDiT-G (2.28).

We acknowledge your concerns as they are critical for the paper’s clarity and improvement. We will revise them in our manuscript.

1. In Line 106, I'm unsure why sigma brings extra computation when the scan path is larger, given that the computation is always fixed.

2. In Line105, it's important to note that flow matching has been utilized in various domains to capture the reader's interest. E.g., boosting diffusion[1], depth estimation[2], motiom[3], even text generation[4].

3. In Line183, One paper is missed.

For Q1,2,3, we admit that our wording is incorrect in this context and revise the sentence as "In this paper, we show that too many scanning orders, e.g., sweep-8 and zigzag-8, may introduce excessive information and lead to worse performance compared to sweep-4". We will cite the suggested papers.

4. In Line 176, I'm curious as to why only B, C, and delta are time-dependent. Why isn't A?

We believe that the authors of Mamba would have a much better answer here. From our point of view, this implementation is made purely to prefer simplicity without sacrificing much performance. For more details, reviewers can check the interpretation in 3.5.2 of the Mamba paper. Simply put, time-dependent delta $\Delta_t$ is sufficient to ensure the selection mechanism for $\bar{A}$ via $\exp(A * \Delta_t)$.

5. In Fig 3a, what does the notation on ZigMa mean?

We are sorry for the missing information; we will update the caption in the revised version. In this case, $\textdagger$ is our reproduced result based on Zigma's official code, and $\textdaggerdbl$ is an adopted result from LFM paper.

6. Regarding line 261, what's the rationale behind sharing those parameters of transformers? Are there any previous works indicating its effectiveness?

As mentioned in L260, we drew inspiration from Zamba for the globally shared attention block. The primary purpose of this component is to reduce the model's parameter count while maintaining the competitive performance of the model by capturing the global dependencies between input tokens. Besides, Zigma also adopts a hybrid mamba-transformer architecture in the design to integrate the text condition for text-to-image generation effectively.

7. In Table 1, it would be better to display training steps. Researchers of generative models tend to compare using training steps rather than epochs.

We acknowledge your comment and update Table 1 in the attached file.

8. In Line342, for sweep4, considering four forwards in a single-forward usually leads to a worse training speed (iter/s) and increased memory usage. It would be interesting to compare these two aspects with other methods.

Actually, in our implementation, we don't run all scans in one forward (such as in VMamba), but we run each order alternatively following Mamba-ND and Zigma paper. Specifically, each Mamba block handles only one direction, and the four scanning orders in sweep4 are evenly distributed throughout the network's depth (number of layers). Sweep4, for example, has four scanning orders: (1) left-to-right, (2) right-to-left, (3) top-to-bottom, and (4) bottom-to-top. Assuming the network has only four layers, layer 1 handles (1), layer 2 handles (2), and so on.

9. In Fig5, is there any reference to the jpeg-8 scan path? How does it defined?

As mentioned in L330, we derived the JPEG-8 scanning order from the scanning order of the JPEG compression algorithm, as described in "A fast JPEG image compression algorithm based on DCT". This also was significantly influenced by the Zigma paper, a pioneering work that put effort into a more sophisticated scanning order and achieved competitive results. Building upon their insights, we extended the concept to the original JPEG scanning orders. We acknowledge the valuable contribution of the Zigma paper and will provide a more detailed explanation of our scanning order development in the camera-ready version for improved clarity.

10. The previously mentioned figure or table should be placed first.

As mentioned in the global response, we thank the reviewer for your attention to detail, which is extremely vital to enhance the clarity and completeness of the paper. We will try our best to reorder the figures and tables so that it would appear before being mentionned.

1. Since the authors argue the incorporation of frequency information can mitigate the problem caused by the scanning order, I urge the authors to conduct an experiment to verify whether the generation performance of DiMSUM model is sensitive to the scanning order.

To substantiate our claims regarding the efficiency of frequency information, we conducted a comprehensive ablation study. This study utilized four scanning orders: (1) bidirection, (2) jpeg-8, (3) sweep-8, and (4) zigzag-8. We trained models on CelebA-HQ at 256x256 resolution for 250 epochs, comparing performance when applying these scanning strategies to: a) spatial domain only or b) both spatial and frequency domains

Table 3 of rebuttal file presents the results of this ablation study. Notably, integrating frequency domain information across all four scanning strategies led to significant performance improvements. This consistent enhancement across various scanning methods provides strong evidence for the effectiveness of our approach in leveraging frequency information.

2. Stronger baselines (such as Pixart-alpha) and the experiments on more datasets (such as conditional generation on MSCOCO) and resolutions are required.

We appreciate the suggestion to explore text-to-image (T2I) generation. However, training competitive text-to-image models presents significant challenges:

• Computational constraints: As a small lab, we lack the extensive resources required for training large-scale T2I models.
• Dataset availability: Recent closure of datasets like LAION-5B limits access to suitable training data for T2I tasks and also would be hard to compare fairly with previous works using those datasets.

It's worth noting that focusing on label-to-image tasks, as we have done, is common in diffusion model research like DiT, EDM, and Consistency Models. Many influential papers in this field have established their contributions through experiments on label-to-image before expanding to text-to-image.

3. From Table 2(a), it seems that the Wavalet transformations would harm the performance (i.e., FID and recall), which cannot demonstrate the effectiveness of the proposed method.

We hypothesize that spatial and frequency signals are not aligned and require careful integration to leverage their information. Naively fusing these domains (e.g., by concatenation) can damage performance due to conflicting or misaligned information.

To address this challenge, we proposed a more sophisticated fusion method using Cross-Attention layers between these spaces. This approach enables the model to effectively combine information from both domains, leveraging their strengths while mitigating potential conflicts. Hence, this fusion technique can enhance the FID from 5.87 to 4.92 in Table 2c submission, contributing to the SoTA result of our method.

1. Our architecture might comprise many components such as Mamba, Transformer, and frequency processing. However, each proposed component is well-motivated and vital to the overall framework, as shown in Table 2a (in submission). The simple Mamba network, without frequency processing and transformer blocks, cannot surpass the performance of the transformer. As shown in section 2.1 submission, the scanning order of the Mamba block has a significant effect. To mitigate its impact, we integrate wavelet frequency with window scanning into architecture. As explained in the global response, wavelet scanning, along with window scanning, can better capture the global information of each wavelet subband, enhancing the overall global information across the whole image. In addition, our model could benefit from learning both high and low frequency [1]. With wavelet frequency and window scanning, our framework reduces FID from 5.27 to 4.92 for CelebA-HQ (Table 2e in submission). Furthermore, recently, hybrid architectures of Transformer and Mamba [2, 3] have demonstrated SoTA performance in NLP tasks. Motivated by these works, we measure the performances of hybrid Transformer-Mamba in vision generative. As shown in Table 2f submission, both independent and shared transformer improves performance, indicating hybrid model's effectiveness.

[1]: Hao Phung, Quan Dao, and Anh Tran. "Wavelet diffusion models are fast and scalable image generators." In CVPR 2023.

[2]: Zyphra unveils zamba: A compact 7b ssm hybrid model, 2024. https://www.zyphra.com/zamba

[3]: O. Lieber et al. Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887, 2024

2. We acknowledge that the presentation of our paper should be clearer and appreciate your feedback. To clarify, the term "manually-defined scanning orders" refers to the heuristic scanning order in the Mamba blocks. In Section 2.1, from lines 94 to 107 (in submission), the scanning order of Mamba is defined as hand-crafted in previous works. Throughout the paper, we mentioned some of these hand-crafted orders, including "bi", "window", "sweep-4/8", and "zigzag-8". These scanning orders are illustrated in Figs 2, 5, and 6 in submission. Although our approach does not completely address the impact of scanning order, our window scanning strategy in wavelet space could efficiently capture the global information within each frequency sub-bands which improves global content of generated image, as explained in global response. Please refer to table 3 (in rebuttal file), our wavelet and window scan component successfully reduce FID score compared to spatial scanning only. Therefore, we consider our frequency scanning method to be synergistic with image space scanning, as it effectively handles the global information of generated image. Besides, completely resolving the limitations of manually-defined scanning is an interesting topic for future research.

3. DIMSUM-L/2 achieves 2.2 seconds latency compared to DIT-L/2 with 3.8 seconds (see table 2 in rebuttal file).

4. In terms of model's type, we would like to revise the terms "Mamba-based" and "Mamba model" to "a hybrid Mamba-Transformer network" for better comprehension. In terms of spatial-frequency scan for transformer, we modify DiT architecture to further incorporate wavelet frequency information similar to DIMSUM. After training spatial-frequency transformer on CelebA-HQ for 1000 epoch, we found that the spatial-frequency transformer failed to learn the human face distribution. The output model only generates noise-pattern images. More rigorous analysis of model is beyond the scope of paper and we would like to leave it for future investigation.

5. We acknowledge that the use of term 'scalable' may have caused some confusion. In the broader context, deep learning’s scalability may refer to capacity to handle various datasets and computational resources while producing correct results in an acceptable period of time [1].

Our model aligns well with this definition of scalability. Our experiments demonstrated that:

• The model adapts to multiple datasets with minimal hyperparameter tuning (similar to DiT paper).
• It achieves competitive performance metrics compared to other methods.
• Inference speed is also faster, as shown in table 2 (in rebuttal file).

Our SoTA results are achieved with a parameter count comparable (or even smaller) to existing models. Refer to Figure 3a and 4a (in submission) and Table 1 (in rebuttal file). This suggests substantial room for further enlargement of our model's parameters, which we anticipate will yield even better improvements across various and bigger datasets (see table 4 in rebuttal file).

6. Our plot intentionally stops at epoch 300, demonstrating the model's capability to converge faster than other methods. DiT does perform well on CelebA-HQ but takes more than 500. Figure 4 (in rebuttal file) illustrates the convergence of DiT at the later epoch. It's noted that DiT and LFM (in figure 3d submission) use the same DiT-L/2 architecture. While DiT uses diffusion loss, LFM uses flow matching. LFM converge faster than DiT. However, our model with flow matching loss, demonstrates even faster convergence than LFM, indicating our architecture enhances convergence rate.

7. Social Impact section will be added in revised paper (see global response).

This section was included to the global response above, however, we include it here for easy reference.

We'd like to clarify that when we refer to DiT-L/2 in the rebuttal, we specifically mean the DiT-L/2 model trained using the flow matching framework like SiT. This choice ensures a fair comparison of inference speed.

This comparison is more equitable because our DiMSUM model also employs the flow matching framework. Crucially, the sampling methods between traditional diffusion methods and flow matching methods differ significantly. By comparing DiMSUM to the flow matching-trained version of DiT (SiT), we ensure that both models use the same sampling approach during inference.

By aligning sampling methods, we can accurately evaluate inference speed differences, ensuring that any variations are due to architectural choices rather than differences in sampling or training processes.

We also sorry for the typo in table 2 of the rebuttal PDF, the parameters of DiMSUM-L/2 should be 460M, not 480M.

Thank you for your response! We provide our answers below:

(1) Why faster?

We agree with the statement “a single Mamba module is faster than a transformer module only when the number of tokens is greater than 1000”. In the provided table below, we can see that Gflops of DIMSUM with 256 token is higher than Gflops of DiT. The reason why our model take less time to generate image is that we use dopri5 to produce image at inference time. We conjecture that DIMSUM could converges to less ODE-curvature solution compared to DiT architecture [1], [2]. Therefore, our architecture use less NFE to achieve the better image quality compared to DiT architecture (this is also indicated by our FID result on several benchmarks like CelebA-HQ and Church). Moreover, with resolution 512 (1024 tokens), we can see that our architecture achieves lower Gflops and inference time. This demonstrate the potential use of DIMSUM in larger benchmark like text-to-image which has larger resolution.

[1] Pooladian, Aram-Alexandre, et al. "Multisample flow matching: Straightening flows with minibatch couplings." arXiv preprint arXiv:2304.14772 (2023).

[2] Lee, Sangyun, Beomsu Kim, and Jong Chul Ye. "Minimizing trajectory curvature of ode-based generative models." International Conference on Machine Learning. PMLR, 2023.

256 (latent size: $32 \times 32$)

512 (latent size: $64 \times 64$)

*Note: In the previous edition, we measured the time latency of all models on two resolutions using the same device and GPU_ID (a single NVIDIA A100). However, for the result of our model upon 512x512, due to resource matter (occupied device), we used a different NVIDIA A100 (different device/GPU_ID). We acknowledged that changing devices could cause unfairness. We fixed it by remeasuring using the same device/GPU_ID (same specs) for a fairer comparison. This issue doesn't affect GFLOPS since we did this separately and already used the same device/GPU_ID. Sorry for the inconvenience.

(2) Compare with SiT

Thank you for noting the SiT model comparisons. To clarify, the SiT model in Table 1 of our paper (FiD=2.15) and the SiT model with FiD=2.06 are the same but with different samplers. We reported the FiD=2.15 result for a fairer NFE comparison, which we'll explain further below.

The SiT authors used two sampling methods for their model. Table 9 of the SiT paper shows results for SiT-XL (cfg=1.5, ODE) with FiD=2.15 and SiT-XL (cfg=1.5, SDE:σ_t) with FiD=2.06. The latter, achieved with a 250 NFE Heun (SDE) solver, is computationally equivalent to a 500 NFE Euler solver, as the Heun solver requires an additional model forward pass per step.

Given the substantial computational requirements and our limited remaining time, evaluating ImageNet using this method is not feasible during the rebuttal phase. However, our comparison using SiT-XL (cfg=1.5, ODE) remains valid. Insights from SiT’s section 3.5 and Table 9 suggest that DiMSUM, using the same flow matching and training framework, would likely see a similar performance boost with the SDE sampler. We will conduct this evaluation and include the results in the revised version, with clear specification of the sampling method for full transparency.

Lastly, we'd like to highlight that our model is considerably smaller than the SiT model (460M vs 675M parameters). This size difference indicates the model scalability for improvement in our approach, as shown by the parameter scaling ablation in Table 4 of our rebuttal. This study reveals that the upscaled DiMSUM-XL/2 slightly outperforms the L/2 version in FID (3.76 vs 3.45).

We acknowledge your concerns as they are critical for the paper’s clarity and improvement. We will revise them in our manuscript.

We'd like to clarify that when we refer to DiT-L/2 in the rebuttal, we specifically mean the DiT-L/2 model trained using the flow matching framework like SiT. This choice ensures a fair comparison of inference speed.

This comparison is more equitable because our DiMSUM model also employs the flow matching framework. Crucially, the sampling methods between traditional diffusion methods and flow matching methods differ significantly. By comparing DiMSUM to the flow matching-trained version of DiT (SiT), we ensure that both models use the same sampling approach during inference.

By aligning sampling methods, we can accurately evaluate inference speed differences, ensuring that any variations are due to architectural choices rather than differences in sampling or training processes.

We also sorry for the typo in table 2 of the rebuttal PDF, the parameters of DiMSUM-L/2 should be 460M, not 480M.