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

Authors: hao  phung
quan  dao
trung  dao
hoang  phan
dimitris n metaxas
anh  tran
vinai  research
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 04:44:14

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

Main paper: 
 introduction: 
 Diffusion models [58, 22] are a trending generative model technique that has gained significant attention in machine learning and computer vision. The core idea behind diffusion models is to learn how to reverse the diffusion process by gradually transforming a simple initial distribution, like Gaussian noise, into a complex data distribution. The flexibility, robust performance, and high-quality outputs make them powerful tools for advancing the state-of-the-art in generative modeling, and large diffusion-based generators have revolutionized the field of image [51, 3] , video [23, 21] , and 3D synthesis [50, 62, 61] . While diffusion models initially rely on UNet architectures, recent methods have shifted gear to build upon transformer backbones. A line of works [48, 14, 11] have shown that transformer-based diffusion models are scalable and consistently offer higher generation quality 2 Related Work 
state space models and their applications in vision tasks: 
 In control engineering and system identification, state space models (SSMs) are described by state variables and first-order differential equations but initially underperformed in deep learning. Recent enhancements, notably by S4 [18] through the use of a HiPPO [16] initialization matrix, have significantly improved SSMs. Subsequent studies [18, 15, 17, 35, 1] show that SSMs can match transformers in long-range sequence modeling with the added benefit of linear time and space complexities. Notably, Mamba [15] has advanced over transformers in NLP by using a time-varying system with context-dependent parameters, enhancing the differentiation of hidden states over long sequences. This positions Mamba as a strong alternative to transformers across various domains.In computer vision, ViM and VMamba [74, 41] are the first works to introduce Mamba as a building block in discriminative tasks. Sequentially, Mamba is explored in many computer vision tasks such as medical imaging [43, 39] , point clouds [70, 34] , and image generation [25, 73] . Similar to vision transformers, images are divided into patches, and the patches are mapped into tokens. The tokens are then arranged in a sequence following a scanning order. In vision transformers, the scanning order does not matter since attention scores are computed between every token pair. However, SSMs consider the order information, introducing an inductive bias about 2D images into the model. Therefore, scanning order is vital in setting vision models' performance. ViM [74] proposed a bidirectional scanning order (sweep-2) for discriminative tasks. VMamba [41] proposed cross-scanning (sweep-4) per each Mamba building block. This cross-scan improves the model performance but costs enormous overhead. MambaND [33] reduces that cost by introducing two methods: interleaved scanning and multi-head scanning. Interleaved scanning, which alternates the scanning order in the sequential blocks, is simpler but gains better performance in discriminative tasks. Recently, Zigma [25] proposed a zigzag-8 scanning order to preserve the locality property (i.e., each token is adjacent to its next and previous tokens). The zigzag-8 scanning order shows faster convergence compared to the bidirectional one. 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 than sweep-4. Instead, sweep-4 offers the best performance (Section 4.4).Diffusion models [22, 57, 56, 51] are an emerging type of generative model that requires a sequential denoising process of several to thousands of steps to sample an image from initial Gaussian noise. Notably, most of them are based on Stochastic Differential Equations (SDE) that require an accumulation of additional stochastic noise at each generation step. Alternatively, there is a line of flow matching methods [36, 40, 44, 5] that emphasize deterministic trajectories from pure Gaussian noise to the target data distribution, favoring a straighter solution path. Their applications span across different tasks like image super-resolution [13] , depth estimation [19] , and motion synthesis [26] . Recent works [44, 31] have proved that diffusion models and flow matching are strongly correlated and can be converted into each other. In this work, we only focus on the simple objective of flow matching for our design.Meanwhile, most methods are originally based on Unet architecture, which utilizes convolution resblock † to capture local information at multiscale resolution. The attention layer is also used, interleaving between resblock layers to capture global information. Recently, the vision transformer [9, 42] has emerged and largely surpassed CNNs in many tasks. For diffusion image modeling, several transformer-based architectures [48, 2, 14] have been recently introduced. Transformerbased architectures capture global information better than Unet ones and can generate higher-quality images. Specifically, UViT [2] replaced convolution resblock layers with transformer blocks and removed downsampling/upsampling blocks. DiT [48] directly replaced Unet with a vision transformer. Inspired by this, MDT [14] and MaskDiT [72] introduced a mask latent modeling approach to better capture contextual information and enhance training efficiency. Although transformer architectures achieve better image generation, these models suffer from quadratic time and memory complexity, slowing down training and inference processes. Recently, with the birth of Flash Attention [8, 7] , both training and inference time of these transformer-based models are significantly reduced thanks to † A residual block is a skip-connection block that learns residual functions relative to the layer inputs.the reduced IO bottleneck. However, the computation time complexity remains quadratic. Recently, the S4 [18] model has been introduced to effectively deal with long-range dependency in the NLP field. Furthermore, the S4 model favors the linear complexity time and space, which is more efficient than the transformer. Among S4 class models, Mamba [15] stands out for its high performance in capturing long-range dependency. In diffusion models, DiffuSSM [65] adopts S4D [17] as a building block for their model and achieves better FID compared to transformer counterparts. Recently, Zigma [25] utilized Mamba for diffusion architecture, using a zigzag scanning order to preserve locality-aware scanning property. Despite showing promising results, Mamba-based diffusion models still struggle to find an optimal scanning scheme to take advantage of the 2D inductive bias from images. We find these approaches stuck in spatial processing, thus failing to incorporate global and local relations. These relations can effectively captured in frequency spectrum, thus we propose to incorporate frequency scanning alongside the existing spatial scanning mechanism.Employing frequency components extracted by Fourier, Cosine, or Wavelet transform in solving vision tasks was common in classical computer vision. Many modern works still find this practice useful in improving the performance of deep neural networks. In perception tasks, several works [37, 38, 67] integrate frequency processing in transformer architecture. NomMer [37] applied a discrete cosine transformer into global attention to efficiently yield synergistic context from both global and local contexts. To improve Masked Image Modeling, Ge-AE [38] introduces an additional frequency decoder using Fourier transform to reconstruct the high-frequency information better. Wave-Vit [67] applies wavelet into self-attention modules to reduce the time and space complexity of the transformer architecture while still preserving the performance. Recent work Simba [47] introduced Fourier-based layers (EinFFT) in combination with Mamba block to replace MLP layers for better channel mixing. To solve the image denoising problem, FreqMamba [71] applied a wavelet and Fourier transformer to process features injected into the Mamba block. In generative modeling, several works [66, 49, 69] corporate wavelet frequencies into generative framework. By explicitly decomposing features/images into high-and low-frequency bands through wavelet transform, the generative model can train stably and converge faster. Furthermore, the high frequencies can be learned more efficiently, leading to a sharper synthesis image. Observing the benefit of wavelet processing in generative modeling, we apply discrete wavelet transform on local features before feeding into the Mamba layers. By using cross-attention to fuse wavelet frequency features and spatial features, our method achieves significant improvement in image synthesis compared to merely spatial feature processing.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: We explicitly mention the limitation and suggest potential directions for future works in section 5.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
method: 
 This section presents DiMSUM, a novel architecture aiming for effective and high-quality image synthesis. Preliminary knowledge will be provided in section 3.1, then overview structure and mechanism of the proposed network (section 3.2), and finally its core components (sections 3.3 and 3.4).State Space Model (SSM). SSM is a new type of sequence model that uses an implicit hidden state h(t) ∈ R N ×L to map the 1D input signal x(t) ∈ R L to its corresponding output signal y(t) ∈ R L . This process is formulated by a parameter matrix A ∈ R N ×N and two projection parametersB ∈ R N ×1 and C ∈ R 1×N : h ′ (t) = Ah(t) + Bx(t), y(t) = Ch ′ (t).For practical usage, the continuous parameters (A, B) are discretized by a time-scale parameter ∆ to produce discrete parameters (A, B), following a zero-order hold (ZOH) rule:A = exp(∆ • A), B = (∆ • A) -1 (exp(∆ • A) -I) • ∆ • B).Hence, the continuous system is rewritten as follows: Mamba then proposes a selective mechanism to enrich the dynamic interactions of different sequential states. In other words, the constant parameters (B, C, ∆) are tuned into input-dependent parameters, enforcing the context awareness of input sequence states:h t = Ah t-1 + Bx t , y t = Ch t .B t = Linear N (x t ), C t = Linear N (x t ), ∆ t = Softplus(Broadcast L (Linear 1 (x t ))),where Linear * (.) is a projection layer to * -dimensional vector, Softplus(.) = log(1 + exp(.)), and Broadcast L (.) means duplicating a single-value vector to L-dimensional vector.Diffusion model. Diffusion models [55, 22, 57, 58, 51] are also known as score-based models that learn the transitional trajectories from a Gaussian noise to signals in the target domain. Most methods are based on Stochastic Differential Equation (SDE), requiring a larger number of function evaluations to generate an image. Recently, Flow matching [36, 40, 5, 44] has proved to be a promising method that finds a deterministic mapping between input Gaussian noise and input data via solving Ordinary Differential Equation (ODE). Given an input data x belonging to the modeling distribution p(x) and a random noise ϵ ∈ N (0, I), the forward process is formulated as:EQUATIONwhere x t is the noise-added data at a time step t ∈ [0, 1] and (α t , σ t ) are time-dependent functions of t. Particularly, these functions are constrained such that α 1 = σ 0 = 0 and α 0 = σ 1 = 1 to produce correct mapping between data x at t = 0 and noise ϵ at t = 1. Specifically, [36, 40, 5] use a simple linear function where α t = 1 -t and σ t = t. We employ Flow matching's training objective to estimate the velocity between noise ϵ and data x:EQUATIONwhere v θ is a velocity estimator implemented by a neural network with parameters θ and c is an input condition (e.g., class or text). If no condition is used, c is set to empty.Wavelet transformation. Among frequency transform techniques, wavelet transform stands out for its simplicity and efficiency. It preserves the structure of image space, with low-frequency subbands representing down-sampled approximations of the input image, while high-frequency ones emphasize local details such as vertical, horizontal, and diagonal edges. Particularly, Haar feature is the most prevalent wavelet transform, consisting of a low-pass filterL = 1 √ 2 [1 1] and a high-pass filter H = 1 √ 2 [-1 1].To decompose an image x ∈ R H×W , it needs to construct 4 kernels LL T , LH T , HL T , HH T , then applies them to the input image to extract corresponding subbands {x LL , x LH , x HL , x HH | x * ∈ R H/2×W/2 }, respectively. This process is called discrete wavelet transform (DWT). Notably, these filters are pairwise orthogonal, so an invertible matrix exists to map the data back to the original image space, coined as discrete inverse wavelet transform (IDWT). Given its benefits, we propose to use wavelet transform to supplement the local structure of frequency components into the process of Mamba, thus leading to enhanced image quality and training convergence, as demonstrated through our empirical experiments in section 4.Inspired by the advancement of Mamba-based diffusion models and frequency-based networks, we design a novel architecture DiMSUM for effective and high-quality image synthesis with the structure For illustration purpose, we plot wavelet representations of an input image but our real process is performed on encoded features of the input. Giving an image of size (8, 8) , for example, it is first decomposed to four wavelet subbands of size (4, 4) where each is further transformed to 2nd-level subbands of size (2, 2). Green dots indicate pixel points within each wavelet subband and a window of size 2 × 2 is used to perform scanning across multiple wavelet subbands like the CNN kernel. presented in Fig. 1 . Similar to Latent Diffusion Models [51] , our method performs image generation on the latent space of a pre-trained encoder † . The method first receives an input image and encodes it to a latent map of size 4 × H × W . It then processes the latent map using our proposed Diffusion Mamba network, whose core is a sequence of DiMSUM blocks, each consisting of DiM blocks that employ a novel Mamba structure with spatial and frequency scanning fusion and a globally weight-shared transformer block. The processed latent is then decoded to the output image.A core component of our approach is the DiM block that relies on a novel Spatial-Frequency Mamba fusion technique. In this section, we will discuss in detail the ideas behind this vital component.Scanning in frequency space. Mamba-based approaches in diffusion models often lack effective scanning schemes for preserving both local and global 2D spatial information. Although several works have proposed different heuristic scanning methods [74, 41, 33, 25] to address this issue, these approaches are insufficient for capturing local pixel dependencies and long-range frequency relationships. Though LocalMamba [28] proposed a window scanning to mimic the convolution kernel, it often underperforms compared to previously mentioned scanning methods as it is limited to the dependencies of nearby-pixels within window. DiMSUM addresses these challenges by decomposing the original image into frequency wavelet subbands. This approach is effective to capture long-range frequency while preserving relations across different subbands. We redesigned the window scanning, where each window corresponds to a subband of the frequency space as in Fig. 3 . Consequently, each window captures the full range of low/high-frequency signals from the original image. This advantage sets us apart from traditional window scanning in image space. As the model progresses through different subbands, it incorporates spatial information represented at various low-to-high frequencies, adding valuable context to the denoising process.Image WaveletWavelet Mamba. We now examine the integration of the wavelet transform into the Mamba structure shown in Fig. 2 . Wavelet Mamba first applies DWT to decompose input features into wavelet subbands. Our main setting uses two-level Haar wavelet to map input into low and high-frequency features. Given input feature x ∈ R C×H×W , first-level wavelet transform is applied to produce 4 wavelet subbands of size R C×H/2×W/2 . Each subband is then further decomposed into second-level wavelet subbands of size R C×H/4×W/4 . This feature is pivotal, as we decompose every wavelet subband to evenly separate an input image into multiple wavelet patches, unlike conventional wavelet transformations that process only LL subbands at the next level. In Wavelet Mamba, we concatenate those subbands to form a 1D sequence, apply window scanning within each subband, and slide across the sequence for feature extraction. The window scanning is inspired by a CNN kernel proposed in [28] with two window sliding directions: left to right and top to bottom. Note that since low-frequency subbands capture the main content of image, it should be input first. Therefore, we do not use reverse scanning orders: right to left and bottom to top. After passing through Wavelet Mamba, the output features are transformed back to input shape by using IDWT twice. With wavelet module, our model can better capture local structure of frequency information. Thus, incorporating Wavelet Mamba with Spatial Mamba can offer better performance, yielding high-quality image generation (Spatial-frequency Mamba in Fig. 1 ).Cross-Attention fusion layer. Given f s and f w are spatial and wavelet features obtained by sweep and window scan. We combine these features using a cross-attention fusion layer as follows:q s , k s , v s = Linear(f s ), q w , k w , v w = Linear(f w ), f out = Linear(Concat(Attn(q s , k w , v w ), Attn(q w , k s , v s ))).More specifically, we first compute each feature's query (q * ), key (k * ), and value (v * ) using linear layers. To fuse the information between spatial and wavelet features, we do cross-attention by swapping the queries (q * ) of spatial and wavelet before applying a self-attention module onto each key, query, and value triplet. Finally, we concat the outputs of two cross attentions by channel followed by a linear projection to obtain the output feature f out (see the last subfigure in Fig. 1 ).Conditional Mamba. Unlike attention modules, conventional Mamba has no explicit mechanism to inject input conditions into its flow. We propose a simple technique that enables Mamba to take in any conditional input c via initializing the very first hidden state with the embedding c instead of zero, as in original Mamba. This can be considered as a type of prior injection into Mamba. Specifically, the recurrent process of Mamba can be rewritten as below:EQUATIONIn conventional Mamba, h -1 is set to zero as there is no previous state at the beginning. Here, we set h -1 = Linear D (c) to inject context prior into flow of Mamba. As shown in ablation, conditional mamba effectively enhances model performance. This is beneficial for both unconditional and conditional generation. For unconditional image generation, we create an auxiliary learnable token to capture image space's global information, similar to vision transformers [10, 59] . For classconditional generation task, we use a class embedding to condition on Mamba. Conditional Mamba is enabled by default in DiM blocks (Fig. 1 ).Since Mamba is better than transformer at long-range dependency [18, 15] but weaker than transformer at in-context learning [46] , we propose a hybrid mamba-transformer architecture which favors both these properties as in recent work Jamba [35] . Motivated by Zamba [1] , we introduce the globally-shared transformer (attention) block. This shared attention block is added after each of four DiM blocks as shown in Fig. 1 since we want to preserve the continuity of the 4-sweep alternative scanning order. By using shared weights, we significantly reduce the number of parameters introduced by different attention blocks. This layer complements the flow of Mamba since transformers excel at extracting global relations without relying on manually defined orders of input sequences, as in Mamba. Hence, with this hybrid architecture, our method effectively addresses Mamba's order dependence while significantly reducing the FID with very few additional parameters.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] Justification: N/A, because our work focuses solely on empirical evaluation to measure the effectiveness of the proposed architecture.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
experiments: 
 Implementation. We established a depth of 20 layers, a base width of 1024, and a patch size of 2 for network configuration (further information on hyperparameters in appendix A). We run experiments on standard datasets: CelebA-HQ [27] , LSUN Church [68] , and ImageNet [52] . For the sampling method, we follow [44, 5, 25] to use adaptive ODE solver 'dopri5' for evaluation. We assessed performance using the Fréchet Inception Distance (FID) [45] and Recall [32] by first generating 50K images and then comparing them with the full reference dataset. We also report GFLOPs and the number of training iterations per second to further illustrate the model efficiency. On the CelebA-HQ dataset at resolutions of 256x256 and 512x512 (Fig. 4 ), our method achieved state-of-the-art FID scores of 4.62 and 6.09, respectively, surpassing the scores reported in recent studies. Furthermore, DiMSUM demonstrated superior recall scores, indicating a greater diversity in the generated samples compared to other methods. This result is particularly impressive, given the majority of baseline methods are based on diffusion processes, which are known for excellent diversity. On LSUN Church (Fig. 5a ), our method outperformed diffusion-based methods and achieved results nearly on par with GAN-based approaches. Moreover, our method recorded the highest recall metric of 0.56, significantly exceeding the recall of GAN methods, thereby underscoring its robustness in generating diverse and high-quality images. As reported in Tables 1, 4, 5a, our method requires less training epochs/iterations to reach the optimal performance compared to the baseline approach, implying a strong and fast convergence. To better illustrate the training convergence comparison, we illustrate in Fig. 4d the performance of different diffusion-based models over training epochs regarding the FID-10K on CelebA-HQ 256. Notably, our proposed method demonstrates a superior convergence speed, rapidly decreasing FID score and stabilizing at a significantly lower value than the other methods like LFM [24] , LDM [51] , and DiT [48] . This rapid descent is particularly evident within the first 150 epochs, after which our method maintains a low FID score and still with sight on decrement, suggesting a high-quality image generation capability. Compared to the learning curves of other methods, our method exhibits a more stable trajectory without significant oscillations between training epochs.In this section, we ablate the design choices for our network, using experiments on the CelebA-HQ 256 dataset. For the starting baseline, we adopt the same training settings from Zigma [25] . We choose sweep-4 with interleave scanning order [33] by default. In Fig. 2a , with our proposed conditional Mamba, the FID score is improved from 6.19 to 5.27, and the same trend is observed for recall. Meanwhile, adding Wavelet Mamba followed by a simple concatenation layer to combine spatial and wavelet features results in a worse score of 5.87 due to the weak alignment between these features. This demonstrates that our proposed cross-attention fusion layer is crucial for performance improvement, fully leveraging wavelet components to achieve a score boosted to 4.92. The performance is further enhanced to 4.66 by incorporating the weight-shared transformer block.Design choices of fusion layer. In Fig. 2c , we report metrics for different fusion layers, ranging from simple linear projection to cross-attention layers. As shown, our proposed cross-attention fusion with swapped query achieves the best results while requiring fewer parameters and GFLOPs than the attention option, with only a marginal increase in computation compared to the linear option.Number of wavelet levels. As shown in Fig. 2b , two wavelet levels provide the best performance on the CelebA-HQ 256 dataset. We argue that the choice of wavelet levels should be based on the input resolution. An input image of size 256 × 256, for instance, is encoded to a compact latent of size 32 × 32, which is further patchified by 2 to the small size of 16 × 16. Hence, applying 3 wavelet levels results in extremely small wavelet subbands of size 2 × 2, leading to reduced performance.Alternative frequency transform. Apart from wavelet transform, we also consider different types of frequency techniques like DCT and Fourier Transform (Fig. 2d ). For DCT, we propose a multi-order JPEG scanning strategy (illustrated in Fig. 8 ), based on JPEG compression [63] instead of the window scanning. For Fourier Features, we directly adopt the EinFFT block from SiMBA [47] . In either case, the performance drops compared with the default choice of wavelet.Transformer layer. In Fig. 2f , we assess the advantage of the transformer layer for both Conditional Mamba and Spatial-Frequency Mamba. As shown, the globally-shared transformer further boosts the performance of our Spatial-Frequency Mamba. In contrast, applying this layer solely to spatial Mamba increases the FID score by 0.13, highlighting the essence of our Spatial-Frequency Mamba in conjunction with the transformer layer. Meanwhile, replacing this shared layer with independent transformers results in a decline of 0.43 in FID and 0.03 in Recall.In Fig. 2e , we compare different scanning orders of Mamba. We keep it simple by using only Mamba block for all experiments without the globally-shared transformer and fusion layer. In Spatial-frequency Mamba, we use "Sweep-4" scanning for spatial features by default and only test other scanning methods for wavelet features. Overall, Sweep-4 performs best when combined with our proposed Conditional Mamba module. We also show that scanning with many-way orders like Sweep-8, Zigzag-8, and Jpeg-8 is not guaranteed to yield better performance than Sweep-4 scanning. In spatial-frequency Mamba, it is demonstrated that our proposed window scanning for wavelet Mamba provides a better outcome than conventional sweep-4 scanning.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: Yes, this information can be found in appendix A and section 4. 
conclusion: 
 Our paper introduces a novel, promising architecture that seamlessly integrates spatial and frequency features into Mamba process. By leveraging wavelet transform within the Mamba framework, our method enhances local structure awareness and ensures efficient spatial and frequency information fusion. This dual-focus strategy improves the detail and quality of generated images and accelerates training convergence. Our comprehensive experiments demonstrate that DiMSUM consistently outperforms state-of-the-art models of comparable size across multiple benchmarks, achieving lower FID scores and higher recall metrics, highlighting its ability to produce diverse and high-fidelity images. The proposed cross-attention fusion layer and globally shared transformer block also contribute to the model's robustness and scalability. Considering the promising results, we anticipate that future research in related domains, such as text-to-image synthesis, will adapt our backbone architecture and achieve comparable improvements in performance.Social impact and limitation. We believe that our proposed network advances the architectural design of state-space models for image generation. This model can be extended to various tasks, such as large-scale text-to-image generation and multimodal diffusion. While there is a risk that our architecture could be misused for malicious purposes, posing a social security challenge, we are confident that this risk can be mitigated with the recent development of security-related research. Hence, the positives can outweigh the negatives, rendering the concern minor.While our method outperforms other diffusion baselines in generation quality and training convergence, we acknowledge areas for improvement. These include designing a multiscale architecture and addressing manually defined scanning orders. Another advanced technique, such as masking training regularization [14, 72] , is orthogonal to our approach and could lead to further improvements. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: Yes, we will release our code and pretrained models on github.Guidelines:• The answer NA means that paper does not include experiments requiring code. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: No, because the evaluation process takes hours on a GPU and is very expensive to compute with multiple seeds. Importantly, we conducted evaluation for several benchmarks in section 4 so running with different seeds are not feasible in our scope.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We discussed the societal imparts in section 5.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).• The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We provide network configuration details in table 6b and hyperparams in table 6a. For the interpolent form of the forward process eq. ( 1), we opt to use a generalized VP form, which is proposed in SiT paper [44] . Particularly, this form defines α t = cos 1 2 πt and σ t = sin 1 2 πt . The choice of transformer block, we adopt a DiT block [48] where we replace the original MLP layer with a Gated MLP layer, similar to SwitchTransformer [12] .Besides, we illustrate the sweep scanning strategy in Fig. 9 for better comprehension. We also visualize the Jpeg scanning orders in Fig. 8 inspired by the Jpeg compression algorithm [63] and the multi-direction scanning from ZigMa paper [25] .Advantages of DiMSUM. First, it's important to highlight that our method outperforms both DiT and SiT while requiring less than a third of the training iterations, achieving the best FID score of 2.11. Compared to other state-space diffusion models, our method outperforms DIFFUSSM-XL, considering a similar training duration. 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.Clarification of scalable term. Given its current association with model parameter scaling in the LLM-dominated landscape, 'scalable' may be mistakenly interpreted as primarily referring to increasing the size of the architecture. However, in the broader context, machine learning/ deep learning algorithms' scalability may refer to their capacity to handle bigger datasets and computational resources while producing correct results in an acceptable period of time.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 4 .Our SoTA results are achieved with a parameter count comparable (or even smaller) to existing models. Refer to Table 4a , 5a and 1. This suggests substantial room for further enlargement of our model's parameters, which we anticipate will yield even greater improvements across various and bigger datasets (see Table 3 ).Why would using only wavelet scanning hurt model performance (as observed in Table 2a )?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 , contributing to the SoTA result of our method. Memory and GFLOPs. The results, as shown in the table 4, 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. Note that the parameter change after changing image size is mainly due to the PatchEmbed layer of the architecture, which both models have.Regarding GFlops, we acknowledge that for 256 × 256 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, proving the efficiency of our method for high-resolution image synthesis. Consequently, at this higher resolution, DiT's GFlops exceed DiMSUM's by approximately 7%. This observation further highlights the strength of Mamba in handling longer context length.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.With these two points, we emphasize upon the importance of our proposed architecture, rather than just the benefits given by the flow matching framework.Speed gain. In table 4 , DiMSUM-L/2 achieves 2.2 seconds latency compared to DIT-L/2 with 3.8 seconds though DiMSUM has larger GFLOPs. Notably, with resolution 512 (around 1024 tokens), the speed gap between our method and DiT becomes more significant and our architecture also has lower Gflops. This demonstrate the potential use of DiMSUM in larger benchmark like text-to-image which has larger resolution.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:• 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.• Adaptive Sampling Efficiency: We employ the dopri5 ODE adaptive solver for sampling from both models, similar to SiT [44] and Zigma [25] . 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.Full training curve of DiT. The plot Fig. 4d 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. Here, we further provide a complete training curve of DiT in Fig. 6 . It's noted that DiT and LFM in Fig. 4d 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.Additional ablation of scanning orders. To substantiate our claims regarding the efficiency of frequency information, we conducted a comprehensive ablation study. This ablation 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 5 shows that 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.Varying NFE. We plot the FID-10K scores of various NFE used for evaluation in Fig. 7 . This shows that increasing NFE beyond 250 leads to minimal or no improvement in the FID scores. This behavior • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.†https://huggingface.co/stabilityai/sd-vae-ft-ema 


Decision: The paper received mixed reviews from four reviewers: one strong accept, one borderline accept, and two borderline rejects. All reviewers appreciated the novelty and effectiveness of the proposed state-space architecture for diffusion models. The authors effectively addressed most of the reviewers' concerns during the rebuttal. However, the two reviewers who gave borderline rejects still have some concerns regarding inference speed and performance with longer training epochs. The authors provided further explanation, but these two reviewers did not respond to the authors' additional explanations.

After reviewing the paper and the rebuttal, AC finds the authors' explanations satisfactory. Given the overall contribution and the positive aspects highlighted, AC believes the paper is worthy of acceptance and should be shared with the research community. This decision has been discussed with SAC. At the same time, AC encourages the authors to add some important rebuttal content into the final camera-ready version and polish their paper by following reviewers' suggestions.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: limitations: See the weaknesses part.

All of my concerns have been addressed. I will increase my score.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: as above
limitations: limitations: as above
paper_summary: paper_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.
strength_weakness: strength_weakness: + 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: 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.
limitations: limitations: Yes, they discussed

Comments: title: Thank you.
Thank you.
Thank you.
Thank you.
Clarification of DiT term in the rebuttal
Clarification of DiT term in the rebuttal
reply
Read the rebuttal and discuss with the authors
comment: We sincerely thank you for your thorough review and your review is insightful and helpful to our paper. We hope our answer could resolve your concern. We wish you a good Neurips.
We sincerely thank you for your thorough review and your review is insightful and helpful to our paper. We hope our answer could resolve your concern. We wish you a good Neurips.
We sincerely thank you for your thorough review and your review is insightful and helpful to our paper. We hope our answer could resolve your concern. We wish you a good Neurips.
We sincerely thank you for your thorough review and your review is insightful and helpful to our paper. We hope our answer could resolve your concern. We wish you a good Neurips.
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.
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.
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.
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$)

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

| 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.
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.
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.
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!
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.
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.
Thanks for the feedback. All of my concerns have been addressed. I will increase my score.
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