appAppendix References (eccv) Package eccv Warning: Package ‘hyperref’ is
loaded with option ‘pagebackref’, which is *not* recommended for camera-ready
version

11institutetext:  CompVis @ LMU Munich, MCML
%TODO␣FINAL:␣Replace␣with␣your␣institution␣list.%\email{lncs@springer.com}https://taohu.me/zigma/

# ZigMa: A DiT-style Zigzag Mamba Diffusion Model

Vincent Tao Hu  Stefan Andreas Baumann  Ming Gui
Olga Grebenkova  Pingchuan Ma  Johannes Fischer  Björn Ommer

###### Abstract

The diffusion model has long been plagued by scalability and quadratic
complexity issues, especially within transformer-based structures. In this
study, we aim to leverage the long sequence modeling capability of a State-
Space Model called Mamba to extend its applicability to visual data
generation. Firstly, we identify a critical oversight in most current Mamba-
based vision methods, namely the lack of consideration for spatial continuity
in the scan scheme of Mamba. Secondly, building upon this insight, we
introduce Zigzag Mamba, a simple, plug-and-play, minimal-parameter burden, DiT
style solution, which outperforms Mamba-based baselines and demonstrates
improved speed and memory utilization compared to transformer-based baselines.
Lastly, we integrate Zigzag Mamba with the Stochastic Interpolant framework to
investigate the scalability of the model on large-resolution visual datasets,
such as FacesHQ $1024\times 1024$ and UCF101, MultiModal-CelebA-HQ, and MS
COCO $256\times 256$.

###### Keywords:

Diffusion Model State-Space Model Stochastic Interpolant

##  1 Introduction

Figure 1: Motivation. Our Zigzag Mamba method improves the network’s position-
awareness by arranging and rearranging the scan path of Mamba in a heuristic
manner.

Diffusion models have demonstrated significant advancements across various
applications, including image processing [67], video analysis [40], point
cloud processing [87], and human pose estimation [29]. Many of these models
are built upon Latent Diffusion Models (LDM)[67], which are typically based on
the UNet backbone. However, scalability remains a significant challenge in
LDMs[41]. Recently, transformer-based structures have gained popularity due to
their scalability [65, 9] and effectiveness in multi-modal training [10].
Notably, the transformer-based structure DiT [65] has even contributed to
enhancing the high-fidelity video generation model SORA [64] by OpenAI.
Despite efforts to alleviate the quadratic complexity of the attention
mechanism through techniques such as windowing [58], sliding [12],
sparsification [18, 46], hashing [19, 74], Ring Attention [53, 14], Flash
Attention [21] or a combination of them [8, 97], it remains a bottleneck for
diffusion models.

On the other hand, State-Space Models [32, 35, 31] have shown great potential
for long sequence modeling, competing with transformer-based methods. Several
methods [30, 32, 27, 69] have been proposed to enhance the robustness [92],
scalability [30], and efficiency [32, 33] of State-Space Models. Among these,
a method called Mamba [30] aims to alleviate these issues through work-
efficient parallel scanning and other data-dependent innovations. However, the
advantage of Mamba lies in 1D sequence modeling, and extending it to 2D images
is a challenging question. Previous works [96, 57] have proposed flattening 2D
tokens directly by computer hierarchy such as row-and-column-major order, but
this approach neglects _Spatial Continuity_ , as shown in Figure 1. Other
works [54, 60] consider various directions in a single Mamba block, but this
introduces additional parameters and GPU memory burden. In this paper, we aim
to emphasize the importance of _Spatial Continuity_ in Mamba and propose
several intuitive and simple methods to enable the application of Mamba blocks
to 2D images by incorporating continuity-based inductive biases in images. We
also generalize these methods to 3D with spatial-temporal factorization on 3D
sequence.

In the end, Stochastic Interpolant [3] provides a more generalized framework
that can uniform various generative models including, Normalizing Flow [16],
diffusion model [70, 39, 72], Flow matching [51, 56, 4], and Schrödinger
Bridge [52]. Previously, some works [61] explore the Stochastic Interpolant on
relatively small resolutions, e.g., $256\times 256$, $512\times 512$. In this
work, we aim to explore it in further more complex scenarios e.g., $1024\times
1024$ resolution and even in videos.

In summary, our contributions are as follows: Firstly, we identify the
critical issue of _Spatial Continuity_ in generalizing the Mamba block from 1D
sequence modeling to 2D image and 3D video modeling. Building on this insight,
we propose a simple, plug-and-play, zero-parameter paradigm named _Zigzag
Mamba (ZigMa)_ that leverages spatial continuity to maximally incorporate the
inductive bias from visual data. Secondly, we extend the methodology from 2D
to 3D by factorizing the spatial and temporal sequences to optimize
performance. Secondly, we provide comprehensive analysis surrounding the Mamba
block within the regime of diffusion models. Lastly, we demonstrate that our
designed _Zigzag Mamba_ outperforms related Mamba-based baselines,
representing the first exploration of Stochastic Interpolants on large-scale
image data ($1024\times 1024$) and videos.

##  2 Related Works

Mamba. Several works [83, 82] have demonstrated that the State-Space Model
possesses universal approximation ability under certain conditions. Mamba, as
a new State-Space Model, has superior potential for modeling long sequences
efficiently, which has been explored in various fields such as medical imaging
[89, 60, 86, 68], image restoration [34, 95], graphs [11], NLP word byte [80],
tabular data [2], point clouds [49], and image generation [25]. Among them,
the most related to us are VisionMamba[96, 57], S4ND [63] and Mamba-ND [48].
VisionMamba[96, 57] uses a bidirectional SSM in discriminative tasks which
incurs a high computational cost. Our method applies a simple alternative
mamba diffusion in generative models. S4ND [63] introduces local convolution
into Mamba’s reasoning process, moving beyond the use of only 1D data. Mamba-
ND [48] takes multi-dimensionality into account in discriminative tasks,
making use of various scans within a single block. In contrast, our focus is
on distributing scan complexity across every layer of the network, thus
maximizing the incorporation of inductive bias from visual data with zero
parameter burden.

Backbones in Diffusion Models. Diffusion models primarily employ UNet-based
[39, 67] and ViT-based [65, 9] backbones. While UNet is known for high memory
demands [67], ViT benefits from scalability [22, 17] and multi-modal learning
[10]. However, ViT’s quadratic complexity limits visual token processing,
prompting studies towards mitigating this issue [84, 21, 12]. Our work,
inspired by Mamba [30], explores an SSM-based model as a generic diffusion
backbone, retaining ViT’s modality-agnostic and sequential modeling
advantages. Concurrently, DiffSSM [90] concentrates on unconditional and class
conditioning within the S4 model [32]. DIS [25] mainly explores the state-
space model on a relatively small scale, which is not the exact focus of our
work. Our work significantly differs from theirs as it primarily focuses on
the backbone design using the Mamba block and extends it to text conditioning.
Furthermore, we apply our method to more complex visual data.

SDE and ODE in Diffusion models. The realm of Score-based Generative Models
encompasses significant contributions from foundational works such as Score
Matching with Langevin Dynamics (SMLD) by Song et al.[71], and the advent of
Diffusion Models with Denoising Score Matching (DDPMs) proposed by Ho et
al.[39]. These methodologies operate within the framework of Stochastic
Differential Equations (SDEs), a concept further refined in the research of
Song et al. [72]. Recent research strides, as exemplified by Karras et al.[42]
and Lee et al.[47], have showcased the efficacy of employing Ordinary
Differential Equation (ODE) samplers for diffusion SDEs, offering significant
reductions in sampling costs compared to traditional approaches that entail
discretizing diffusion SDEs. Furthermore, within the domain of Flow Matching
methodologies [51] and Rectified Flow architectures [55], both SMLD and DDPMs
emerge as specialized instances under distinct paths of the Probability Flow
ODE framework [72]. These models typically utilize velocity field
parameterizations employing the linear interpolant, a concept that finds
broader applications in the Stochastic Interpolant framework [3], with
subsequent generalizations extending to manifold settings [13]. The SiT model
[61] scrutinizes the interplay between interpolation methods in both sampling
and training contexts, albeit in the context of smaller resolutions such as
$512\times 512$. Our research endeavors to extend these insights to a larger
scale, focusing on the generalization capabilities for 2D images of
$1024\times 1024$ and 3D video data.

##  3 Method

In this section, we begin by providing background information on State-Space
Models [32, 35, 31], with a particular focus on a special case known as Mamba
[30]. We then highlight the critical issue of _Spatial Continuity_ within the
Mamba framework, and based on this insight, we propose the Zigzag Mamba. This
enhancement aims to improve the efficiency of 2D data modeling by
incorporating the continuity inductive bias inherent in 2D data. Furthermore,
we design a basic cross-attention block upon Mamba block to achieve text-
conditioning. Subsequently, we suggest extending this approach to 3D video
data by factorizing the model into spatial and temporal dimensions, thereby
facilitating the modeling process. Finally, we introduce the theoretical
aspects of stochastic interpolation for training and sampling, which underpin
our network architecture.

###  3.1 Background: State-Space Models

State Space Models (SSMs) [32, 35, 31] have been proven to handle long-range
dependencies theoretically and empirically [33] with linear scaling w.r.t
sequence length. In their general form, a linear state space model can be
written as follows:

$\displaystyle x^{\prime}(t)$
$\displaystyle=\mathbf{A}(t)x(t)+\mathbf{B}(t)u(t)$ $\displaystyle y(t)$
$\displaystyle=\mathbf{C}(t)x(t)+\mathbf{D}(t)u(t),$

mapping a 1-D input sequence $u(t)\in\mathbb{R}$ to a 1-D output sequence
$y(t)\in\mathbb{R}$ through an implicit N-D latent state sequence
$x(t)\in\mathbb{R}^{n}$. Concretely, deep SSMs seek to use stacks of this
simple model in a neural sequence modeling architecture, where the parameters
$\mathbf{A},\textbf{B},\mathbf{C}$ and $\mathbf{D}$ for each layer can be
learned via gradient descent.

Recently, Mamba [30] largely improved the flexibility of SSMs by relaxing the
time-invariance constraint on SSM parameters, while maintaining computational
efficiency. By employing a work-efficient parallel scan, Mamba mitigates the
impact of the sequential nature of recurrence, whereas fusing GPU operations
removes the requirement to materialize the expanded state. In this paper, we
focus on exploring the scanning scheme of Mamba in diffusion models to
maximize the use of inductive-bias from multi-dimensional visual data.

Figure 2: ZigMa. Our backbone is structured in L layers, mirroring the style
of DiT [65]. We use the single-scan Mamba block as the primary reasoning
module across different patches. To ensure the network is positionally aware,
we’ve designed an arrange-rearrange scheme based on the single-scan Mamba.
Different layers follow pairs of unique rearrange operation $\Omega$ and
reverse rearrange $\bar{\Omega}$, optimizing the position-awareness of the
method.

###  3.2 Diffusion Backbone: Zigzag Mamba

DiT-Style Network. We opt to use the framework of ViT by AdaLN [65] rather
than the skip-layer focused U-ViT structure [9], as ViT has been validated as
a scalable structure in literature [64, 10, 17]. Considering the
aforementioned points, it informs our Mamba network design depicted in Figure
4. The core component of this design is the Zigzag Scanning, which will be
explained in the following paragraph.

Zigzag Scanning in Mamba. Previous studies [81, 90] have used bidirectional
scanning within the SSM framework. This approach has been expanded to include
additional scanning directions [54, 57, 91] to account for the characteristics
of 2D image data. These approaches unfold image patches along four directions,
resulting in four distinct sequences. Each of these sequences is subsequently
processed together through every SSM. However, since each direction may have
different SSM parameters ($\mathbf{A}$, $\mathbf{B}$, $\mathbf{C}$, and
$\mathbf{D}$), scaling up the number of directions could potentially lead to
memory issues. In this work, we investigate the potential for amortizing the
complexity of the Mamba into each layer of the network.

Our approach centers around the concept of token rearrangement before feeding
them into the Forward Scan block. For a given input feature $\mathbf{z}_{i}$
from layer $i$, the output feature $\mathbf{z}_{i+1}$ of the Forward Scan
block after the rearrangement can be expressed as:

$\displaystyle{\mathbf{z}}_{\Omega_{i}}$
$\displaystyle=\texttt{arrange}({\mathbf{z}}_{i},\Omega_{i}),$ (1)
$\displaystyle\bar{{\mathbf{z}}}_{\Omega_{i}}$
$\displaystyle=\texttt{scan}({\mathbf{z}}_{\Omega}),$ (2)
$\displaystyle{\mathbf{z}}_{i+1}$
$\displaystyle=\texttt{arrange}(\bar{{\mathbf{z}}}_{\Omega_{i}},\bar{\Omega_{i}}),$
(3)

$\Omega_{i}$ represents the 1D permutation of layer $i$, which rearranges the
order of the patch tokens by $\Omega_{i}$, and $\Omega_{i}$ and
$\overline{\Omega}_{i}$ represent the reverse operation. This ensures that
both $\mathbf{z}_{i}$ and $\mathbf{z}_{i+1}$ maintain the sample order of the
original image tokens.

(a) sweep-scan

(b) zigzag-scan

(c) zigzag-scan with 8 schemes

Figure 3: The 2D Image Scan. Our mamba scan design is based on the sweep-scan
scheme shown in subfigure (a). From this, we developed a zigzag-scan scheme
displayed in subfigure (b) to enhance the continuity of the patches, thereby
maximizing the potential of the Mamba block. Since there are several possible
arrangements for these continuous scans, we have listed the eight most common
zigzag-scans in subfigure (c).

Now we explore the design of the $\Omega_{i}$ operation, considering
additional inductive biases from 2D images. We propose one key properties:
_Spatial Continuity_. Regarding Spatial Continuity, current innovations of
Mamba in images [96, 57, 54] often squeeze 2D patch tokens directly following
the computer hierarchy, such as row-and-column-major order. However, this
approach may not be optimal for incorporating the inductive bias with
neighboring tokens, as illustrated in Figure 3. To address this, we introduce
a novel scanning scheme designed to maintain spatial continuity during the
scan process. Additionally, we consider space-filling, which entails that for
a patch of size $N\times N$, the length of the 1D continuous scanning scheme
should be $N^{2}$. This helps to efficiently incorporate tokens to maximize
the potential of long sequence modeling within the Mamba block.

To achieve the aforementioned property, we heuristically design eight possible
space-filling continuous schemes111We also experimented with more complex
continuous space-filling paths, such as the Hilbert space-filling curve [62].
However, empirical findings indicate that this approach may lead to
deteriorated results. For further detailed comparisons, please refer to the
Appendix., denoted as ${\mathbf{S}}_{j}$ (where $j\in[0,7]$), as illustrated
in Figure 3. While there may be other conceivable schemes, for simplicity, we
limit our usage to these eight. Consequently, the scheme for each layer can be
represented as $\Omega_{i}={\mathbf{S}}_{\\{i\%8\\}}$, where $\%$ denotes the
modulo operator.

Figure 4: The Detail of our Zigzag Mamba block.  The detail of Mamba Scan is
shown in Figure 2. The condition can include a timestep and a text prompt.
These are fed into an MLP, which separately modulates the Mamba scan for long
sequence modeling and cross attention for multi-modal reasoning.

Deploying text-condition on Zigzag Mamba. While Mamba offers the advantage of
efficient long sequence modeling, it does so at the expense of the attention
mechanism. As a result, there has been limited exploration into incorporating
text-conditioning in Mamba-based diffusion models. To address this gap, we
propose a straightforward cross-attention block with skip layers built upon
the Mamba block, as illustrated in Figure 4. This design not only enables long
sequence modeling but also facilitates multi-token conditioning, such as text-
conditioning. Furthermore, it has the potential to provide interpretability
[15, 75, 38], as cross-attention has been utilized in diffusion models.

Generalizing to 3D videos by factorizing spatial and temporal information. In
previous sections, our focus has been on the spatial 2D Mamba, where we
designed several spatially continuous, space-filling 2D scanning schemes. In
this section, we aim to leverage this experience to aid in designing
corresponding mechanisms for 3D video processing. We commence our design
process by extrapolating from the conventional directional Mamba, as depicted
in Figure 5. Given a video feature input ${\mathbf{z}}\in\mathbb{R}^{B\times
T\times C\times W\times H}$, we propose three variants of the Video Mamba
Block for facilitating 3D video generation.

(a) Sweep-scan: In this approach, we directly flatten the 3D feature
$\mathbf{z}$ without considering spatial or temporal continuity. It’s worth
noting that the flattening process follows the computer hierarchy order,
meaning that no continuity is preserved in the flattened representation.

(b) 3D Zigzag: Compared with we formulate the 2D zigzag in previous
subsections, we follow the similar design to generalize it to 3D Zigzag to
keep the continuity in 2D and 3D simultaneously. Potentially, the scheme has
much more complexity. We heuristically list 8 schemes as well. However, we
empirically find that this scheme will lead to suboptimal optimization.

(c) Factorized 3D Zigzag = 2D Zigzag + 1D Sweep: To address the suboptimal
optimization issue, we propose to factorize the spatial and temporal
correlations as separate Mamba blocks. The order of their application can be
adjusted as desired, for example, "sstt" or "ststst", where "s" represents the
spatial-zigzag Mamba and "t" represents the temporal-zigzag Mamba. For a 1D
temporal sweep, we simply opt for forward and backward scanning, since there
is only one dimension on the time axis.

Computation Analysis. For a visual sequence
$\mathbf{T}\in\mathbb{R}^{1\times\text{M}\times\text{D}}$, the computation
complexity of global self-attention and $k$-direction mamba and our zigzag
mamba are as follows:

$\displaystyle\zeta(\text{self-
attention})=4\text{M}\text{D}^{2}+2\text{M}^{2}\text{D},$ (4)
$\displaystyle\zeta(\text{k-mamba})=k\times[3\text{M}(2\text{D})\text{N}+\text{M}(2\text{D})\text{N}^{2}],$
(5)
$\displaystyle\zeta(\text{zigzag})=3\text{M}(2\text{D})\text{N}+\text{M}(2\text{D})\text{N}^{2},$
(6)

where self-attention exhibits quadratic complexity with respect to sequence
length M, while Mamba exhibits linear complexity (N is a fixed parameter, set
to 16 by default). Here, $k$ represents the number of scan directions in a
single Mamba block. Therefore, $k$-mamba and zigzag share linear complexity
with respect to self-attention. Moreover, our zigzag method can eliminate the
$k$ series, further reducing the overall complexity.

Upon completing the design of the Zigzag Mamba network for improved visual
inductive-bias integration, we proceed to combine it with a new diffusion
framework, as illustrated below.

Figure 5: The 3D Video Scan. (a) We illustrate the bidirectional Mamba with
the sweep scan, where the spatial and temporal information is treated as a set
of tokens with a computer-hierarchy order. (b) For the 3D zigzag-scan, we aim
to maximize the potential of Mamba by employing a spatial continuous scan
scheme and adopting the optimal zigzag scan solution, as depicted in Figure 3.
(c) We further separate the reasoning between spatial and temporal
information, resulting in a factorized combination of 2D spatial scan
($\Omega$) plus a 1D temporal scan ($\Omega^{{}^{\prime}}$) scheme.

###  3.3 Diffusion Framework: Stochastic Interpolant

#### Sampling based on vector ${\mathbf{v}}$ and score ${\mathbf{s}}$.

Following [3, 76], the time-dependent probability distribution
$p_{t}({\mathbf{x}})$ of ${\mathbf{x}}_{t}$ also coincides with the
distribution of the reverse-time SDE [6]:

$d{\mathbf{X}}_{t}={\mathbf{v}}({\mathbf{X}}_{t},t)dt+\frac{1}{2}w_{t}{\mathbf{s}}({\mathbf{X}}_{t},t)dt+\sqrt{w_{t}}d\bar{\mathbf{W}}_{t},$
(7)

where $\bar{\mathbf{W}}_{t}$ is a reverse-time Wiener process, $w_{t}>0$ is an
arbitrary time-dependent diffusion coefficient,
${\mathbf{s}}({\mathbf{x}},t)=\nabla\log p_{t}({\mathbf{x}})$ is the score,
and ${\mathbf{v}}({\mathbf{x}},t)$ is given by the conditional expectation

$\displaystyle{\mathbf{v}}({\mathbf{x}},t)$
$\displaystyle=\mathbb{E}[\dot{\mathbf{x}}_{t}|{\mathbf{x}}_{t}={\mathbf{x}}],$
(8)
$\displaystyle=\dot{\alpha}_{t}\mathbb{E}[{\mathbf{x}}_{*}|{\mathbf{x}}_{t}={\mathbf{x}}]+\dot{\sigma}_{t}\mathbb{E}[{\boldsymbol{\varepsilon}}|{\mathbf{x}}_{t}={\mathbf{x}}],$

where $\alpha_{t}$ is a decreasing function of $t$, and $\sigma_{t}$ is an
increasing function of $t$. Here, $\dot{\alpha}_{t}$ and $\dot{\sigma}_{t}$
denote the time derivatives of $\alpha_{t}$ and $\sigma_{t}$, respectively.

As long as we can estimate the velocity ${\mathbf{v}}({\mathbf{x}},t)$ and/or
score ${\mathbf{s}}({\mathbf{x}},t)$ fields, we can utilize it for the
sampling process either by probability flow ODE [72] or the reverse-time SDE
(7). Solving the reverse SDE (7) backwards in time from
${\mathbf{X}}_{T}={\boldsymbol{\varepsilon}}\sim\mathcal{N}(0,\mathbf{I})$
enables generating samples from the approximated data distribution
$p_{0}(\mathbf{x})\sim p({\mathbf{x}})$. During sampling, we can perform
direct sampling from either ODE or SDEs to balance between sampling speed and
fidelity. If we choose to conduct ODE sampling, we can achieve this simply by
setting the noise term ${\mathbf{s}}$ to zero.

In [3], it shows that one of the two quantities
${\mathbf{s}}_{\theta}({\mathbf{x}},t)$ and
${\mathbf{v}}_{\theta}({\mathbf{x}},t)$ needs to be estimated in practice.
This follows directly from the constraint

$\displaystyle{\mathbf{x}}$
$\displaystyle=\mathbb{E}[{\mathbf{x}}_{t}|{\mathbf{x}}_{t}={\mathbf{x}}],$
(9)
$\displaystyle=\alpha_{t}\mathbb{E}[{\mathbf{x}}_{*}|{\mathbf{x}}_{t}={\mathbf{x}}]+\sigma_{t}\mathbb{E}[{\boldsymbol{\varepsilon}}|{\mathbf{x}}_{t}={\mathbf{x}}],$

which can be used to re-express the score ${\mathbf{s}}({\mathbf{x}},t)$ in
terms of the velocity ${\mathbf{v}}({\mathbf{x}},t)$ as

$\displaystyle{\mathbf{s}}({\mathbf{x}},t)$
$\displaystyle=\sigma_{t}^{-1}\frac{\alpha_{t}{\mathbf{v}}({\mathbf{x}},t)-\dot{\alpha}_{t}{\mathbf{x}}}{\dot{\alpha}_{t}\sigma_{t}-\alpha_{t}\dot{\sigma}_{t}}.$
(10)

Thus, ${\mathbf{v}}({\mathbf{x}},t)$ and ${\mathbf{s}}({\mathbf{x}},t)$ can be
mutually conversed. We illustrate how to compute them in the following.

Estimating the score ${\mathbf{s}}$ and the velocity ${\mathbf{v}}$. It has
been shown in score-based diffusion models [72] that the score can be
estimated parametrically as ${\mathbf{s}}_{\theta}({\mathbf{x}},t)$ using the
loss

$\mathcal{L}_{\mathrm{s}}(\theta)=\int_{0}^{T}\mathbb{E}[\|\sigma_{t}{\mathbf{s}}_{\theta}({\mathbf{x}}_{t},t)+{\boldsymbol{\varepsilon}}\|^{2}]{\mathrm{d}}t.$
(11)

Similarly, the velocity ${\mathbf{v}}({\mathbf{x}},t)$ can be estimated
parametrically as ${\mathbf{v}}_{\theta}({\mathbf{x}},t)$ via the loss

$\displaystyle\mathcal{L}_{\mathrm{v}}(\theta)$
$\displaystyle=\int_{0}^{T}\mathbb{E}[\|{\mathbf{v}}_{\theta}({\mathbf{x}}_{t},t)-\dot{\alpha}_{t}{\mathbf{x}}_{*}-\dot{\sigma}_{t}{\boldsymbol{\varepsilon}}\|^{2}]{\mathrm{d}}t,$
(12)

where $\theta$ represents the Zigzag Mamba network that we described in the
previous section, we adopt the linear path for training, due to its simplicity
and relatively straight trajectory:

$\displaystyle\alpha_{t}=1-t,\qquad$ $\displaystyle\sigma_{t}=t.$ (13)

We note that any time-dependent weight can be included under the integrals in
both (11) and (12). These weight factors play a crucial role in score-based
models when $T$ becomes large [45, 44]. Thus, they provide a general form that
considers both the time-dependent weight and the stochasticity.

##  4 Experiment

In this section, we begin by detailing the experimental setup concerning image
and video datasets, as well as our training details. Subsequently, we delve
into several in-depth analyses aimed at elucidating the rationale behind our
method design across various resolutions. Finally, we present our results
obtained from higher-resolution and more complex datasets.

###  4.1 Dataset and Training Detail

Image Dataset. To explore the scalability in high resolution, we conduct
experiments on the FacesHQ $1024\times 1024$. The general dataset that we use
for training and ablations is FacesHQ, a compilation of CelebA-HQ [88] and
FFHQ [43], as employed in previous work such as [24, 26].

For text-conditioned generation, we conduct the experiments on the MultiModal-
CelebA $256^{2}$,$512^{2}$ [88] and MS COCO $256\times 256$ [50] datasets.
Both datasets are composed of text-image pairs for training. Typically, there
are 5 to 10 captions per image in COCO and MultiModal-CelebA. We convert
discrete texts to a sequence of embeddings using a CLIP text encoder [66]
following Stable Diffusion [67]. Then these embeddings are fed into the
network as a sequence of tokens.

Video dataset. UCF101 dataset is consists of 13,320 video clips, which are
classified into 101 categories. The total length of these video clips is over
27 hours. All these videos are collected from YouTube and have a fixed frame
rate of 25 FPS with the resolution of $320\times 240$. We randomly sample
continuous 16 frames and resize the frames to $256\times 256$.

Training Details. We uniformly use AdamW [59] optimizer with $1e-4$ learning
rate. For extracting latent features, we employ off-the-shelf VAE encoders. To
mitigate computational costs, we adopted a mixed-precision training approach.
Additionally, we applied gradient clipping with a threshold of 2.0 and a
weight decay of 0.01 to prevent NaN occurrences during Mamba training. Most of
our experiments were conducted on 4 A100 GPUs, with scalability exploration
extended to 16 and 32 A100 GPUs. For sampling, we adopt the ODE sampling for
speed consideration. For further details, please refer to the Appendix 9.5.

Table 1: Ablation of Scanning Scheme Number. We evaluate various zigzag scanning schemes. Starting from a simple ‘‘Sweep’’ baseline, we consistently observe improvements as more schemes are implemented. MultiModal-CelebA$256$ |  MultiModal-CelebA$512$
---|---
|  FID${}^{\text{5k}}$ |  FDD${}^{\text{5k}}$ |  KID${}^{\text{5k}}$ |  FID${}^{\text{5k}}$ |  FDD${}^{\text{5k}}$ |  KID${}^{\text{5k}}$
Sweep | 158.1 | 75.9 | 0.169 | 162.3 | 103.2 | 0.203
Zigzag-1 | 65.7 | 47.8 | 0.051 | 121.0 | 78.0 | 0.113
Zigzag-2 | 54.7 | 45.5 | 0.041 | 96.0 | 59.5 | 0.079
Zigzag-8 | 45.5 | 26.4 | 0.011 | 34.9 | 29.5 | 0.023

###  4.2 Ablation Study

Scan Scheme Ablation. We provide several important findings based on our
ablation studies on MultiModal-CelebA dataset in various resolutions in Table
1. Firstly, switching the scanning scheme from sweep to zigzag led to some
gains. Secondly, as we increased the zigzag scheme from 1 to 8, we saw
consistent gains. This indicates that alternating the scanning scheme in
various blocks can be beneficial. Finally, the relative gain between Zigzag-1
and Zigzag-8 is more prominent at higher resolutions ($512\times 512$, or
longer sequence token number) compared to lower resolutions ($256\times 256$,
or shorter sequence token number), this shows the great potential and more
efficient inductive-bias incorporation in longer sequence number.

Spatial Continuity is Critical. We first explore the importance of spatial
continuity in Mamba design by grouping patches of size $N\times N$ into
various sizes: $2\times 2$, $4\times 4$, $8\times 8$, and $16\times 16$,
resulting in groups of patch sizes $N/2\times N/2$, $N/4\times N/4$,
$N/8\times N/8$, and $N/16\times N/16$, respectively. Then, we apply our
designed Zigzag-8 scheme at the group level instead of the patch level. Figure
6 illustrates that with increased spatial continuity, notably improved
performance is achieved. Furthermore, we compare our approach with random
shuffling of $N\times N$ patches, revealing notably inferior performance under
random shuffling conditions. All of these results collectively indicate that
spatial continuity is a critical requirement when applying Mamba in 2D
sequences.

Figure 6: Spatial Continuity Analysis. As we incrementally enlarge the patch
group size, the continuous segment of the patch also expands. This enhances
spatial continuity, which we find improves FID on MultiModal-CelebA 256, 512
dataset.

Ablation study about the Network and FPS/GPU-Memory. In Figure 7 (a,b),we
analyze the forward speed and GPU memory usage while varying the global patch
dimensions from $32\times 32$ to $196\times 196$. For the speed analysis, we
report Frame Per Second (FPS) instead of FLOPS, as FPS provides a more
explicit and appropriate evaluation of speed 222https://github.com/state-
spaces/mamba/issues/110#issuecomment-1916464012. For simplicity, we uniformly
apply the zigzag-1 Mamba scan scheme and use batch size=1 and patch size=1 on
an A100 GPU with 80GB memory. It’s worth noting that all methods share nearly
identical parameter numbers for fair comparison. We primarily compare our
method with two popular transformer-based Diffusion backbones, U-ViT [9] and
DiT [65]. It is evident that our method achieves the best FPS and GPU
utilization when gradually increasing the patching number. U-ViT demonstrates
the worst performance, even exceeds the memory bounds when the patch number is
196. Surprisingly, DiT’s GPU utilization is close to our method, which
supports our backbone choice of DiT from a practical perspective.

Furthermore, in Figure 7 (c) we conduct an ablation study on GPU memory and
FPS among different variants of our method. We find that our method incurs
nearly zero FPS and GPU memory burden when gradually increasing the Mamba Scan
Schemes. This analysis provides insights into the comparative performance and
efficiency of different diffusion backbones, highlighting the advantages of
our proposed method.

Order Receptive Field. We propose a new concept in Mamba-based structure for
multidimensional data. Given that various spatially-continuous zigzag paths
may exist in multidimensional data, we introduce the term _Order Receptive
Field_ which denotes the number of zigzag paths _explicitly_ employed in the
network design.

Ablation study about the Order Receptive Field and FPS/GPU-Memory. As depicted
in Fig. 10, Zigzag Mamba consistently maintains its GPU memory consumption and
FPS rate, even with a gradually increasing Order Receptive Field. In contrast,
our primary baseline, Parallel Mamba, along with variants like Bidirectional
Mamba and Vision Mamba [57, 96], experience a consistent decrease in FPS due
to increased parameters. Notably, Zigzag Mamba, with an Order Receptive Field
of 8, can perform faster without altering parameters.

(a) FPS v.s. Patch Number.

(b) GPU Memory v.s. Patch Number.

(c) GPU usage of variants of our method.

(d) FPS v.s. Patch Size.

Figure 7: (a, b).GPU Memory usage and FPS between our method and transformer-
based methods(U-VIT [9] and DiT [65]). (c). GPU memory usage and FPS among
various variants of our method. (d). The relationship between patch size and
FID under various resolutions.

(a) Order Receptive Field v.s. FPS.

(b) Order Receptive Field v.s. GPU Memory.

Figure 8: The ablation study about Order Receptive Field, FPS, GPU Memory.

Patch size. We conducted an ablation study on patch sizes ranging from 1, 2,
4, to 8 in Figure 7 (d), aiming to explore their behaviors under the framework
of Mamba. The results reveal that the FID deteriorates as the patch size
increases, aligning with the common understanding observed in the field of
transformers [23, 78]. This suggests that smaller patch sizes are crucial for
optimal performance.

Table 2: Main Results on MS-COCO dataset. Our method consistently outperforms the baseline and can achieve even better results when the training scale is increased. Variants |  FID${}^{\text{5k}}$
---|---
Sweep | 195.1
Zigzag-1 | 73.1
Bidirection Mamba [96] | 60.2
Zigzag-8 | 41.8
Zigzag-8 $\times$ 16GPU | 33.8
Table 3: Main result on FacesHQ-1024 dataset with 4,094 tokens in latent space. Our method can outperform the baseline and can achieve even better results when the training scale is increased.  Method |  FID${}^{\text{5k}}$ |  FDD${}^{\text{5k}}$
---|---|---
Bidirection Mamba-16GPU [96] | 51.1 | 66.3
Zigzag-Mamba -16GPU | 37.8 | 50.5
Zigzag-Mamba -32GPU | 26.6 | 31.2

###  4.3 Main Result

Main Result on 1024$\times$1024 FacesHQ. To elaborate on the scalability of
our method within the Mamba and Stochastic Interpolant framework, we provide
comparisons on a high-resolution dataset (1024$\times$1024 FacesHQ) in Table
3. Our primary comparison is against Bidirectional Mamba, a commonly used
solution for applying Mamba to 2D image data [57, 96]. With the aim of
investigating Mamba’s scalability in large resolutions up to 1,024, we employ
the diffusion model on the latent space of $128\times 128$ with a patch size
of 2, resulting in 4,096 tokens. The network is trained on 16 A100 GPUs.
Notably, our method demonstrates superior results compared to Bidirectional
Mamba. Details regarding loss and FID curves can be found in Appendix 9.3.
While constrained by GPU resource limitations, preventing longer training
duration, we anticipate consistent outperformance of Bidirectional Mamba with
extended training duration.

COCO dataset. To further compare the performance of our method, we also
evaluate it on the more complex and common dataset MS COCO. We compare with
the Bidirection Mamba as the baseline in Table 2. It should be noted that all
methods share nearly identical parameter numbers for fair comparison. We
trained all methods using 16 A100 GPUs. please check Appendix 9.5 for details.
As depicted in Table 2, our Zigzag-8 method outperforms Bidirectional Mamba as
well as Zigzag-1. This suggests that amortizing various scanning schemes can
yield significant improvements, attributed to better incorporation of the
inductive bias for 2D images in Mamba.

UCF101 dataset. In Table 4, we present our results on the UCF101 dataset,
training all methods using 4 A100 GPUs, with further scalability exploration
conducted using 16 A100 GPUs. We mainly compare our method consistantly with
Bidirection Mamba [96]. For the choice of the 3D Zigzag Mamba, please refer to
Appendix 9.5. For Factorized 3D Zigzag Mamba in video processing, we deploy
the _sst_ scheme for factorizing spatial and temporal modeling. This scheme
prioritizes spatial information complexity over temporal information,
hypothesizing that redundancy exists in the temporal domain. Our results
consistently demonstrate the superior performance of our method across various
scenarios, underscoring the intricacy and effectiveness of our approach.

Figure 9: Visualization of various resolutions on FacesHQ $1024\times 1024$ and MultiModal-CelebA $512\times 512$. Our generated samples present high fidelity across various resolutions.  Table 4: Video Scan Scheme on UCF101 dataset. Our method outperforms the baseline and can achieve even better results when the training scale is increased.  Method |  Frame-FID${}^{\text{5k}}$ |  FVD${}^{\text{5k}}$
---|---|---
Bidirection Mamba [96] -4GPU | 256.1 | 320.2
3D Zigzag Mamba -4GPU | 238.1 | 282.3
Factorized 3D Zigzag Mamba -4GPU | 216.1 | 210.2
Bidirection Mamba [96] -16GPU | 146.2 | 201.1
Factorized 3D Zigzag Mamba -16GPU | 121.2 | 140.1

Visualization. We demonstrate the image visualization of our best results on
FacesHQ 1024 and MultiModal-CelebA 512 in Figure 9. For the visualization of
videos, please refer to Appendix 9.1. It is evident that the visualization is
visually pleasing across various resolutions, indicating the efficacy of our
methods.

##  5 Conclusion

In this paper, we present the Zigzag Mamba Diffusion Model, developed within
the Stochastic Interpolant framework. Our initial focus is on addressing the
critical issue of spatial continuity. We then devise a Zigzag Mamba block to
better utilize the inductive bias in 2D images. Further, we factorize the 3D
Mamba into 2D and 1D Zigzag Mamba to facilitate optimization. We empirically
design various ablation studies to examine different factors. This approach
allows for a more in-depth exploration of the Stochastic Interpolant theory.
We hope our endeavor can inspire further exploration in the Mamba network
design.

##  6 Acknowledgement

We would like to thank Timy Phan, Yunlu Chen, Di Wu, Divin Yan for the
extensive proofreading.

## References

  * [1] Agarwal, N., Suo, D., Chen, X., Hazan, E.: Spectral state space models. arXiv (2023)
  * [2] Ahamed, M.A., Cheng, Q.: Mambatab: A simple yet effective approach for handling tabular data. arXiv (2024)
  * [3] Albergo, M.S., Boffi, N.M., Vanden-Eijnden, E.: Stochastic interpolants: A unifying framework for flows and diffusions. arXiv (2023)
  * [4] Albergo, M.S., Vanden-Eijnden, E.: Building normalizing flows with stochastic interpolants. arXiv (2022)
  * [5] Ali, A., Zimerman, I., Wolf, L.: The hidden attention of mamba models. arXiv (2024)
  * [6] Anderson, B.D.: Reverse-time diffusion equation models. Stochastic Processes and their Applications (1982)
  * [7] Anthony, Q., Tokpanov, Y., Glorioso, P., Millidge, B.: Blackmamba: Mixture of experts for state-space models. arXiv (2024)
  * [8] Ao, S., Zhao, W., Han, X., Yang, C., Liu, Z., Shi, C., Sun, M., Wang, S., Su, T.: Burstattention: An efficient distributed attention framework for extremely long sequences. arXiv (2024)
  * [9] Bao, F., Li, C., Cao, Y., Zhu, J.: All are worth words: a vit backbone for score-based diffusion models. CVPR (2023)
  * [10] Bao, F., Nie, S., Xue, K., Li, C., Pu, S., Wang, Y., Yue, G., Cao, Y., Su, H., Zhu, J.: One transformer fits all distributions in multi-modal diffusion at scale. arXiv (2023)
  * [11] Behrouz, A., Hashemi, F.: Graph mamba: Towards learning on graphs with state space models. arXiv (2024)
  * [12] Beltagy, I., Peters, M.E., Cohan, A.: Longformer: The long-document transformer. arXiv (2020)
  * [13] Ben-Hamu, H., Cohen, S., Bose, J., Amos, B., Grover, A., Nickel, M., Chen, R.T., Lipman, Y.: Matching normalizing flows and probability paths on manifolds. In: ICML (2022)
  * [14] Brandon, W., Nrusimha, A., Qian, K., Ankner, Z., Jin, T., Song, Z., Ragan-Kelley, J.: Striped attention: Faster ring attention for causal transformers. arXiv preprint arXiv:2311.09431 (2023)
  * [15] Chefer, H., Gur, S., Wolf, L.: Transformer interpretability beyond attention visualization. In: CVPR (2021)
  * [16] Chen, R.T., Rubanova, Y., Bettencourt, J., Duvenaud, D.K.: Neural ordinary differential equations. NeurIPS (2018)
  * [17] Chen, S., Xu, M., Ren, J., Cong, Y., He, S., Xie, Y., Sinha, A., Luo, P., Xiang, T., Perez-Rua, J.M.: Gentron: Delving deep into diffusion transformers for image and video generation. arXiv (2023)
  * [18] Child, R., Gray, S., Radford, A., Sutskever, I.: Generating long sequences with sparse transformers. arXiv (2019)
  * [19] Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., et al.: Rethinking attention with performers. arXiv (2020)
  * [20] Crowson, K., Baumann, S.A., Birch, A., Abraham, T.M., Kaplan, D.Z., Shippole, E.: Scalable high-resolution pixel-space image synthesis with hourglass diffusion transformers. arXiv (2024)
  * [21] Dao, T., Fu, D., Ermon, S., Rudra, A., Ré, C.: Flashattention: Fast and memory-efficient exact attention with io-awareness. NeurIPS (2022)
  * [22] Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A.P., Caron, M., Geirhos, R., Alabdulmohsin, I., et al.: Scaling vision transformers to 22 billion parameters. In: ICML (2023)
  * [23] Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al.: An image is worth 16x16 words: Transformers for image recognition at scale. In: ICLR (2021)
  * [24] Esser, P., Rombach, R., Ommer, B.: Taming transformers for high-resolution image synthesis. In: CVPR (2021)
  * [25] Fei, Z., Fan, M., Yu, C., Huang, J.: Scalable diffusion models with state space backbone. arXiv (2024)
  * [26] Fischer, J.S., Gui, M., Ma, P., Stracke, N., Baumann, S.A., Ommer, B.: Boosting latent diffusion with flow matching. arXiv (2023)
  * [27] Fu, D.Y., Dao, T., Saab, K.K., Thomas, A.W., Rudra, A., Ré, C.: Hungry hungry hippos: Towards language modeling with state space models. arXiv (2022)
  * [28] Gong, H., Kang, L., Wang, Y., Wan, X., Li, H.: nnmamba: 3d biomedical image segmentation, classification and landmark detection with state space model. arXiv (2024)
  * [29] Gong, J., Foo, L.G., Fan, Z., Ke, Q., Rahmani, H., Liu, J.: Diffpose: Toward more reliable 3d pose estimation. In: CVPR (2023)
  * [30] Gu, A., Dao, T.: Mamba: Linear-time sequence modeling with selective state spaces. arXiv (2023)
  * [31] Gu, A., Goel, K., Gupta, A., Ré, C.: On the parameterization and initialization of diagonal state space models. NeurIPS (2022)
  * [32] Gu, A., Goel, K., Ré, C.: Efficiently modeling long sequences with structured state spaces (2021)
  * [33] Gu, A., Johnson, I., Goel, K., Saab, K., Dao, T., Rudra, A., Ré, C.: Combining recurrent, convolutional, and continuous-time models with linear state space layers. NeurIPS (2021)
  * [34] Guo, H., Li, J., Dai, T., Ouyang, Z., Ren, X., Xia, S.T.: Mambair: A simple baseline for image restoration with state-space model. arXiv (2024)
  * [35] Gupta, A., Gu, A., Berant, J.: Diagonal state spaces are as effective as structured state spaces. NeurIPS (2022)
  * [36] He, W., Han, K., Tang, Y., Wang, C., Yang, Y., Guo, T., Wang, Y.: Densemamba: State space models with dense hidden connection for efficient large language models. arXiv (2024)
  * [37] He, X., Cao, K., Yan, K., Li, R., Xie, C., Zhang, J., Zhou, M.: Pan-mamba: Effective pan-sharpening with state space model. arXiv (2024)
  * [38] Hertz, A., Mokady, R., Tenenbaum, J., Aberman, K., Pritch, Y., Cohen-Or, D.: Prompt-to-prompt image editing with cross attention control. arXiv (2022)
  * [39] Ho, J., Jain, A., Abbeel, P.: Denoising diffusion probabilistic models. In: NeurIPS (2020)
  * [40] Ho, J., Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., Fleet, D.J.: Video diffusion models. In: ARXIV (2022)
  * [41] Huang, Z., Zhou, P., Yan, S., Lin, L.: Scalelong: Towards more stable training of diffusion model via scaling network long skip connection. NeurIPS (2024)
  * [42] Karras, T., Aittala, M., Aila, T., Laine, S.: Elucidating the design space of diffusion-based generative models. In: NeurIPS (2022)
  * [43] Karras, T., Laine, S., Aila, T.: A style-based generator architecture for generative adversarial networks. In: CVPR (2019)
  * [44] Kingma, D., Salimans, T., Poole, B., Ho, J.: Variational diffusion models. In: NeurIPS (2021)
  * [45] Kingma, D.P., Gao, R.: Understanding the diffusion objective as a weighted integral of elbos. arXiv (2023)
  * [46] Kitaev, N., Kaiser, Ł., Levskaya, A.: Reformer: The efficient transformer. arXiv (2020)
  * [47] Lee, S., Kim, B., Ye, J.C.: Minimizing trajectory curvature of ode-based generative models. ICML (2023)
  * [48] Li, S., Singh, H., Grover, A.: Mamba-nd: Selective state space modeling for multi-dimensional data. arXiv (2024)
  * [49] Liang, D., Zhou, X., Wang, X., Zhu, X., Xu, W., Zou, Z., Ye, X., Bai, X.: Pointmamba: A simple state space model for point cloud analysis. arXiv preprint arXiv:2402.10739 (2024)
  * [50] Lin, T.Y., Maire, M., Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., Zitnick, C.L.: Microsoft coco: Common objects in context. In: ECCV (2014)
  * [51] Lipman, Y., Chen, R.T., Ben-Hamu, H., Nickel, M., Le, M.: Flow matching for generative modeling. ICLR (2023)
  * [52] Liu, G.H., Chen, T., So, O., Theodorou, E.: Deep generalized schrödinger bridge. NeurIPS (2022)
  * [53] Liu, H., Zaharia, M., Abbeel, P.: Ring attention with blockwise transformers for near-infinite context. arXiv (2023)
  * [54] Liu, J., Yang, H., Zhou, H.Y., Xi, Y., Yu, L., Yu, Y., Liang, Y., Shi, G., Zhang, S., Zheng, H., et al.: Swin-umamba: Mamba-based unet with imagenet-based pretraining. arXiv (2024)
  * [55] Liu, X., Gong, C., Liu, Q.: Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv (2022)
  * [56] Liu, X., Gong, C., Liu, Q.: Flow straight and fast: Learning to generate and transfer data with rectified flow. ICLR (2023)
  * [57] Liu, Y., Tian, Y., Zhao, Y., Yu, H., Xie, L., Wang, Y., Ye, Q., Liu, Y.: Vmamba: Visual state space model. arXiv (2024)
  * [58] Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., Guo, B.: Swin transformer: Hierarchical vision transformer using shifted windows. In: ICCV (2021)
  * [59] Loshchilov, I., Hutter, F.: Decoupled weight decay regularization. In: ICLR (2019)
  * [60] Ma, J., Li, F., Wang, B.: U-mamba: Enhancing long-range dependency for biomedical image segmentation. arXiv (2024)
  * [61] Ma, N., Goldstein, M., Albergo, M.S., Boffi, N.M., Vanden-Eijnden, E., Xie, S.: Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. arXiv (2024)
  * [62] McKenna, D.M.: Hilbert curves: Outside-in and inside-gone. Mathemaesthetics, Inc (2019)
  * [63] Nguyen, E., Goel, K., Gu, A., Downs, G., Shah, P., Dao, T., Baccus, S., Ré, C.: S4nd: Modeling images and videos as multidimensional signals with state spaces. NeurIPS (2022)
  * [64] OpenAI: Sora: Creating video from text (2024), https://openai.com/sora
  * [65] Peebles, W., Xie, S.: Scalable diffusion models with transformers. arXiv (2022)
  * [66] Radford, A., Kim, J.W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al.: Learning transferable visual models from natural language supervision. In: ICML (2021)
  * [67] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., Ommer, B.: High-resolution image synthesis with latent diffusion models. In: CVPR (2022)
  * [68] Ruan, J., Xiang, S.: Vm-unet: Vision mamba unet for medical image segmentation. arXiv (2024)
  * [69] Smith, J.T., Warrington, A., Linderman, S.W.: Simplified state space layers for sequence modeling. arXiv (2022)
  * [70] Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., Ganguli, S.: Deep unsupervised learning using nonequilibrium thermodynamics. In: ICML (2015)
  * [71] Song, Y., Ermon, S.: Generative modeling by estimating gradients of the data distribution. arXiv (2019)
  * [72] Song, Y., Sohl-Dickstein, J., Kingma, D.P., Kumar, A., Ermon, S., Poole, B.: Score-based generative modeling through stochastic differential equations. In: ICLR (2021)
  * [73] Stein, G., Cresswell, J., Hosseinzadeh, R., Sui, Y., Ross, B., Villecroze, V., Liu, Z., Caterini, A.L., Taylor, E., Loaiza-Ganem, G.: Exposing flaws of generative model evaluation metrics and their unfair treatment of diffusion models. NeurIPS (2023)
  * [74] Sun, Z., Yang, Y., Yoo, S.: Sparse attention with learning to hash. In: ICLR (2021)
  * [75] Tang, R., Liu, L., Pandey, A., Jiang, Z., Yang, G., Kumar, K., Stenetorp, P., Lin, J., Ture, F.: What the daam: Interpreting stable diffusion using cross attention. arXiv (2022)
  * [76] Tong, A., Malkin, N., Fatras, K., Atanackovic, L., Zhang, Y., Huguet, G., Wolf, G., Bengio, Y.: Simulation-free schr$\backslash$" odinger bridges via score and flow matching. arXiv (2023)
  * [77] Unterthiner, T., van Steenkiste, S., Kurach, K., Marinier, R., Michalski, M., Gelly, S.: Fvd: A new metric for video generation. ICLR Workshop (2019)
  * [78] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, Ł., Polosukhin, I.: Attention is all you need. In: NeurIPS (2017)
  * [79] Wang, C., Tsepa, O., Ma, J., Wang, B.: Graph-mamba: Towards long-range graph sequence modeling with selective state spaces. arXiv (2024)
  * [80] Wang, J., Gangavarapu, T., Yan, J.N., Rush, A.M.: Mambabyte: Token-free selective state space model. arXiv (2024)
  * [81] Wang, J., Yan, J.N., Gu, A., Rush, A.M.: Pretraining without attention. arXiv (2022)
  * [82] Wang, S., Li, Q.: Stablessm: Alleviating the curse of memory in state-space models through stable reparameterization. arXiv (2023)
  * [83] Wang, S., Xue, B.: State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. NeurIPS (2024)
  * [84] Wang, W., Ma, S., Xu, H., Usuyama, N., Ding, J., Poon, H., Wei, F.: When an image is worth 1,024 x 1,024 words: A case study in computational pathology. arXiv (2023)
  * [85] Wang, Z., Ma, C.: Semi-mamba-unet: Pixel-level contrastive cross-supervised visual mamba-based unet for semi-supervised medical image segmentation. arXiv (2024)
  * [86] Wang, Z., Zheng, J.Q., Zhang, Y., Cui, G., Li, L.: Mamba-unet: Unet-like pure visual mamba for medical image segmentation. arXiv (2024)
  * [87] Wu, L., Wang, D., Gong, C., Liu, X., Xiong, Y., Ranjan, R., Krishnamoorthi, R., Chandra, V., Liu, Q.: Fast point cloud generation with straight flows. In: CVPR (2023)
  * [88] Xia, W., Yang, Y., Xue, J.H., Wu, B.: Tedigan: Text-guided diverse face image generation and manipulation. In: CVPR (2021)
  * [89] Xing, Z., Ye, T., Yang, Y., Liu, G., Zhu, L.: Segmamba: Long-range sequential modeling mamba for 3d medical image segmentation. arXiv (2024)
  * [90] Yan, J.N., Gu, J., Rush, A.M.: Diffusion models without attention. arXiv (2023)
  * [91] Yang, Y., Xing, Z., Zhu, L.: Vivim: a video vision mamba for medical video object segmentation. arXiv (2024)
  * [92] Yu, A., Nigmetov, A., Morozov, D., Mahoney, M.W., Erichson, N.B.: Robustifying state-space models for long sequences via approximate diagonalization. arXiv (2023)
  * [93] Zhang, T., Li, X., Yuan, H., Ji, S., Yan, S.: Point could mamba: Point cloud learning via state space model. arXiv (2024)
  * [94] Zhang, Z., Liu, A., Reid, I., Hartley, R., Zhuang, B., Tang, H.: Motion mamba: Efficient and long sequence motion generation with hierarchical and bidirectional selective ssm. arXiv (2024)
  * [95] Zheng, Z., Wu, C.: U-shaped vision mamba for single image dehazing. arXiv (2024)
  * [96] Zhu, L., Liao, B., Zhang, Q., Wang, X., Liu, W., Wang, X.: Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv (2024)
  * [97] zhuzilin: Ring flash attention. https://github.com/zhuzilin/ring-flash-attention (2024)

##  7 Limitations and Future Work

Our method relies solely on the Mamba Block with a DiT-style layout and
conditioning manner. However, a potential limitation of our work is that we
cannot exhaustively list all possible spatial continuous zigzag scanning
schemes given a specific global patch size. Currently, we set these scanning
schemes empirically, which may lead to sub-optimal performance. Additionally,
due to GPU resource constraints, we were unable to explore longer training
durations, although we anticipate similar conclusions.

For future work, we aim to delve into various applications of the Zigzag
Mamba, leveraging its scalability for long-sequence modeling. This exploration
may lead to improved utilization of the Mamba framework across different
domains and applications.

##  8 Impact Statement

This work aims to enhance the scalability and unlock the potential of the
Mamba algorithm within the framework of diffusion models, enabling the
generation of large images with high-fidelity. By incorporating our cross-
attention mechanism into the Mamba block, our method can also facilitate text-
to-image generation. However, like other endeavors aimed at enhancing the
capabilities and control of large-scale image synthesis models, our approach
carries the risk of enabling the generation of harmful or deceptive content.
Therefore, ethical considerations and safeguards must be implemented to
mitigate these risks.

##  9 Appendix

Table 5: The ablation about Position Embedding(PE) on unconditional MultiModal-CelebA256. Learnable Positional Encoding (PE) outperforms Sinusoidal PE, and Sinusoidal PE is superior to No PE.  PE | Scan |  FID${}^{\text{5k}}$
---|---|---
\-- | sweep2 | 21.33
\-- | zigzag8 | 14.27
Sinusoidal | sweep2 | 18.47
Sinusoidal | zigzag8 | 14.03
Learnable | sweep2 | 16.38
Learnable | zigzag8 | 13.32
Table 6: The ablation about Hilbert and Zigzag scan path under various Order Receptive Field ( ORF ) on unconditional MultiModal-CelebA256.. Scan-ORF |  FID${}^{\text{5k}}$
---|---
hilbert-8 | 27.38
hilbert-2 | 61.67
zigzag-2 | 15.45
zigzag-8 | 13.32
Table 7: Different Order Receptive Field( ORF ) of our Zigma on unconditional MultiModal-CelebA256. ORF of ZigZag |  FID${}^{\text{5k}}$
---|---
2 | 15.45
4 | 13.46
6 | 13.42
8 | 13.32
Table 8: Various methods for text-to-image generation on the MultiModal-CelebA 256 dataset. Method |  FID${}^{\text{5k}}$ |  FDD${}^{\text{5k}}$ |  KID${}^{\text{5k}}$
---|---|---|---
In-Context | 61.1 | 39.1 | 0.061
Cross-Attention | 45.5 | 26.4 | 0.011
Table 9: Details of ZigMa Model Variants. We follow previous works [65, 9, 23]
model configurations for the Small (S), Base (B) and Large (L) variants; we
also introduce an XLarge (XL) config as our largest model. CA denotes the
cross-attention for text-to-image conditioning.

Model | Layers $N$ | Hidden size $d$ | #params
---|---|---|---
ZigMa-S | 12 | 384 | 31.3M
ZigMa-B | 12 | 768 | 133.8M
ZigMa-L | 24 | 1024 | 472.5M
ZigMa-XL | 28 | 1152 | 1058.7M
CA-ZigMa-S | 12 | 384 | 59.2M
CA-ZigMa-B | 12 | 768 | 214.1M
CA-ZigMa-L | 24 | 1024 | 724.4M
CA-ZigMa-XL | 28 | 1152 | 1549.8M

###  9.1 Visualization

FacesHQ $1024\times 1024$ uncurated visualization in Fig. 16.

MS-COCO uncurated visualization. We visualize the samples in Fig. 15.

(a) Model Complexity v.s. FPS.

(b) Model Complexity v.s. GPU Memory.

Figure 10: The ablation study about Model Complexity, FPS, GPU Memory.  Figure
11: The FID trends comparing the Hilbert scan, Sweep scan, and our Zigzag
scan. The y-axis is logarithmic scale.

(a) size=8

(b) size=16

Figure 12: The Hilbert space-filling curve with various sizes.  Figure 13: The
demonstration of Peano Curve. The figure is borrowed from
https://en.wikipedia.org/wiki/Peano_curve.

###  9.2 New Result about the Scanning Scheme and Position Embedding.

We also conduct basic ablations on various factors, including position
embedding and various Hilbert space-filling curves. Unlike the experiments in
the main paper, we perform these experiments on unconditional MultiModal-
CelebA256 dataset for a uniform comparison. We train the network for 100,000
steps.

The ablation about the Position Embedding. As shown in Table 5, the learnable
embedding performs better than the Sinusoidal embedding, which in turn
performs better than no position embedding. In various cases, our zigzag
method surpasses the baselines. Notably, our performance remains almost
unchanged whether we use the Sinusoidal position embedding or no position
embedding. This suggests that our method can better incorporate spatial
inductive-bias compared to our baseline. Finally, using the learnable position
embedding provides further, albeit marginal, gains suggesting that better
position embedding exists even within our zigzag scan scheme.

Exploration of Hilbert space-filling curve. Primarily, we ablate the Hilbert
scan curve [62], as depicted in Figure 12. There are also eight variants of
this scan considering different angles and starting points. We rearrange them
in a similar manner to our Zigzag scan. All parameters are kept consistent for
a fair comparison. We utilize the Gilbert algorithm
333https://github.com/jakubcerveny/gilbert to guarantee that the Hilbert curve
remains continuous across any square size. We train our network on single
A100-SXM4-80GB for 120k iterations. We evaluate the FID on 5,000 images for a
fixed step, the FID curve is demonstrated in Figure 11.

While the Hilbert space-filling curve offers increased locality compared to
our zigzag scan and maintains continuity, its complex structure appears to
hinder the SSM’s ability to work on the flattened sequence, resulting in a
worse inductive bias than our zigzag curve on natural images. Therefore, we
hypothesize that structure may hold greater significance than locality in
generative tasks.

#### Hilbert Curve is difficult to optimize.

We show the result in Table 6. We can observe that the performance of the
Hilbert scan path drops significantly, even if we decrease the Order of
Receptive Field (ORF). This confirms the assumption that the Hilbert scan path
is difficult to optimize, even when considering only two different schemes of
the Hilbert scan.

The influence of various Order Receptive Field ( ORF ). We show the influence
of the Order Receptive Field ( ORF ) in Table 7. We observe that increasing
the ORF to 4 improves performance. The performance appears to plateau around
ORF=4. Since there is no additional parameter or GPU memory burden, we
recommend practitioners to choose the largest ORF( 8 in our case).

#### Another Interpretation: Zigzag scan is the simplest Peano curve.

Our Zigzag scan can be seen as the simplest case of Peano Curve as shown in
Figure 13.

###  9.3 New Result of 2D visual data

The variants of our ZigMa Models. We list the variants of our model in Tab. 9.
We use the Base (B) Model as the default. Applying the cross-attention model
is optional, as this module can introduce some parameter and speed burdens.
However, any advancements in attention optimization can be seamlessly
integrated into our model.

Ablation study about the Model Complexity and FPS/GPU-Memory. As shown in Fig.
10. Our method can achieve much better parameter efficiency when incorporating
the receptive order. The receptive order refers to the cumulative spatial-
continuous zigzag scan path in 2D images, which we’ve incorporated into the
Mamba as an inductive bias. We list the parameter consumption when we
gradually increase the receptive order in Fig. 10. The receptive order refers
to the cumulative spatial-continuous zigzag scan path in 2D images, which
we’ve incorporated into the Mamba as an inductive bias.

Loss and FID curve. The training loss curve and the FID curve are demonstrated
in Fig. 14. The loss and FID show the same trend, with our Zigzag Mamba
consistently outperforming other baselines like Sweep-1 and Sweep-2.

(a) Loss trend of the MultiModal-CelebA256.

(b) FID trend of the MultiModal-CelebA256.

(c) Loss trend of the MultiModal-CelebA512.

(d) FID trend of the MultiModal-CelebA512.

Figure 14: The loss and FID trend under various resolutions on dataset
MultiModal-CelebA. Sweep-1 and Sweep-2 are the Mamba scans without spatial
continuity, while Zigzag-8 represents our method. This is the direct log from
weight-and-bias (wandb).

In-context v.s. Cross Attention We compare our cross-attention with in-context
attention in Tab. 8. For in-context attention, we concatenate the text tokens
with the image tokens and feed them into the Mamba block. Our results
demonstrate that in-context attention performs worse than our cross-attention.
We hypothesize that this is due to the discontinuity between the text tokens
and the image patch tokens.

###  9.4 New result of 3D Visual Data

The choice of the 3D Zigzag Mamba. For Factorized 3D Zigzag Mamba in video
processing, we deploy the _sst_ scheme for factorizing spatial and temporal
modeling. This scheme prioritizes spatial information (ss)complexity over
temporal information (t), hypothesizing that redundancy exists in the temporal
domain. There are numerous other possible combinations of s and t to explore,
which we leave for future work.

###  9.5 More Details

More related works Several works [83, 82] have demonstrated that the State-
Space Model possesses universal approximation ability under certain
conditions. Mamba, as a new State-Space Model, has superior potential for
modeling long sequences efficiently, which has been explored in various fields
such as medical imaging [89, 60, 86, 68, 28], image restoration [34, 95],
graphs [11, 79], NLP word byte [80], tabular data [2], human motion synthesis
[94], point clouds [49, 93], image generation [25], semi-supervised learning
[85],interpretability [5], image dehazing [95] and pan sharpening [37]. It has
been extended to Mixture of Experts [7], spectral space [1], multi-dimension
[48, 63, 96, 57] and dense connection [36]. Among them, the most related to us
are VisionMamba[96, 57], S4ND [63] and Mamba-ND [48]. VisionMamba[96, 57] uses
a bidirectional SSM in discriminative tasks which incurs a high computational
cost. Our method applies a simple alternative mamba diffusion in generative
models. S4ND [63] introduces local convolution into Mamba’s reasoning process,
moving beyond the use of only 1D data. Mamba-ND [48] takes multi-
dimensionality into account in discriminative tasks, making use of various
scans within a single block. In contrast, our focus is on distributing scan
complexity across every layer of the network, thus maximizing the
incorporation of inductive bias from visual data with zero parameter burden.

Double-Indexing Issue for $\Omega_{i}$. As shown in Fig. 2. We need to arrange
and rearrange operation that needs to conduct indexing along the token number
dimension to achieve spatial-continuous mamba reasoning, as the indexing can
be time-consuming when considering the large token numbers, We can formulate
the arrange and rearrange operation as follows:

$\displaystyle{\Omega^{{}^{\prime}}_{i}}$
$\displaystyle=\bar{\Omega}_{i-1}\cdot\Omega_{i},$ (14)
$\displaystyle{\mathbf{z}}_{i+1}$
$\displaystyle=\texttt{scan}({\mathbf{z}}_{\Omega^{{}^{\prime}}_{i}}),$ (15)

where $\bar{\Omega}_{-1}=I$, this assumes that the Mamba-based networks are
permutation equivariant to the order of the tokens. They require 50% fewer
indexing operations, a point which we reiterate here for clearer comparison:

$\displaystyle{\mathbf{z}}_{\Omega_{i}}$
$\displaystyle=\texttt{arrange}({\mathbf{z}}_{i},\Omega_{i}),$ (17)
$\displaystyle\bar{{\mathbf{z}}}_{\Omega_{i}}$
$\displaystyle=\texttt{scan}({\mathbf{z}}_{\Omega}),$ (18)
$\displaystyle{\mathbf{z}}_{i+1}$
$\displaystyle=\texttt{arrange}(\bar{{\mathbf{z}}}_{\Omega_{i}},\bar{\Omega_{i}}),$
(19)

Evaluation Metrics. For image-level fidelity, we use established metrics such
as Fréchet Inception Distance (FID) and Kernel Inception Distance (KID),
following previous works. However, since studies [73, 20] have shown that FID
does not fully reflect human-based opinions, we also adopt the Fréchet DINOv2
Distance (FDD) using the official repository 444https://github.com/layer6ai-
labs/dgm-eval. Our method primarily involves sampling 5,000 real and 5,000
fake images to compute the related metrics.

We primarily consider two metrics for video fidelity evaluation: framewise FID
and Fréchet Video Distance (FVD) [77]. We sample 200 videos and compute the
respective metrics based on these samples.

The training parameters of various datasets are listed in Tab. 10. We don’t
apply any position encoding because Mamba, unlike Transformer, is not
permutation invariant. Therefore, its position is automatically encoded by its
order in Mamba. For the COCO dataset, a weight decay of 0.01 can contribute to
marginal FID gains (approximately 0.8).

The conditioning of timestep and prompt. The conditioning process is
illustrated in Algorithm 1.

Table 10: Hyperparameters and number of parameters for our network in various
datasets. All models are trained on a single A100 with 40GB of VRAM using a
bfloat16 of accelerator package.

| FacesHQ 1024 | MS-COCO 256 | MultiModal-CelebA 512 | UCF-101
---|---|---|---|---
Autoencoder $f$ | 8 | 8 | 8 | 8
$z$-shape | $4\times 128\times 128$ | $4\times 32\times 32$ | $4\times 64\times 64$ | $16\times 4\times 32\times 32$
Model size | 133.8M | 133.8M | 133.8M | 133.8M
Patch size | 2 | 1 | 1 | 2
Channels | 768 | 768 | 768 | 768
Depth | 12 | 12 | 12 | 12
Optimizer | AdamW | AdamW | AdamW | AdamW
Batch size/GPU | 8 | 8 | 4 | 8
GPU num | 32 | 32 | 16 | 16
Learning rate | 1e-4 | 1e-4 | 1e-4 | 1e-4
weight decay | 0 | 0 | 0 | 0
EMA rate | 0.9999 | 0.9999 | 0.9999 | 0.9999
Warmup steps | 0 | 0 | 0 | 0
A100-hours | 768 | 768 | 384 | 384

Algorithm 1 Mamba Block.

⬇

def mamba_block(x, t, c = None):

# x: input data, shape [B, (W x H) , C] or [B, (T x W x H) , C]

# t: timestep, (B, C)

# c: condition, (B, D, C)

x = reshape(x) # (B, K, C)

def _mamba(x):

x = rearrange(x) # rearrange by a zigzag manner

x = mamba(x)

x = rearrange_back(x)# rearrange back by a zigzag manner

m, n = AdaLN( t )

x = _mamba( x * m + n ) + x

if c is not None:

p, q = AdaLN( c )

x = cross_attention( x * p + q ) + x

return x

Figure 15: The Uncurated Visualization of MS-COCO dataset. The first row is
illustrated with pairs of images and their captions, while the remaining rows
only images.  Figure 16: Uncurated Visualization of FacesHQ dataset.