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

We introduce a novel state-space architecture for diffusion models, effectively harnessing spatial and frequency information to enhance the inductive bias towards local features in input images for image generation tasks. While state-space networks, including Mamba, a revolutionary advancement in recurrent neural networks, typically scan input sequences from left to right, they face difficulties in designing effective scanning strategies, especially in the processing of image data. Our method demonstrates that integrating wavelet transformation into Mamba enhances the local structure awareness of visual inputs and better captures long-range relations of frequencies by disentangling them into wavelet subbands, representing both low- and high-frequency components. These wavelet-based outputs are then processed and seamlessly fused with the original Mamba outputs through a cross-attention fusion layer, combining both spatial and frequency information to optimize the order awareness of state-space models which is essential for the details and overall quality of image generation. Besides, we introduce a globally-shared transformer to supercharge the performance of Mamba, harnessing its exceptional power to capture global relationships. Through extensive experiments on standard benchmarks, our method demonstrates superior results compared to DiT and DIFFUSSM, achieving faster training convergence and delivering high-quality outputs. The codes and pretrained models are released at https://github.com/VinAIResearch/DiMSUM.git.

Paper

References (77)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer C5pc4/10 · confidence 5/52024-06-29

Summary

This paper introduces a state-space architecture for diffusion models that enhances local feature detection in image generation by combining spatial and frequency information. Traditional state-space networks like Mamba struggle with visual data processing, but integrating wavelet transformation improves local structure awareness. The method fuses wavelet-transformed outputs with original Mamba outputs and incorporates a globally-shared transformer to capture global relationships. Experiments show this approach achieves faster training convergence and high-quality outputs, demonstrating state-of-the-art results.

Strengths

+ The proposed method works on the frequency domain seems to be a good idea compared to existing mamba diffusion approaches. + The architecture looks like getting inspiration from DiT and transferring to the Mamba style may be interesting. + Performance on several datasets shows the promising method.

Weaknesses

+ 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. + 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. + The scalability is important for this architecture, but it is missing, which would weaken the paper. + 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.

Questions

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. Q2. How is the performance w.r.t more sampling steps? A curve is a good demonstration.

Rating

4

Confidence

5

Soundness

3

Presentation

2

Contribution

2

Limitations

Yes, they discussed

Authorsrebuttal2024-08-13

Thank you for your response! We provide answers below: Q1. 512 latency 256 (latent size: $32 \times 32$) | Method | Time | Params | | --- | --- | --- | | DiMSUM-L/2 | 2.20s | 460M | | DiT-L/2 | 3.80s | 458M | 512 (latent size: $64 \times 64$) | Method | Time | Params | | --- | --- | --- | | DiMSUM-L/2 | 2.86s | 461M | | DiT-L/2 | 4.78s | 459M | *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$) | Method | MEM | GFlops | | --- | --- | --- | | DiMSUM-L/2 | 2.42G | 84.49 | | DiT-L/2 | 2.30G | 80.74 | 512 (latent size: $64 \times 64$) | Method | MEM | GFlops | | --- | --- | --- | | DiMSUM-L/2 | 2.46G | 337.48 | | DiT-L/2 | 2.34G | 361.14 | 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.

Reviewer C5pc2024-08-13

1. Could authors confirm that the proposed method is faster mainly because of using less NFE in sampling? And will it be slower if both methods are used under the same NFE? 2. How many NFEs are used concretely for the tables just newly provided for each method (DiT and DiMSUM)? This is to ensure the understanding that the advantage of speed comes from the Mamba thing or just using fewer NFEs. Reading the whole paper, I just had in my mind that the model with the Mamba hybrid Transformer is faster than the pure Transformer, but it is not if they just use the same NFEs. And of course, their proposed method satisfied the need for more Memory and more parameters used.

Authorsrebuttal2024-08-13

Thanks for your quick response! We provides our answers to your questions as follows: Q1. Observing the Memory and GFLOPS table above, it's true to claim DiMSUM-L/2 shows slower inference speed compared to its counterpart for 256x256 images using the same NFE (due to bigger GFLOPS), however, there are two crucial points to consider: 1. Scaling Efficiency: When we increase the image size to 512x512, as evident from the Memory and GFLOPS table, our model actually requires fewer GFLOPS at this higher resolution, thanks to its slower latency scaling which we mentioned above. Consequently, for 512x512 images and larger, DiMSUM-L/2 would outperform its counterpart in speed given the same NFE. 2. Adaptive Sampling Efficiency: We employ the dopri5 ODE adaptive solver for sampling from both models. This solver dynamically adjusts the NFE based on the initial noise and diffusion model characteristics, using the minimum NFE necessary to achieve optimal image quality. Notably, DiMSUM requires fewer NFE to meet the dopri5 stopping condition while still achieving a significantly better FID score than DiT. **We hypothesize that our proposed hybrid architecture converges to a better solution with less curvature, enabling high-fidelity image production with fewer NFEs.** With these two points, we emphasize upon the importance of our proposed architecture, rather than just the benefits given by the flow matching framework. Q2. About the exact number of NFE: Since dopri5 requires different NFEs depending on initial noise, we sample 100 random noises and generate images for both models using dopri5 to measure the NFE. The average NFEs for DIMSUM and DIT at resolution 256 are 61 and 143. For resolution 512, the average NFEs for DIMSUM and DIT are 82 and 138.

Authorsrebuttal2024-08-13

Clarification of DiT term in the rebuttal

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.

Reviewer nhSz8/10 · confidence 5/52024-06-30

Summary

This paper introduces a novel architecture for diffusion models that leverages spatial and frequency information to emphasize local features in image generation. It integrates wavelet transformation into the state-space networks, such as Mamba, to enhance the awareness of local structures in visual inputs. The outputs are fused with original Mamba outputs using a cross-attention fusion layer to optimize order awareness, crucial for image generation quality. A globally-shared transformer is added to boost Mamba's performance. The method achieves state-of-the-art results on standard benchmarks, with faster training convergence and high-quality outputs.

Strengths

- Overall, I appreciate this paper. It selects the most promising technical stacks—combining flow matching and Mamba—to address generative tasks. - The paper is well-organized, featuring valuable methodology design and experimental ablations. - It's interesting to see that the community has started considering state exploration in Mamba, as indicated in Line245. The initialization of the state is critical. We need to examine the code to understand how to implement it efficiently. See https://github.com/state-spaces/mamba/issues/101 for more details.

Weaknesses

- In Line 106, I'm unsure why sigma brings extra computation when the scan path is larger, given that the computation is always fixed. - 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]. - In Line 176, I'm curious as to why only B, C, and delta are time-dependent. Why isn't A? - In Line183, One paper is missed. [5] - In Fig 3a, what does the notation on ZigMa mean? - Regarding line 261, what's the rationale behind sharing those parameters of transformers? Are there any previous works indicating its effectiveness? - In Table 1, it would be better to display training steps. Researchers of generative models tend to compare using training steps rather than epochs. - 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. - In Fig5, is there any reference to the jpeg-8 scan path? How does it defined? Suggestions: - The previously mentioned figure or table should be placed first. [1]. Boosting Latent Diffusion with Flow Matching [2]. DepthFM: Fast Monocular Depth Estimation with Flow Matching [3]. Motion Flow Matching for Human Motion Synthesis and Editing [4]. Flow Matching for Conditional Text Generation in a Few Sampling Steps [5], Latent Space Editing in Transformer-based Flow Matching Overall, I think the exploration of using Mamba and flow matching is the right way to go, while I have some technical concerns about it. I am eager to hear the authors’ feedback.

Questions

as above

Rating

8

Confidence

5

Soundness

4

Presentation

4

Contribution

4

Limitations

as above

Reviewer eitf5/10 · confidence 4/52024-07-08

Summary

This paper proposes a novel Mamba-based diffusion model, DiMSUM, which leverages both spatial and frequency information to enhance visual generation capabilities. Specifically, DiMSUM applies wavelet transformation to the input signals, decomposing them into wavelet subbands. By employing a query-swapped cross-attention technique, DiMSUM effectively integrates spatial and frequency information. Furthermore, DiMSUM incorporates several transformer blocks into the Mamba model, enriching the global context features.

Strengths

1. The paper is well-motivated. It observes the difficulties caused by the manually-defined scanning orders in tradition Mamba-based diffusion models, and uses transformation in the frequency domain to mitigate this problem. 2. The paper is well-written and easy to follow.

Weaknesses

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. 2. Stronger baselines (such as [1]) and the experiments on more datasets (such as conditional generation on MSCOCO) and resolutions are required. 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. [1] Chen, Junsong, et al. "Pixart-$\alpha $: Fast training of diffusion transformer for photorealistic text-to-image synthesis." arXiv preprint arXiv:2310.00426 (2023).

Questions

See the weaknesses part. Meanwhile, I recommend the authors to provide sufficient experiments to support your method. And I would consider increasing my score if the above concerns can be addressed.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

See the weaknesses part. All of my concerns have been addressed. I will increase my score.

Reviewer RLU74/10 · confidence 4/52024-07-09

Summary

This paper proposes to employ a Mamba-transformer hybrid framework with wavelet-spatial scanning for image generation. The authors claim that the scanning in the frequency domain can model the long-range relations of tokens in 2D spatial space. The experiments and ablation studies have shown the effectiveness of the proposed method.

Strengths

- This paper is easy to follow. - The implementation details and the codes are included in this submission. - The performance on multiple datasets and sufficient ablation studies have verified the effectiveness of the proposed method.

Weaknesses

- The proposed method is too complex. Scanning the image and frequency space makes the framework complicated. The authors further introduced transformer blocks in the mamba framework, which added more complexity to the entire method. - Why do the authors mention "manually-defined scanning orders" in their abstract? This issue is not discussed in the subsequent contents of the paper. Moreover, it cannot be addressed by the scanning in the frequency domain. The representation of an image in the frequency domain also has two dimensions, so the challenges for the scanning on the frequency domain should be the same as the scanning on the image space. Moreover, the scanning directions of the proposed method are also manually defined, even a window scanning is involved. - The comparison of the real inference speed between the proposed framework and transformers is not included in this paper. - Although the authors claim that they propose a Mamba-based diffusion model. However, this model still includes many transformer blocks, so this is a hybrid model rather than a Mamba model. Some claims in this paper like ''mamba-based'' or ''Mamba models'' should be revised. Moreover, is the spatial-frequency scan also powerful on the pure transformer model? Does the Mamba-transformer beat the pure transformer model given the same input setting? - The title of this paper claims that the proposed method is scalable. However, the parameter counts of the proposed model is smaller than 1.0B, so this title may not be proper.

Questions

- In Figure 3(d), why does DiT perform so poorly on CelebA-HQ? Are there any bugs? DiT works very well on many large-scale datasets.

Rating

4

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

Although this paper uses public datasets, at least a section for social impact should be included in this paper. Moreover, this papers does not discuss the limitations of the proposed method.

Authorsrebuttal2024-08-13

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$) | Method | Time | Params | GFlops | | --- | --- | --- | --- | | Ours-L/2 | 2.20s | 460M | 84.49 | | DiT-L/2 | 3.80s | 458M | 80.74 | 512 (latent size: $64 \times 64$) | Method | Time | Params | GFlops | | --- | --- | --- | --- | | Ours-L/2 | 2.86s | 461M | 337.48 | | DiT-L/2 | 4.78s | 459M | 361.14 | *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.

Authorsrebuttal2024-08-13

Clarification of DiT term in the rebuttal

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.

Area Chair 27KU2024-08-09

Read the rebuttal and discuss with the authors

Hi Reviewers, The authors have submitted their rebuttal responses addressing the concerns and feedback you provided. We kindly request that you review their responses and assess whether the issues you raised have been satisfactorily addressed. If there are any areas where you believe additional clarification is needed, please do not hesitate to engage in further discussion with the authors. Your insights are invaluable to ensuring a thorough and constructive review process. Best AC

Reviewer eitf2024-08-10

Thanks for the feedback. All of my concerns have been addressed. I will increase my score.

Authorsrebuttal2024-08-10

We sincerely thank you for the insightful review and supportive feedback to make the paper more complete. Your support is highly appreciated, and we're glad that our responses have addressed your concerns.

Reviewer nhSz2024-08-12

reply

Thanks for the author's reply. I have also read the other reviewer's opinion. My concerns are fully resolved, so I will increase my score to 8 to reflect this. I encourage the authors to consolidate the results on ImageNet to compare with related baselines with Zigma.

Authorsrebuttal2024-08-12

Thank you for raising your concerns as well as for the detailed review of our paper's presentation. Addressing those definitely improves our paper generally. We will further follow your valuable recommendation!

Reviewer C5pc2024-08-12

Thanks for the rebuttal, it addressed some concerns, still remained some unclear points. 1) How is the latency comparison with 512x512? 2) Memory use is an important metric when claiming the efficiency of the proposed method (in Line 127), could authors provide Memory Use and Gflop besides a comparison with others for 256 and 512 resolution they showed in Tab. 2 rebuttal? 3) Could the authors explain some insights as to why their method is faster than DiT although it has more parameters than DiT? 4) In the proposed method, the input image is further processed with several handcrafted features of wavelet transforms and scanning operations, before patchifying and feeding into the model, which will also be considered to take more processing time to slow down the speed, are they accounted for in the inference speed when making a comparison for both 256 and 512 resolutions, and are they used the same NFE to compare? 5) In the paper, it mentioned transformers-based methods, including MDT, and MaskDiT, and the proposed method obviously cannot be comparable with these models, but the abstract claims it achieves state-of-the-art, which I believe is misleading and inappropriate.

Reviewer RLU72024-08-13

Thank you for your reply. Your rebuttal has solved most of my concerns, but some are still not addressed. The remaining concerns almost overlap with Reviewer C5pc's new concerns. For example, (1) why your mamba-transformer is faster than DiT? The rebuttal shows that DIMSUM is faster than DiT on ImageNet 256 (256 tokens considering the downsampler of the pre-trained autoencoder), however, a single Mamba module is faster than a transformer module only when the number of tokens is greater than 1000 (please refer to Fig. 8 in [1]); (2) the performance of DiMSUM after longer training on ImageNet 256 is 2.11 which is worse than SiT (FiD=2.06) [2], though both of these methods use flow matching. Thus, the performance gain of Mamba remains ambiguous (at least marginal). [1] Gu A, Dao T. Mamba: Linear-time sequence modeling with selective state spaces[J]. arXiv preprint arXiv:2312.00752, 2023. [2] Ma N, Goldstein M, Albergo M S, et al. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers[J]. arXiv preprint arXiv:2401.08740, 2024.

Authorsrebuttal2024-08-13

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.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC