DiMSUM: Diffusion Mamba - A Scalable and Unified Spatial-Frequency Method for Image Generation
Authors: Hao Phung
Quan Dao
Trung Tuan Dao
Hoang Phan
Dimitris N. Metaxas
Anh Tuan Tran
Venue: NeurIPS 2024 poster (2024)
Keywords: Diffusion models
Mamba
Wavelet transformation
Primary area: generative_models
TL;DR: A new Mamba architecture for image diffusion models

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.

DiMSUM Diffusion Mamba - A Scalable and Method for Image Unified Generation
Hao Phung"15 133
Quan Dao 12f
Trung Dao
Hoang Phan
Dimitris N. Metaxas
Anh Tran
VinAI Research
Ruttgers University
Commell University
*New York University
htp26@cormel1.co edu
quan. daogrutgers edu
v. trungdt21evinai.io v. v.anhte162@vinai.ise
hvp20110nyu.edu
dnm@cs rutgers. edu
Abstract
We introduce novel state-pace architecture for diffusion models, effectively harnessing spatial and frequency information to enhance the inductive bias to- wards local features input images for image generation tasks. While state-space networks, including Mamba, revolutionary advancement in recurrent neural networks, typically sean input sequences from left to right, they face difficul- ties in designing effecctive 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 frequencies by disentangling them into wavelet subbands, representing both low- and figh-freeuency components. These wavelet- based outputs are then processed and seamlessly fused with the original Mamba outputs through cross-attention fusion layer, combining both spatial and fre- quency information to optimize the order awareness of state-space models which is essential for the details and overall quality of image generation. Besides, introduce globally-shared transformer to supercharge the performance of Mamha, 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 de- livering high-quality outputs. The codes and pretrained models are released at https://githhnn www Comm/VinAIReseaan ctt.///iiATTeeearrr//////M git
Introduction
are a trending generative model technique that has gained significant attention in machine learning and computer vision. The core idea behind diffusion models to learn how to reverse the diffusion process by gradually transforming 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-ar in generative modeling, and large diffusion- generators have revolutionized the field of image [5]] [5]]] 3] video [23] 21. and 3D synthesis (50 62 61). While diffusion models initially rely UNet architectures, recent methods have shifted gear to build upon transformer backbones. A line works [48] [4] [1]] have shown
Diffusion models 558 22
that transformer- based diffusion models are scalable and consistently offer higher generation quality
Equal contributions.
Work done while at VinAI.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

than the UNet- based counterparts. Even the most common open-source text-to-image tool, Stable Diffusion, has switched to use transformers in their upcoming release [53]. Hence, transformers are becoming the new backbone standard for diffusion models. The power of this structure lies in the attention mechanism for capturing richer in-context relations. However, transformers have the drawback of a costly quadratic complexity, which might hinder their feasibility for high-dimensional data.
While transformers are taking over state-of-the-art diffusion backbones, a novel technique called state- space models (SSMs) has suddenly arrived, promising a better alternative. SSMs [8] I6 [5] have revolutionized the NLP field, favoring linear time complexity and excelling at long long-context modeling. This type of network bears similarities to the recurrent process of RNNs while being capable of fully operating in parallel like convolutional networks. SSM promises to surpass transformers in most tasks, prioritizing compute efficiency, such as long-sequence modeling. Mamba [15] is a special type of SSM that offers greater quality by introducing time conditioning and context dependency to the hidden states. In the context of computer vision, within a very short period, this architecture has been used to address variety of problems, including image perception [74 28 33], image restoration 120 71. and image generation [41 74 65 25]. In diftusion- image generation, Diffusion State
Space models (DIFFUSSM) 655 already surpass their transformer-based counterparts.
Though showing many advantages, Mamba still has critical weakness when processing 2D imagery data. Like vision transformers. images are divided into patches and then mapped into tokens. Mamba processes these tokens following a specific scanning order, thus introducing an inductive bias about 2D images into the model. Specifically, this order greatly impacts the interplay between image tokens, thereby affecting model performance. This characteristic unfavorable, particularly when transformers have no such order-dependency issue. Many vision-based Mamba studies have focused on solving this problem on proposed advanced scanning mechanisms like bi-directional [74), cross- scanning [4 [33, or 8-directions zigzag [25]. Despite improving performance, these scanning techniques still fail to capture global and long-range relations and do not fully release Mamba's potential.
In this paper, we enhance Mamba-based diffusion models, specifically focusing on image generation. Previous models failed to address the scanning order issue due to their exclusive reliance on spatial processing, overlooking crucial long-range relations that manifest in the frequency spectrum. We suggest a novel approach integrating frequency scanning with the conventional spatial scanning mechanism. Although initial work in Mamba has explored this combination for limited task of image deraining [71), it lacked a comprehensive analysis of the effective integration of these features.
Motivated by the above observation, this paper introduces DiMSUM, novel architecture that harnesses Mamba's 's power to unlock diffusion models' generation capabilities. Our approach enhances sensitivity to local structures and long-range dependencies by integrating wavelet transforms and spatial information. Using a query-swapped cross attennaton technique, we dynamically synergize spatial and frequency information, accelerating convergence and improving image synthesis quality. Consequently, this boosts image quality and enhances the efficiency and scalability of the training.
Additionally, we incorporate globally shared transformer blocks address global context integration, a limitation of traditional Mamba models. The block can also be viewed as a token-mixing layer that enriches global relations among image tokens, addressing the weak inductive bias of the manually defined scanning orders in the original Mamba. Hence, DiMSUM can maintain high performance even with larger, more complex datasets. Extensive experiments show that DiMSUM achieves state-of-thh-as FID scores and recall, setting new benchmark in generative image modeling.
In summary, our contributions lie three-fold: (1) A novel Mamba architecture for diffusion models that leverages both spatial and frequency features to enhance the awareness of local structures within input images, leading to better image generation. (2) We interleave globally shared transformer blocks per a certain number of Mamba blocks. The transformer with a strong capacity for capturing global relationships significantly boosts generation results when integrating with Mamba. The transformer can also be seen as an order-invariant mixing layer that complements Mamba's 's loose assumptions about the order of 2D data. (3) Superior results across image generation benchmarks like ImageNet, Cel and LSUN Church. Additionally, our method maintains comparable GFLOPs and parameters with existing diffusion architectures while offering faster training convergence.
2

---

2 Related Work
2.1 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 HiPPO [16] initialization matrix, have significantly improved SSMs. Subsequent studies [18] [3]
show that SSMs can match transformers in long-range sequence modeling with the added benefit of linear time and space complexities. Notably, Mamba [1]] has advanced over transformers in NLP by using 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.
17 35
In computer vision, ViM and VMamba [7| 41 are the first works to introduce Mamba 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 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 inductive bias about 2D images into the model. Therefore, scanning order is vital in setting vision models' performance. ViM [74] proposed bidirectional scanning order (sweep-2) for discriminative tasks. VMamba 44 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
discriminative tasks, Recently, Zigma [255 proposed a zigzag-8 scanning order to preserve the locality property (i.e., each token is adjacent its next and previous tokens). The zigzag-8 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.4p
2.2 Diffusion Architecture
Diffusion models 22 57 56,55 51 are an emerging type of generative model that requires a sequential denoising process of several thousands of steps to sample image from initial Gaussian noise. Notably, most of them are based on Stochastic Differential Equations (SDE) that require an aceu- mulation of additional stochastic noise each generation step. Alternatively, there a line of flow matching methods [36 40 44 5 that emphasize deterministic trajectories from pure Gaussian noise to the target data distribution, favoring straighter solution path. Their applications span across different tasks like image super- resoutreso 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 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, have been recently introduced. Transformer- based architectures capture global information better than Unet ones and can generate higher-quality images. Specificallyy UViT (2) replaced convolution resblock layers with transformer blocks and removed hownsampling/ussappingg blocks. DiT (88 directly replaced Unet with vision transformer. Inspired by this, MDT [4] 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 77. both training and inference time of these transformer-based models are significantly reduced thanks to
several transformer--aat based architectures [48, 2 [4]
A residual block is conneccnnttoonn block that learns residual functions relative to the layer inputs.
3

---

the reduced IO bottleneck. However, the computation time complexity remains quadratic. Recently, the S4 [8] 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 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-bas 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 effectively captured in frequency spectrum, thus we propose to incorporate frequency scanning alongside the existing spatial scanning mechanism.
2.3 Frequency-basee networks
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 gigeffreyeeccy 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, Famba [71] applied 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 fcatures/imagee into high- and low- 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 image synthesis compared to merely spatial feature processing.
3 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.
3.1 Preliminary
State Space Model (SSM). SSM is a new type of sequence model that uses an implicit hidden state h(t) E NxL. map the 1D input signal z(t) E R² to its corresponding output signal v(() E R.L. This process is formulated by a parameter matrix A E JNNXN and two projection parameters B E Nx1 and C E RIXN
h'(t) Ah(() + Bx((), v(t) = Ch'(t.
For practical usage, the continuous parameters (A, B) are discretized by a simel parameter A to produce discrete parameters (A, following zero-order hold (ZOH) rule:
A = exp( A A), B AA A) (exp(A A A) I) A -B).
Hence, the continuous system is rewritten as follows:
h Aht- BTg, yt Ch+.

---

Ball MSP
Uneer Pastaae
negga
Blokk
Comsila lasin
Lieer
Mardo
Mirde
Scals
Dwl
Pazzil
Evbed
Vayda
Newe
Danedin line
VUP
MLP
BEicE
Candinn
Cordianna
Fynke!
Diflusion Mamba
DIM Block
Spatial frecuency Mamba
Creas-Attation layer
Figure 1: Overview of DIMSUM architecture.
Mamba then proposes selective mechanism to enrich the dynamic interactions of different sequential states. In other words, the constant parameters (B, C, A) are tuned into input-dependent parameters, enforcing the context awareness of input sequence states:
(aailluu(Brodcastt... (Linearr (xx))); Broadcastz()) means duplicating single-value vector to L-dimensional vector.
B. L (zz), C, Linear A
where Linear, (-) is a projection layer to t-dimensional vector, Softplus(-) log(1 + exp(-)), and
Diffusion model. Diffusion models [55] 22 57 58 51 are also known as score-based models that learn the transitional trajectories from Gaussian noise to signals in the target domain. Most methods are based on Stochastic Differential Equation (SDE), reguiring a larger number of function evaluations 44 has proved to be promising method
generate an image. Recently, Flow matching [6 that finds 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 random noise E (0, I), the forward process is formulated as:
It IZ2 + ED
where v4 is the noise-added data time step E [0, 1] and (as, σe) are time-dependent functions of t. Particularly, these functions are constrained such that a₁ = Oo = 0 and 0 = 1 to produce correct mapping between data x at and noise e & 1. Specifically, [36 40 use simple
linear function where a; = t and or = t. We employ Flow matching's 's training objective to estimate the velocity between noise and data a::
i argmin Et,x, [l]=]
To (0 c,t)
where Uh velocity estimator implemented by a neural network with parameters 0 and c is input condition (e.g., class or text). If no condition is used, e is set to empty.
Wavelet transformation. Among frequency transform techniques, wavelet transform stands out for its simplicity and efficiency. preserves the structure of image space, with cow-frequeccy subbands representing down-sampled approximations of the input image, while high- fighffreyuency ones emphasize local details such as vertical, horizontal, and diagonal edges. Particularly, Haar feature the most prevalent wavelet transform, consisting of low-pass filter + [1 and
high-pass filter H 4 kernels LLT, LHT,, "H'" HHT, then applies them to the input image to extract corresponding subbands {ØLL, LLH, DHL, THH Ix E /2//2x/22), respectively. This process called discrete wavelet transform (DWTT). Notably, these filters pairwise orthogonal, 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 wavelet transform to supplement the local structure of frequency components into the process Mamba, thus leading to enhanced image quality and training convergence, as demonstrated through our empirical experiments in section 4
V2 [-1 1]. To decompose image R it needs to construct
3.2 Overview of the proposed network
Inspired by the advancement of Mamba-hased diffusion models and frequency- networks, we design novel architecture DiMSUM for effective and high-quality image synthesis with the structure
5

---

Horizonaal scarming
ib
ib
T
Input image
Wavelet level
Wavelet lovel 2
Figure 2: Illustration of Wavelet Mamba (Best view in color). For illustration purpose, we plot wavelet representations of input image but our real process is performed encoded features of the input. Giving an image of size (8,8), 8). for example, 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 x is used to perform scanning across multiple wavelet subbands like the CNN kernel.
presented in Fig. Similar to Latent Diffusion Models S1 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 x 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.
3.3
DiM block
core component of our approach the DiM block that relies on novel Featial-Frequecyy 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
to address this issue, these approaches are insufficient for capturing local pixel dependencies and long-range frequency relationships. Though LocalMamba (28) proposed window scanning 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.
33 25
DiMSUM addresses these challenges by decom- posing the original image into frequency wavelet subbands. This approach effective to capture long-range frequency while preserving relations across different subbands. We redesigned the window scanning, where each window corre- sponds to a subband of the frequency space as in Fig. 3 Consequently, each window captures the full range of low/high-frequenge signals from the original image. This advantage sets us apart from traditional window scanning in image space. As the model progresses through differ- ent subbands, it incorporates spatial information represented various low-to-high frequencies, adding valuable context to the denoising pro- cess.
Image
ib
Wavelet
Figure 3: Comparison of window scanning on im- age and wavelet space. For illustration, one-level wavelet transformation is applied and each sub- band is half the resoluton of original image.
Wavelet Mamba. We now examine the integra- tion of the wavelet transform into the Mamba structure shown in Fig. Wavelet Mamba first applies DWT to decompose input features into wavelet subbands. Our main setting uses two-level Haar wavelet map input into low and high-frequency features. Given input feature x E RCxHXX xWW first-level wavelet transform is applied to produce 4 wavelet subbands of size CC H/2 Each
http:://huggiinnnecee..sssaaaaal ft-eaa

---

subband is then further decomposed into second-leve wavelet subbands of size RC* This feature pivotal, as we decompose every wavelet subband to evenly separate an imput image into multiple wavelet patches, unlike conventional wavelet transformations that process only LLL subbands the next level. In Wavelet Mamba, we concatenate those subbands form ID sequence, apply window scanning within each subband, and slide across the sequence for feature extraction. The window scanning is inspired by CNN kernel proposed in (28) with two window sliding directions: left to right and top 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 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 Sspatial-frequence Mamba in Fig.
fusion layer. Given fa and fa are spatial and wavelet features obtained by sweep and window sean, We combine these features using cross-attention fusion layer as follows:
qa k.s Ug = Linear(fe),
Gur kws Vw = Linear(fw))
fout- Linear Linear (Concat(Attn(qe. kw Uw) Attn h: Ut))):
More specifically, we first compute each feature's query (g-), key (k:)), and value (v:,) using linear layers. To fuse the information between spatial and wavelet features, 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 linear projection to obtain the output feature Sout (see the last subfigure in Fig. D
do cross-attention by
Conditional Mamba. Unlike attention modules, conventional Mamba has explicit mechanism to inject input conditions into its flow. We propose a simple technique that enables Mamba to take in any conditional imput via initializing the very first hidden state with the embedding 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 be rewritten as below:
ho Ah- + B
he
Abi,1 + B
you Cha
1/ Ch.
In conventional Mamba, h-1 is set to zero as there is no previous state at the beginning. Here, we set h_1 Linearp(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 ss global information, similar to vision transformers [10 59). For class- conditional generation task, we use a class embedding to condition on Mamba. Conditional Mamba is enabled by default in DiM blocks (Fig.
3.4 Globally-shared attention block
Since Mamba is better than transformer at long-range dependency [1] 15 but weaker than trans- former at in-context learning [46), we propose a hybrid mambattransformer architecture which favors both these properties as in recent work Jamba [35). Motivated by Zamba (]]. we introduce the globally-shared transformer (attention) block. This shared attention block is added after each of four DiM blocks as shown in Fig. 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 intro- duced 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 's order dependence while significantly reducing the FID with very few additional parameters.
4 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), ISSUN Church and Ima (5). For the sampling method, we follow 44
25
to use adaptive ODE solver "dopri5' for evaluation. We assessed
7

---

Model
NFEL
FID
Recall"
Epochs
Ours
61
4.62
0.52
225
Zigmat
85
63
7.66
0.40
400
LFM-8
0.46
LDM-4 51 (ADM)
500
5.11
0.49
600
Model
NFEL
FID,
Recallf
Epechs
(ADM)
Ours
82
6.09
0.46
165
LDM-8 (DIT)S
250
10.21
500
LFM-8
94
6.35
0.41
500
LSGM CB
23
7.22
1K
6,40
0.35
WiveDDil 49
2
5.94
0.37
500
DDGAN
64
2
8.43
0.33
400
RDUOT 8
5.60
0.38
(c) Size 512 x 512.
Score SDE 58
4000
7.23
6.2K
(a) Size 256 x 256. is our reproduced result and $ is adopted results from LFM paper.
10
MDD11K
0k5
10t
300
1550
200
250
300
Eboch
(d) Comparison of FID-10K over training epochs.
(b) Qualitative results.
Figure 4: Result of our model versus others upon CelebA-HQ. is our reproduced result based on Zigma's 's official code, and is an adopted result from LFM paper.
Model
NFF
FID.
Recallf
Epochs
Our
73
3.76
0.56
395
DIFFUSSM
AS
250
3.94
LFM-8 55
90
5,54
0.48
500
LDM 51
250
4.02
0.52
400
WileDil
5.06
0.40
500
DDPM 22
1000
7.89
640
StylGANN M0
4,21
StyloGAN2
0.36
13K
(b) Qualitative results. Figure 5: Result of our model versus others upon LSUN Church 256 x 256 dataset.
(a) Quantitative results.
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 training iterations per second to further illustrate the model efficiency.
4.1 Image Generation
Table I: Class-conditional image generation on ImageNet 256 x 256 dataset.
On the CelebA-HQ dataset at resolutions of 256x256 and 512x512 (Fig. 4. our method achieved state-of-the-aan FID scores of 4.62 and 6.09, respectively, surpassing the scores reported in recent studies. Furthermore, DiMSUM demon- strated superior recall scores, indicating a greater diversity the generated samples compared to other methods. This result is particularly impres- sive, 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 GANN-ased approaches. Moreover, our method recorded the highest recall metric of 0.56, significantly exceed- ing the recall of GAN methods, thereby underscoring its robustness generating diverse and high-quality images.
Model
Recallt
Purams
elbers
Epoch
Ours
$10
Oun-GG
460M
936K
SSM bused
DIFFUSSM-XC XL 65
9.07
0.64
673M
2578K 256
515
SLG
2.28
0.56
673M
28 25788
256
515
UNat-boued LDM-4
0.62
400M
200
LDM-4G
400M
178K
Transfrmen based
DIT-122 AX
23.53
458M
400K x 256
801
DIT-XL/2
9.62
0.67
675M
7000K
256
1.4K
DIT-XL/-GG
2.27
0.57
675M
7000K
256
1.4K
SIT-X1-2
255
1.4K
STT-L22GG
256
1.4K
BigGaa-deep 4
6.9
0.28
160M
StykGGAN-X. 54
2,30
0.53
166M
25000K x 256
4K

---

FiD
Recallf
Purams
GFLOPs
#Lv
FIDJ
Recalft
GFLOPs
Baseline
6.19
0.46
413M
51.65
Conditional Mamba
5.27
0.49
446M
$1.69
5.09
0.50
84.48
Wavelet Mamba (w/ concat)
5.87
0.47
394M
56.54
4.65
0.52
84.49
Cross Atlention fusson layer
436M
5.23
0.49
84.50
Shared transformer block
4.65
0.52
459M
84.49
(b) Wavelet levels.
(a) Network components.
FID4
Recalif
GHLOPs
FID4
Recall
Params
GFLOPs
Linear
6,00
0.47
411M
60.83
Atlention
5.05
0.50
461M
73.71
DCT
5.53
0.50
436M
67.33
CAFL (swap q)
4,92
0.50
436M
EinFFT
5.63
0.48
371M
66.96
CAFL (swap k)
5.45
0.48
436M
67.27
Wavelet
4.92
0.50
436M
62.42
(c) Fusion layers. CAFL. means Ettss-AAtention Fusion Layer.
(d) Frequency types.
Order
FID
Recallt
iterss
Conditional Mamba Only
FID.
Recallf
Params
6,39
0.44
2.06
Conditional Mamba Only
Sweep-4
5.27
0.49
2,06
Sweep-8
5,53
0.48
1.97
GS
5,40
0.49
469M
78.30
Zigrag.8. 8
Jpeg-8 Window
6,26
0.45
1.97
Mamba
10.88
0.36
2.05
ldp
5.08
0.49
397M
63.37
GS
4.65
0.52
454M
84.49
Spatial-frequency Mamba
Sweep-4 4 Sweep-4 5.41
0.49
1.54
(f) Transformer layers. GS stands for globally shared. Idp stands for Independent.
Sweep-4 4| Window
4.92
0.50
1.54
(e) Scanning orders.
Table 2: Ablation studies CelebA-HQ 256 x 256 dataset at epoch 250.
On the ImageNetlk 256 dataset, methodology attains formidable FID of 2.26 using guidance scale of 1.4, surpassing the DiT models across comparable and larger model configurations, such as DI/-XL/2. This performance superiority extends to other benchmarks, including the SSM-based DIFFUSSM-XL-C model. Although our model yields results similar to those of the SiT model, our model is approximately 30% smaller in size.
4.2 Training convergence
As reported in Tables 1 5a. our method requires less training epochs/iterations to reach the optimal performance compared the baseline approach, implying strong and fast convergence. To better illustrate the training convergence comparison, we illustrate in Fig. 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 significantly lower value than the other methods like LFM[24], LDM(5T) and DiT[48). This rapid descent is particularly evident within the first 150 epochs, after which our method maintains 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 more stable trajectory without significant oscillations between training epochs.
4.3 Ablation of network design
In this section, we ablate the design choices for our network, using experiments 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 proposed conditional Mamba, the FID score is improved from 69 to 5.27, and the same trend observed for recall. Meanwhile, adding Wavelet Mamba followed by 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 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 CelehA-HQ 256 dataset. We argue that the choice of wavelet levels should be based on the input resolution. An input image of size 256 x 256, for instance, is encoded to a compact latent of size 32 x 32, which further patchified by 2 to the small size of 16 x 16. Hence, applying wavelet levels results in extremely small wavelet subbands of size 2 x 2, leading to reduced performance.
also consider different types of frequency techniques like DCT and Fourier Transform (Fig. 2d). For DCT, we propose a multi- JPEG scanning strategy (illustrated in Fig. [ based on JPEG compression [63] instead of the window scanning. For Fourier Features, we directly adopt the FinFFT block from SiMBA 47 the performance drops compared with the default choice of wavelet.
Alternative frequency transform. Apart from wavelet transform,
In either case,
Transformer layer: In Fig. 21] 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 Opatial-Frequent
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 decline of 0.43 in FID and 0.03 in Recall.
4.4 Ablation of scanning orders
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, 8, Zigzag-8, 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 better outcome than conventional sweep-4 scanning.
5
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 stte-oofthhe-as 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- cext-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 rext-to-image generation and multimodal diffusion. While there is a risk that our architecture could be misused for malicious purposes, posing social security challenge, we are confident that this risk can be mitigated with the recent development of security rreated research. Hence, the positives can outweigh the negatives, rendering the concern minor.
While method outperforms other diffusion baselines in generation quality and training conver- gence, 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.
10

---

Acknowledgements
Research partially funded by research grants Prof. Dimitris Metaxas from NSF: 2310966, 2235405, 2212301, 2003874, 1951890, AFOSR 23RT0630, and NIH 2R01HL
References
[1] Zyphra unveils zamba: A compact 7b ssm hybrid model, 2024. https//wwww.cybrr..oommmmaaa [Online release: April 16th 2024].
[2] F. Bao, S. Nie. K. Xue, Y. Cao, C. Li, H. Su, and J. Zhu. All are worth words: A vit backbone for diffusion models. In CVPR, 2023.
[3] Betker, G. Goh, L. Jing, T. Brooks, J. Wang, L. Li, L. Ouyang, J. Zhuang, J. Lee, Y. Guo, al. Improving image generation with better captions. Computer Science. https:///ddn openai. com/papers/dallle-3 pdf. 2(3):8, 2023.
[4] A. Brock, J. Donahue, and K. Simonyan. Large scale GAN training for high fidelity natural image synthesis. In International Conference Learning Representations, 2019.
[5] Q. Dao, H. Phung, B. Nguyen, and A. Tran. Flow matching in latent space, 2023.
[6] Q. Dao, B. Ta, T. Pham, and A. Tran. Robust diffusion gan using cemi-unbalanced optimal transport. arXiv preprint arXiv:2311.17100 2023.
[7] T. Dao. 2lashAttention-2: Faster attention with better parallelism and work partitioning. 2023.
[8] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, 2022.
[9] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, Uszkoreit, and N. Houlsby. An image is worth 16x16 words:
Transformers for image recognition at scale. ICLR, 2021.
[10] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[11] Esser, S. Kulal, Blattmann, Entezari, J. Müller, H. Saini, Y. Levi, D. Lorenz, A. Sauer, F. Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. arXiv preprint arXiv:2403. 032006 2024.
[12] W. Fedus, Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-3, 2022.
[13] J. S. Fischer, M. Gui, P. Ma, N. Stracke, S. A. Baumann, and B. Ommer. Boosting latent diffusion with flow matching, arXiv preprint 21X.0223 2023.
[14] S. Gao, P. Zhou, M.-M. Cheng, and Yan. Masked diffusion transformer is strong image synthesizer. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2023.
[15] A. Gu and T. Dao. Mamba: Linear time sequence modeling with selective state spaces, arXiv preprint arXiv:23/2.00752 2023.
[16] A. Gu, T. Dao, S. Ermon, A. Rudra, and C. Ré. Hippo: Recurrent memory with optimal
polynomial projections. Advances neural information processing systems, 33:1474414887, 2020.
[17] A. Gu, K. Goel, A. Gupta, and C. Ré. On the parameterization and initialization of diagonal state space models. Advances Neural Information Processing Systems, 15335971-35983, 2022.

---

[18] A. Gu, K. Goel, and C. Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022.
[19] M. Gui, J. S. Fischer, U. Prestel, Ma, D. Kotovenko, O. Grebenkova, S. A. Baumann, V. T. Hu, and B. Ommer. Depthfm: Fast monocular depth estimation with flow matching, arXiv preprint arXiv:2403. 1378888 2024.
[20] H. Guo, J. Li, T. Dai, Z. Ouyang, X. Ren, and S.-T. Xia. Mambair: A simple baseline for image restoration with state-space model. arXiv preprint arXiv:2402. /56488 2024.
[21] Ho. W. Chan, C. Saharia, J. Whang, R. Gao, A. Gritsenko, D. P. Kingma, B. Poole, M. Norouzi, D. J. Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:22/0.02300: 2022.
[22] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851 2020.
[23] Ho, T. Salimans, A. A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet. Video diffusion models. In ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022.
[24] T. Hu, D. W. Zhang, P. Mettes, M. Tang, D. Zhao, and C. G. Snoek. Latent space editing in transformer- based flow matching. In AAAI, 2024.
[25] V. T. Hu, S. A. Baumann, M. Gui, O. Grebenkova, P. Ma. J. Fischer, and B. Ommer. Zigma: Zigzag mamba diffusion model. arXiv preprint arXiv:24033. 10:::200333880,, 2024.
[26] V. T. Hu, W. Yin, P. Ma, Y. Chen, B. Fermando, Y. M. Asano, E. Gavves, P. Mettes, B. Ommer, and C. G. Snoek. Motion flow matching for human motion synthesis and editing. arXiv preprint arXiv:2312.08895 2023.
[27] H. Huang, z. li, R. He, Z. Sun, and T. Tan. Introvae: Introspective variational autoencoders for photographic image synthesis. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Desa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
[28] Huang, X. Pei, You, Wang. C. Qian, and C. Xu. Localmamba: Visual state space model with windowed selective scan. arXiv preprint 2xXx::2433.09338 2024.
[29] T. Karras, M. Aittala, J. Hellsten, S. Laine, J. Lehtinen, and T. Aila. Training generative adversarial networks with limited data. In Advances in neural information processing systems, 2020.
[30] T. Karras, S. Laine, and T. Aila. A style-based generator architecture for generative adversar- ial networks. In Proceedings of the IEEE/CVF conference on computer vision and puttern recognition, pages 4401-4410, 2019.
[31] D. P. Kingma and R. Gao. Understanding diffusion objectives as the ELBO with simple data augmentation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[32] T. Kynkäänniemi, Karras, S. Laine, J. Lehtinen, and T. Aila. Improved precision and recall metric for assessing generative models. Advances in Neural Information Processing Systems, 32, 2019.
[33] Li, H. Singh, and A. Grover. Mamba-nd: Selective state space modeling for multi-dimenniionll data. arXiv preprint arXiv:4402.05892, 2024.
[34] D. Liang, X. Zhou, X. Wang, X. Zhu, W. Xu, Z. Zou, X. Ye, and X. Bai. Pointmamba: A simple state space model for point cloud analysis. arXiv preprint wxXi:22402.10739 2024.
[35] O. Lieber, B. Lenz, Bata, G. Cohen, J. Osin, I. Dalmedigos, E. Safahi, S. Meirom, Y. Belinkov, S. Shalev-Shwartz, al. Jamba: A hybrid mrams language model. arXiv preprint arXiv:2403.19888 19887, 2024.
[36] Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2023.
12

---

[37] H. Liu, X. Jiang, X. Li, Z. Bao. D. Jiang, and B. Ren. Nommer: Nominate synergistic context in vision transformer for visual recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12073-12082, 2022.
[38] H. Liu, X. Jiang. X. Li, A. Guo. Y. Hu, D. Jiang, and B. Ren. The devil is in the frequency: Geminated gestalt autoencoder for self-supervised visual pre-training. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 1649-1656, 2023.
[39] J. Liu, H. Yang. H.-Y. Zhou, Y. Xi, L. Yu, Y. Yu, Y. Liang, G. Shi, S. Zhang, H. Zheng, et al. Swin-umamba: Mamba-based unet with imagenet-based pretraining. arXiv preprint arXiv:2402.03322 2024.
[40] X. Liu, C. Gong, and Q. Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. Feb. 2023.
[41] Y. Liu, Y. Tian, Y. Zhao, H. Yu, L. Xie, Y. Wang, Q. Ye, and Y. Liu. Vmamba: Visual state space model. ArXiv, 04/22001101666 2024.
[42] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo. Swin transformer: Hierarchieal vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
[43] J. Ma, F. Li, and B. Wang. U-mamba: Enhancing long-range dependency for biomedical image segmentation. arXiv preprint arXiv:2401.04722, 2024.
[44] N. Ma, M. Goldstein, M. S. Albergo, N. M. Boffi, E. Vanden-Eijnden and S. Xie. Sit: Exploring flow and diffusion-basee generative models with scalable interpolant transformers. arXiv preprint arXiv:2401.08740 2024.
[45] M. F. Naeem, S. J. Oh, Y. Uh Y. Choi, and Yoo. Reliable fidelity and diversity metries for generative models. ArXiv, abs/2002.099977, 2020.
[46] J. Park, J. Park, Z. Xiong, N. Lee, J. Cho, S. Oymak, K. Lee, and D. Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks. arXiv preprint 202.:4444. 2024.
[47] B. N. Patro and V. S. Agneeswaran. Simba: Simplified mamba-based architecture for vision and multivariate time series. arXiv preprint arXiv:2403. 2024.
[48] W. Peebles and S. Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022.
[49] H. Phung, Q. Dao, and A. Tran. Wavelet diffusion models are fast and scalable image generators. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition (CVPR), 2023.
[50] B. Poole, A. Jain, J. T. Barron, and B. Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209. 0et.22299.4998. 14988, 2022.
[51] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[52] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. S. Bernstein, A. C. Berg, and L. Fei-Fei. Imagenet large scale visual recognition challenge. International Journal of Computer Vision, 115:211 - 252, 2014.
[53] A. Sauer, F. Boesel, T. Dockhom, A. Blattmann, P. Esser, and R. Rombach. Fast high resolutonon image synthesis with latent adversarial diffusion distillation. arXiv preprint arXiv: /2015, 2024.
[54] A. Sauer, Schwarz, and A. Geiger. Stylegan-xl: Scaling stylegan to large diverse datasets. In ACM SIGGRAPH 2022 conference proceedings, pages 1-10, 2022.
13

---

[55] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning. pages 2265. PMLR, 2015.
[56] J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
[57] Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
[58] Y. Song, J. generative modeling through stochastic differential equations. International Conference on Learning Representations, 2020.
D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score- based
[59] H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, and H. Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pages 10347-10357. PMLR, 2021.
[60] A. Vahdat, K. Kreis, and J. Kautz. Score-based generative modeling in latent space. In
A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Informa- lion Processing Systems, 2021.
[61] V. Voleti, C.-H. Yao, M. Boss, A. Letts, D. Pankratz, D. Tochilkin, C. Laforte, R. Rombach, and V. Jampani. Sv3d: Novel multi-view synthesis and 3d generation from a single image using latent video diffusion. arXiv preprint arXiv:2403. 12008, 2024.
[62] 7. Wang. C. Lu, Y. Wang, F. Bao, C. Li, H. Su, and J. Zhu. Prolificdreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. Advances in Neural Information Processing Systems, 36, 2024.
[63] W. Xiao, N. Wan, A. Hong, and X. Chen. A fast jpeg image compression algorithm based on det. In 2020 IEEE International Conference on Smart Cloud (SmariCloud), pages 106-110, 10, 2020.
[64] Z. Xiao, K. Kreis, and A. Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. In International Conference on Learning Representations, 2022.
[65] J. N. Yan, J. Gu, and A. M. Rush. Diffusion models without attention. arXiv preprint arXiv:2311.18855 2023.
[66] M. Yang, 7. Wang, 7. Chi, and W. Feng. Wavegan: Frequency-aware gan for high-fidelity few-shot image generation. In European Conference on Computer Vision, pages 1-17. Springer, 2022.
[67] T. Yao, Y. Pan, Y. Li, C.-W. Ngo, and T. Mei. Wave-vit: Unifying wavelet and transformers for visual representation learning. In European Conference on Computer Vision, pages 328-345. Springer, 2022.
[68] F. Yu, Y. Zhang, S. Song, A. Seff, and J. Xiao. Lsun: Construction of large-scale image dataset using deep learning with humans in the loop. ArXiv, 055/11066033665. 2015.
[69] X. Yuan, L. Li, J. Wang, Z. Yang. K. Lin, Z. Liu, and L. Wang. Spatial-frequem u-net for denoising diffusion probabilistic models. arXiv preprint arXiv:2307.14644 2023.
[70] T. Zhang, X. Li, H. Yuan, S. Ji, and S. Yan. Point cloud mamba: Point cloud learning via state space model. arXiv preprint arXiv:2403.0076. 2024.
[71] Z. Zhen, Y. Hu, and Z. Feng. Freqmamba: Viewing mamba from a frequency perspective for image deraining. arXiv preprini arXiv:2404.09477, 2024.
[72] H. Zheng, W. Nie, A. Vahdat, and A. Anandkumar. Fast training of diffuusion models with masked transformers. arXiv preprint arXiv:2306. 09300, 2023.
[73] Zhengcong, Mingyuan, Y. Changqian, and H. Jusnshi. Scalable diffusion models with state space backbone. arXiv preprint, 2024.
14

---

[74] L. Zhu, B. Liao, Q. Zhang, X. Wang. W. Liu, and X. Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint arXiv:2401.09417, 2024.
15

---

Table 3: Scaling DIMSUM's parameters LSUN Church
Model
FID
Epochs
Params
DiMSUM-L/2
3.76
395
460M
DiMSUM XL/2
3,45
340
718M
DIFFUSSM
3.94
673M
StyleGAN
4.21
StyleGAN2
3.86
13K
A Training details
We provide network configuration details table (bl and hyperparams in table 6a For the interpolent form of the forward process eq. (I), we opt to use generalized VP form, which is proposed in SiT paper [44]. Particularly, this form defines a₁, = cos (ttl and o; sin (t). The choice of transformer block, we adopt DiT block 48 where we replace the original MLP layer with Gated
MLP layer, similar to
Besides, we illustrate the sweep scanning strategy in Fig. 9 for better comprehension. We also visualize the Jpeg scanning orders in Fig. iinspired by the Jpeg compression algorithm (63 and the multi-direction scanning from ZigMa paper [25].
B Discussion
Advantages of DIMSUM. First, it's important highlight that our method outperforms both DiT and SiT while requiring less than third the training iterations, achieving the best FID score of 2.11. Compared to other state-space diffusion models, method outperforms DIFFUSSM-XL. considering similar training duration. Notably, our method also uses 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 L.MM-dominted 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 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 achieved with a parameter count comparable (or smaller) to existing models. Refer to Table and [ This suggests substantial room for further enlargement of our model's parameters, which 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 Table 2c contributing to the SoTA result of our method.
16

---

Table 4: Speed and GFLOPs comparison. Single-sample generation was performed. and all tests were conducted on an NVIDIA A100 40GB GPU.
Method
Time
MEM
Params
GFlops
256 (latent size:
32
32)
Ours-L/2 DIT-L/2
2.20s
2.42G
460M
84,49
3.80s
2.30G
458M
80.74
512 (latent size:
64
64)
Ours-L/2 DIT-L/2
2.86s 4.78s
2,46G
461M
337.48
2.34G
459M
361.14
C Speed Analysis
Memory and GFLOPs. The results, as shown in the table reveal that DiMSUM-L/2's memory usage is slightly higher than its counterpart. This increase 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 x 256 images, DiMSUM produces about 4% more GFlops than DiT. However, 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 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.
Analysis. Observing the Memory and GFLOPS in table table 4 it's true to claim DiMSUM-L/2 shows slower inference speed compared to its counterpart for 256x256 images using the same NEE (due to bigger GFLOPS), however, there are two erucial 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 SiT [44] and Zigma [25). This solver dynamically adjusts the NFE based on the initial noise and diffusion model characteristics, using the minimum NEE necessary to achieve optimal image quality. Notably, DiMSUM requires fewer NFE to meet the dopri5 stopping condition while still achieving significantly better FID score than DIT. We hypothesize that our proposed hybrid architecture converges better solution with less curvature, enabling high-fidelity image production with fewer NFEs.
17

---

DIT
102
FID-10K
101
0
100
200
300
400
500
Epoch
Figure Training curve of DiT after training longer on CelebA, supporting Fig. 4d in the manuscript.
Table Additional ablation of scanning orders involves using window scanning for wavelet blocks and examining the sensitivity of DiMSUM to different scanning orders for spatial blocks. For simplicity, no shared-transformer block is used.
Order
freq spatial
spatial
Bi
5.47
6.39
Jpeg
5.74
6.26
sweep-8
5.23
5.53
zigma-8
5.71
6,26
sweep-4
4.92
5.28
With these two points, we emphasize upon the importance of our proposed architecture, rather than just the benefits given by the flow matching framework.
D More quantitative results
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. 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 claims regarding the efficiency of frequency information, we conducted 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 T 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 NFF, 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 improvement in the FID scores. This behavior
18

---

6.05
Heun
Euler
6.00
5.95
5.90
FID-10K
5.85
5.80
5.75
5.70
5.65
50
100
150
200
250
500
Step
Figure 7: FID-10K of varying NFE on CelebA 256.
Table 6:
and network config of our DiMSUM network.
(a) Hyper- pareerppraamters
(b) Network config
CelebA 256 & 512
Church
ImageNet
Config
Value
Learning rate
le-4
5e-5
le-4
Depth
16
01,32 Batch size Droppath Max-grad-norm
0.9, 0.999 64 & 32
0.9, 0.999
0.9, 0.999
Hidden size
1024
128
704
Patch size
Use learnable absolute
True
2
2
positional embedding
Label-dropout
0.
0.
0.15
Attention every k layer
4
Epochs
250 & 165
395
320
GPUs (A100)
2
Train days
0.89 3.2
3.42
12
consistent with the observation of the dow-matching Fig. of FM paper 36 which has been shown to require fewer NFEs compared to other SDE-based methods.
E More qualitative examples
present our uncurated generated samples of CelebA-HQ 256 Fig.
CelebA-HQ 512 Fig.
LSUN Church in Fig. [1 and ImageNet in Fig. 18 13
19
20
19

---

Horizontal
Vertical
Horizontal
Vertical
ib
ib
ib
ii
Reverse Horizontal 0
Reverse Vertical
Reverse Horizontal
Reverse Vertical 1
Figure 8: Illustration of Jpeg-8 scanning orders.
Sweep-8
Sweep-4
Bidirectional
Horizontal 0
Vertical 0
Horizontal 1
Vertical
ib
i6
ib
i6
Reverse Horizontal
Reverse Vertical
Reverse Horizontal
Reverse Vertical
Figure 9: Illustration of Sweep scanning orders.
F
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our contributions are accurately reflected in the abstract and in the summary of introduction section section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
20

---

n
Figure 10: Uncurated generated samples of CelebA-HQ 256.
21

---

shu stck
ittestookk
Figure 11: Uncurated generated samples of LSUN Church 256.
22

---

Figure 12: Uncurated generated samples of CelebA-HQ 512.
23

---

P
altoooseeeco bon Cerone
Figure 14: Uncurated generated samples of ImageNet 256 class 387 (lesser panda, red panda, panda, bear cat, cat bear, Ailurus ful- gens) with cfg=4.0.
Figure 13: Uncurated generated samples of ImageNet 256 class 360 (otter) with cfg=4.0.
24

---

Figure 15: Uncurated generated samples of ImageNet 256 class 108 (sea anemone) with cfg=4.0.
Figure 16: Uncurated generated samples of ImageNet 256 class 393 (anemone fish) with cfg=4.0.
25

---

Figure 17: Uncurated generated samples of ImageNet 256 class 417 (balloon) with cfg=4.0.
Figure 18: Uncurated generated samples of ImageNet 256 class 974 (geyser) with cfg=4.0.
26

---

Figure 20: Uncurated generated samples of ImageNet 256 class 21 (bald eagle, Amer- ican eagle, Haliaeetus leucocephalus) with cfg=1.5.
Figure 19: Uncurated generated samples of ImageNet 256 class 250 (Siberian husky) with cfg=1.5.
27

---

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
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
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create 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 few datasets with 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, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails 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, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
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 cross- referenced.
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 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.
4. Experimental Result Reproducibility
28

---

Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that 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 and section
Guidelines:
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 a specific model and empirical evaluation, 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 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 (c.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., 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 have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions 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.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
mablic/guidess http://guuiies//oeubbmmiiis CoodeeSubissiiooool111..
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https //nips.co/ http://wwwi/peeeaeeee//eettaaaaa
for more details.
29

---

The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
reproducible, they
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, we provided them in Implementation part of section and Training details in appendix
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: No, because the evaluation process takes hours on GPU and is very ex- pensive to compute with multiple seeds. Importantly, 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 accompanied by emror bars, confi- dence 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 the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
30

---

Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: Yes, this is presented in appendix A
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Yes, we totally confirm to it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws 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
Guidelines:
The answer NA means that there no socictal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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 partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate 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 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.
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 system learns from feedback over time, improving the efficiency and accessibility of ML).
31

---

11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [Yes]
Justification: Yes, our code will be released with Licences.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Yes, we have properly cited the code, data and models used.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasete has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Our released source code is accompanied by instruction document Details about implementation can be found section section 4
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
32

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjeets
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere 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.
33


Reviews:
[{"review_id": "UJQm4BvNem", "reviewer": "Reviewer_RLU7", "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.", "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": "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.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- This paper is easy to follow.\n\n- The implementation details and the codes are included in this submission.\n\n- 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.\n\n- 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. \n\n- The comparison of the real inference speed between the proposed framework and transformers is not included in this paper. \n\n- 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?\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1. Our architecture might comprise many components such as Mamba, Transformer, and frequency processing. However, each proposed component is well-motivated and vital to the overall framework, as shown in Table 2a (in submission). The simple Mamba network, without frequency processing and transformer blocks, cannot surpass the performance of the transformer. As shown in section 2.1 submission, the scanning order of the Mamba block has a significant effect. To mitigate its impact, we integrate wavelet frequency with window scanning into architecture. As explained in the global response, wavelet scanning, along with window scanning, can better capture the global information of each wavelet subband, enhancing the overall global information across the whole image. In addition, our model could benefit from learning both high and low frequency [1]. With wavelet frequency and window scanning, our framework reduces FID from 5.27 to 4.92 for CelebA-HQ (Table 2e in submission). Furthermore, recently, hybrid architectures of Transformer and Mamba [2, 3] have demonstrated SoTA performance in NLP tasks. Motivated by these works, we measure the performances of hybrid Transformer-Mamba in vision generative. As shown in Table 2f submission, both independent and shared transformer improves performance, indicating hybrid model's effectiveness.\n\n[1]: Hao Phung, Quan Dao, and Anh Tran. \"Wavelet diffusion models are fast and scalable image generators.\" In CVPR 2023.\n\n[2]: Zyphra unveils zamba: A compact 7b ssm hybrid model, 2024. https://www.zyphra.com/zamba\n\n[3]: O. Lieber et al. Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887, 2024\n\n2. We acknowledge that the presentation of our paper should be clearer and appreciate your feedback. To clarify, the term \"manually-defined scanning orders\" refers to the heuristic scanning order in the Mamba blocks. In Section 2.1, from lines 94 to 107 (in submission), the scanning order of Mamba is defined as hand-crafted in previous works. Throughout the paper, we mentioned some of these hand-crafted orders, including \"bi\", \"window\", \"sweep-4/8\", and \"zigzag-8\". These scanning orders are illustrated in Figs 2, 5, and 6 in submission.\nAlthough our approach does not completely address the impact of scanning order, our window scanning strategy in wavelet space could efficiently capture the global information within each frequency sub-bands which improves global content of generated image, as explained in global response. Please refer to table 3 (in rebuttal file), our wavelet and window scan component successfully reduce FID score compared to spatial scanning only. Therefore, we consider our frequency scanning method to be synergistic with image space scanning, as it effectively handles the global information of generated image. Besides, completely resolving the limitations of manually-defined scanning is an interesting topic for future research.\n\n3. DIMSUM-L/2 achieves 2.2 seconds latency compared to DIT-L/2 with 3.8 seconds (see table 2 in rebuttal file).\n\n4. In terms of model's type, we would like to revise the terms \"Mamba-based\" and \"Mamba model\" to \"a hybrid Mamba-Transformer network\" for better comprehension. In terms of spatial-frequency scan for transformer, we modify DiT architecture to further incorporate wavelet frequency information similar to DIMSUM. After training spatial-frequency transformer on CelebA-HQ for 1000 epoch, we found that the spatial-frequency transformer failed to learn the human face distribution. The output model only generates noise-pattern images. More rigorous analysis of model is beyond the scope of paper and we would like to leave it for future investigation. \n\n5. We acknowledge that the use of term 'scalable' may have caused some confusion. In the broader context, deep learning’s scalability may refer to capacity to handle various datasets and computational resources while producing correct results in an acceptable period of time [1].\n\nOur model aligns well with this definition of scalability. Our experiments demonstrated that: \n- The model adapts to multiple datasets with minimal hyperparameter tuning (similar to DiT paper).\n- It achieves competitive performance metrics compared to other methods.\n- Inference speed is also faster, as shown in table 2 (in rebuttal file).\n\nOur SoTA results are achieved with a parameter count comparable (or even smaller) to existing models. Refer to Figure 3a and 4a (in submission) and Table 1 (in rebuttal file). This suggests substantial room for further enlargement of our model's parameters, which we anticipate will yield even better improvements across various and bigger datasets (see table 4 in rebuttal file).\n\n6. Our plot intentionally stops at epoch 300, demonstrating the model's capability to converge faster than other methods. DiT does perform well on CelebA-HQ but takes more than 500. Figure 4 (in rebuttal file) illustrates the convergence of DiT at the later epoch. It's noted that DiT and LFM (in figure 3d submission) use the same DiT-L/2 architecture. While DiT uses diffusion loss, LFM uses flow matching. LFM converge faster than DiT. However, our model with flow matching loss, demonstrates even faster convergence than LFM, indicating our architecture enhances convergence rate.\n\n7. Social Impact section will be added in revised paper (see global response).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xCB4fGCykH", "reviewer": "Reviewer_eitf", "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.", "questions": "See the weaknesses part.\n\nMeanwhile, 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": "See the weaknesses part.\n\nAll of my concerns have been addressed. I will increase my score.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe 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. \n2.\tThe paper is well-written and easy to follow.", "weaknesses": "1.\tSince 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. \n2.\tStronger baselines (such as [1]) and the experiments on more datasets (such as conditional generation on MSCOCO) and resolutions are required. \n3.\tFrom 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. \n\n\n[1] Chen, Junsong, et al. \"Pixart-$\\alpha $: Fast training of diffusion transformer for photorealistic text-to-image synthesis.\" arXiv preprint arXiv:2310.00426 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nTo substantiate our claims regarding the efficiency of frequency information, we conducted a comprehensive ablation study. This study utilized four scanning orders: (1) bidirection, (2) jpeg-8, (3) sweep-8, and (4) zigzag-8. We trained models on CelebA-HQ at 256x256 resolution for 250 epochs, comparing performance when applying these scanning strategies to: a) spatial domain only or b) both spatial and frequency domains\n\nTable 3 of rebuttal file presents the results of this ablation study. Notably, integrating frequency domain information across all four scanning strategies led to significant performance improvements. This consistent enhancement across various scanning methods provides strong evidence for the effectiveness of our approach in leveraging frequency information.\n\n2. Stronger baselines (such as Pixart-alpha) and the experiments on more datasets (such as conditional generation on MSCOCO) and resolutions are required.\n\nWe appreciate the suggestion to explore text-to-image (T2I) generation. However, training competitive text-to-image models presents significant challenges:\n\n- Computational constraints: As a small lab, we lack the extensive resources required for training large-scale T2I models.\n- Dataset availability: Recent closure of datasets like LAION-5B limits access to suitable training data for T2I tasks and also would be hard to compare fairly with previous works using those datasets.\n\nIt's worth noting that focusing on label-to-image tasks, as we have done, is common in diffusion model research like DiT, EDM, and Consistency Models. Many influential papers in this field have established their contributions through experiments on label-to-image before expanding to text-to-image.\n\n3. 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.\n\nWe 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.\n\nTo address this challenge, we proposed a more sophisticated fusion method using Cross-Attention layers between these spaces. This approach enables the model to effectively combine information from both domains, leveraging their strengths while mitigating potential conflicts. Hence, this fusion technique can enhance the FID from 5.87 to 4.92 in Table 2c submission, contributing to the SoTA result of our method.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KBlK9fcS2b", "reviewer": "Reviewer_nhSz", "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.", "questions": "as above", "limitations": "as above", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Overall, I appreciate this paper. It selects the most promising technical stacks—combining flow matching and Mamba—to address generative tasks.\n- The paper is well-organized, featuring valuable methodology design and experimental ablations.\n- 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.\n- 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].\n- In Line 176, I'm curious as to why only B, C, and delta are time-dependent. Why isn't A?\n- In Line183, One paper is missed. [5]\n- In Fig 3a, what does the notation on ZigMa mean?\n- Regarding line 261, what's the rationale behind sharing those parameters of transformers? Are there any previous works indicating its effectiveness?\n- In Table 1, it would be better to display training steps. Researchers of generative models tend to compare using training steps rather than epochs.\n- 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.\n- In Fig5, is there any reference to the jpeg-8 scan path? How does it defined?\n\nSuggestions:\n\n- The previously mentioned figure or table should be placed first.\n\n[1]. Boosting Latent Diffusion with Flow Matching\n\n[2]. DepthFM: Fast Monocular Depth Estimation with Flow Matching\n\n[3]. Motion Flow Matching for Human Motion Synthesis and Editing\n\n[4]. Flow Matching for Conditional Text Generation in a Few Sampling Steps\n\n[5], Latent Space Editing in Transformer-based Flow Matching\n\nOverall, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1. In Line 106, I'm unsure why sigma brings extra computation when the scan path is larger, given that the computation is always fixed.\n\n2. 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].\n\n3. In Line183, One paper is missed.\n\nFor Q1,2,3, we admit that our wording is incorrect in this context and revise the sentence as \"In this paper, we show that too many scanning orders, e.g., sweep-8 and zigzag-8, may introduce excessive information and lead to worse performance compared to sweep-4\". We will cite the suggested papers.\n\n4. In Line 176, I'm curious as to why only B, C, and delta are time-dependent. Why isn't A?\n\nWe believe that the authors of Mamba would have a much better answer here. From our point of view, this implementation is made purely to prefer simplicity without sacrificing much performance. For more details, reviewers can check the interpretation in 3.5.2 of the Mamba paper. Simply put, time-dependent delta $\\Delta_t$ is sufficient to ensure the selection mechanism for $\\bar{A}$ via $\\exp(A * \\Delta_t)$. \n\n5. In Fig 3a, what does the notation on ZigMa mean?\n\nWe are sorry for the missing information; we will update the caption in the revised version. In this case, $\\textdagger$ is our reproduced result based on Zigma's official code, and $\\textdaggerdbl$ is an adopted result from LFM paper.\n\n6. Regarding line 261, what's the rationale behind sharing those parameters of transformers? Are there any previous works indicating its effectiveness?\n\nAs mentioned in L260, we drew inspiration from Zamba for the globally shared attention block. The primary purpose of this component is to reduce the model's parameter count while maintaining the competitive performance of the model by capturing the global dependencies between input tokens. Besides, Zigma also adopts a hybrid mamba-transformer architecture in the design to integrate the text condition for text-to-image generation effectively.\n\n7. In Table 1, it would be better to display training steps. Researchers of generative models tend to compare using training steps rather than epochs.\n\nWe acknowledge your comment and update Table 1 in the attached file.\n\n8. 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.\n\nActually, in our implementation, we don't run all scans in one forward (such as in VMamba), but we run each order alternatively following Mamba-ND and Zigma paper. Specifically, each Mamba block handles only one direction, and the four scanning orders in sweep4 are evenly distributed throughout the network's depth (number of layers). Sweep4, for example, has four scanning orders: (1) left-to-right, (2) right-to-left, (3) top-to-bottom, and (4) bottom-to-top. Assuming the network has only four layers, layer 1 handles (1), layer 2 handles (2), and so on.\n\n9. In Fig5, is there any reference to the jpeg-8 scan path? How does it defined?\n\nAs mentioned in L330, we derived the JPEG-8 scanning order from the scanning order of the JPEG compression algorithm, as described in \"A fast JPEG image compression algorithm based on DCT\". This also was significantly influenced by the Zigma paper, a pioneering work that put effort into a more sophisticated scanning order and achieved competitive results. Building upon their insights, we extended the concept to the original JPEG scanning orders. We acknowledge the valuable contribution of the Zigma paper and will provide a more detailed explanation of our scanning order development in the camera-ready version for improved clarity.\n\n10. The previously mentioned figure or table should be placed first.\n\nAs mentioned in the global response, we thank the reviewer for your attention to detail, which is extremely vital to enhance the clarity and completeness of the paper. We will try our best to reorder the figures and tables so that it would appear before being mentionned.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hOAYKblnxl", "reviewer": "Reviewer_C5pc", "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.", "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.\n\nQ2. How is the performance w.r.t more sampling steps? A curve is a good demonstration.", "limitations": "Yes, they discussed", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "+ The proposed method works on the frequency domain seems to be a good idea compared to existing mamba diffusion approaches.\n\n+ The architecture looks like getting inspiration from DiT and transferring to the Mamba style may be interesting.\n\n+ 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. \n\n+ 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. \n\n+ The scalability is important for this architecture, but it is missing, which would weaken the paper.\n\n+ 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1. Training longer looks like getting an overfitting curve, an analysis needs to be done. This behavior is contradict to Transformer-based diffusion DiT, MDT, MaskDiT, etc.\n\nThank you for pointing out, we confirm that our method does have overfitting. After converging to the best FID10K score of 5.49 at epoch 250, the model begins to exhibit an overfitting pattern. We also observe the same overfitting problem for transformer-based diffusion models like DiT and MDT. Figure 3 in the rebuttal file demonstrates that overfitting does occur in transformer-based models, albeit at a later stage, starting from epoch 800 onward, due to their slower convergence rate compared to our model.\n\nWe hypothesize that overfitting is not solely determined by model architecture but by various factors. For instance, the Improved DDPM [2] paper showed that with identical architecture in Appendix F, changing the noise scheduling from linear to cosine could lead to overfitting in same UNET architecture. Based on our empirical experiments above, we think that the overfitting effect is not only attributable to our architecture but rather a broader phenomenon in diffusion models. For MDT, overfitting emerges much later, around epoch 1100, due to its masking strategy, which could serves as a form of regularization to delay the onset of overfitting.\n\n2. 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.\n\nDIMSUM-L/2 achieves 2.2 seconds latency compared to DIT-L/2 with 3.8 seconds (see table 2 in rebuttal file).\n\n3. The scalability is important for this architecture, but it is missing, which would weaken the paper.\n\nFor this question, we provide our answers in the global response.\n\n4. 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.\n\nTo answer this question, we break it down into two parts: (1) not clearly outperforming DiT and (2) underperforming compared to MDT.\n\nFirst, it is important to highlight that our method achieves competitive performance with DiT and DIFFUSSM-XL, requiring significantly fewer training epochs—specifically, 4x fewer than DiT and 2x fewer than DIFFUSSM-XL. Notably, our method also uses a smaller network size of 460M parameters, compared to 675M of DiT while still demonstrating strong generation capacity and faster training convergence. In the updated ImageNet table provided in the PDF, we show that our method can further surpass the performance of other diffusion baselines when trained for a similar training duration as DIFFUSSM-XL. Notably, our training iterations are still less than a third of those required by DiT and SiT, yielding the best FID score of 2.11.\n\nThough our method does not outperforms the results of MDT yet, we believe that MDT is orthogonal to our proposed architecture, where a masking scheme is introduced to enhance further the contextual learning ability of diffusion models, including our method. It is an interesting topic to combine with our network for boosting model performance as mentioned in the limitation of global response.\n\n5. 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.\n\nTo clarify, Euler solver and DDIM are the same, except that Euler operates on continuous time intervals and removes the additional stochastic noise in each denoising step. In the manuscript, we followed SiT, Zigma, LFM works on using adaptive solver like \"dopri5\" for evaluation. In table 5 of the attached PDF, we show that FID scores of adaptive ODE solver and Euler solver (or DDIM sampling) are similar with small numerical difference, reconfirming the fairness of the benchmarking with DiT.\n\n6. Q2. How is the performance w.r.t more sampling steps? A curve is a good demonstration.\n\nWe plot the FID-10K scores of various NFE used for evaluation in Fig. of the attached pdf. This shows that increasing NFE beyond 250 leads to minimal or no improvement in the FID scores. This behavior is consistent with the observation of the flow-matching in Figure 7 of FM paper, which has been shown to require fewer NFEs compared to other SDE-based methods.\n\n[1] Lipman, Yaron, et al. \"Flow matching for generative modeling.\" ICLR 2023.\n\n[2] Nichol, Alexander Quinn, and Prafulla Dhariwal. \"Improved denoising diffusion probabilistic models.\" International conference on machine learning. PMLR, 2021.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We address the reviewers' comments below, referring to them as: R1(RLU7), R2(eitf), R3(nhSz), R4(C5pc). We sincerely thank all reviewers for their valuable feedback. We appreciate the positive comments on clear writing and good flow (R1, R2, R3), well-defined motivation to investigate and integrate Mamba with Transformer (R2, R3, R4), our novel approach of integrating frequency space into Mamba (R4), and thorough and extensive experiments (R1, R4).

**Scanning in Frequency Space**:

We would like to further elaborate on the motivation and effectiveness of our frequency scanning method. Mamba-based approaches in diffusion models often struggle with efficiently scanning patches while maintaining local and global 2D-spatial information. Although several works have proposed sophisticated scanning methods to address this issue, these approaches are complex and computationally expensive. Conversely, some methods have adopted local scanning (i.e., windowing) strategies [3] to improve model latency and throughput. Still, these often underperform compared to previously mentioned scanning methods as it is limited to the dependencies of nearby-pixels within window.

DiMSUM addresses these challenges by breaking the original image into frequency wavelet subbands (each subband is half of resolution of image). This approach is efficient to capture long-range frequency while preserving information across different subbands. We redesigned the window scanning, where each window corresponds to a subband of the frequency space (Fig 1 rebuttal PDF). Consequently, each window captures low/high frequency signal of original image. Since each frequency sequence is 1/4 of image sequence, negative impact of scanning is reduced. As the model progresses through different subbands, it incorporates spatial information represented at various low-to-high frequencies, adding valuable context to the process. This key difference distinguishes our method from traditional window scanning in image space.

We conducted two key ablation studies to validate this:

1. Table 3 (rebuttal PDF) shows models using both frequency and spatial domains consistently outperform spatial-only models across various scanning orders, demonstrating the robustness and effectiveness of frequency domain intergration.
2. Table 2e (submission) confirms our window-scanning strategy's efficacy in the frequency domain, achieving the best FID of $4.92$ on CelebA-HQ at $256 \times 256$ resolution.

These studies substantiate the effectiveness of our frequency domain integration and scanning strategy.

**Scalable term**

 We acknowledge that the use of the term 'scalable' may have caused some confusion, given its current association with model parameter scaling in the LLM-dominated landscape. 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 2 (in rebuttal file).

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

**Presentation**:

We appreciate the reviewers' suggestions for improving our presentation. We will address the following in our revised version:
1. Inconsistent/Misunderstanding terms (e.g., "manually-defined scanning orders," "Mamba-based" (R1)
2. Addition of social impacts and limitations sections (R1).
3. Claim about the computation of larger scan path (R3).
4. Missing citations and annotation definitions in captions (R3).
5. Figures and tables orderings (R3).
6. Updated ImageNet table with training iterations instead of epochs as suggested by (R3) and the new results of our method when trained for the same number of iterations as DIFFUSSM. Additionally, we apologize for the typo regarding the training epochs for DIFFUSSM. It should be corrected from 1.4K to 515.

**Social impacts and limitations**:

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 negative ones, rendering the concern minor.

While our method outperforms other diffusion baselines in generation quality and training convergence, we acknowledge areas for improvement. These include enhancing multiscale feature paradigms and multi-scale diffusion loss robustness [1, 2], and addressing manually-defined scanning orders. We're also intrigued by recent masked-training diffusion methods (e.g., MaskDiT, MDT), which reduce computational requirements and potentially improve global context learning. Integrating these orthogonal approaches with our work could yield further improvements.

[1] Crowson, et al. "Scalable high-resolution pixel-space image synthesis with hourglass diffusion transformers." In ICML, 2024.

[2] Hoogeboom, et al. "simple diffusion: End-to-end diffusion for high resolution images." In ICML, 2023.

[3] Huang, et al. "Localmamba: Visual state space model with windowed selective scan." In Arxiv, 2024.

Meta-review: 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.
Decision: Accept (poster)