MaskMamba: A Hybrid Mamba-Transformer Model for Masked Image Generation

## Abstract

Abstract Image generation models have encountered challenges related to scalability and quadratic complexity, primarily due to the reliance on Transformer-based backbones. In this study, we introduce MaskMamba, a novel hybrid model that combines Mamba and Transformer architectures, utilizing Masked Image Modeling for non-autoregressive image synthesis. We meticulously redesign the bidirectional Mamba architecture by implementing two key modifications: (1) replacing causal convolutions with standard convolutions to better capture global context, and (2) utilizing concatenation instead of multiplication, which significantly boosts performance while accelerating inference speed. Additionally, we explore various hybrid schemes of MaskMamba, including both serial and grouped parallel arrangements. Furthermore, we incorporate an in-context condition that allows our model to perform both class-to-image and text-to-image generation tasks. Our MaskMamba outperforms Mamba-based and Transformer-based models in generation quality. Notably, it achieves a remarkable 54.44 % percent 54.44 54.44\% 54.44 % improvement in inference speed at a resolution of 2048 × 2048 2048 2048 2048\times 2048 2048 × 2048 over Transformer.

## 1 Introduction

In recent years, the field of generative image models in computer vision has witnessed significant advancements, particularly in class-to-image (Gao et al. (); Sun et al. (); Sauer et al. ()) and text-to-image tasks (Yu et al. (; ); Bao et al. ()). Traditional autoregressive generative models, such as VQGAN (Esser et al. ()) and LlamaGen (Sun et al. ()), demonstrate excellent performance in class-conditional generation. In the realm of text-conditional generation, models like Parti (Yu et al. (; )) and DALL-E (Ramesh et al. ()) convert images into discrete tokens using the image tokenizer and project the encoded text features to caption embeddings via an additional MLP (Chen et al. ()), operating in an autoregressive manner for both training and inference. Concurrently, non-autoregressive methods, including MAGE (Li et al. ()) and MUSE (Chang et al. ()), leverage Masked Image Modeling, transforming images into discrete tokens during training and predicting randomly masked tokens.

Another prominent approach to image generation involves diffusion models (Song & Ermon (); Song et al. (); Ho et al. (); Dhariwal & Nichol (); Nichol et al. ()), such as LDM (Rombach et al. ()) with an UNet backbone. Although these models demonstrate high generation quality, their convolutional neural network architecture imposes constraints that hinder scalability. To address this challenge, Transformer-based generative models, such as DiT (Peebles & Xie ()), enhance global modeling capabilities through attention mechanisms and significantly improve generation quality. However, the computational complexity of attention mechanisms increases quadratically with sequence length, which constrains both training and inference efficiency.

Mamba (Gu & Dao ()) presents a state-space model (Gu et al. (; )) characterized by linear time complexity, offering substantial advantages in managing long sequence tasks. Contemporary image generation efforts, including DiM (Teng et al. ()), ZigMa (Hu et al. ()), and diffuSSM (Yan et al. ()), primarily replace the original Transformer block with a Mamba module. These models enhance both efficiency and scalability. Nevertheless, generating images based on diffusion models typically requires hundreds of iterations, which can be prohibitively time-consuming.

Figure: Figure 1: Examples of class-conditional (top) and text-conditional (bottom) image generation using MaskMamba-XL.
Refer to caption: x1.png

To eliminate the quadratic complexity growth with sequence length in Transformer models and the excessive generation iterations in autoregressive models, we introduce MaskMamba that integrates Mamba and Transformer architectures and utilizes non-autoregressive Masked Image Modeling (Ni et al. (); Lezama et al. ()) for image synthesis. We meticulously redesign Bi-Mamba (Mo & Tian (); Zhu et al. ()) to render it suitable for masked image generation by replacing the causal convolution with standard convolution. Meanwhile, we select concatenation instead of multiplication in the final stage of Bi-Mamba to reduce computational complexity, notably improving the inference speed by $17.77\%$ compared to Bi-Mamba (Zhu et al. ()).

We further investigate various MaskMamba hybrid schemes, including serial and grouped parallel schemes (Shaker et al. ()). In serial schemes, we explore alternating layer-by-layer arrangements, as well as placing the Transformer in the last $N/2$ layers. For grouped parallel schemes, we assess the effects of partitioning the model into two or four groups along the channel dimension. Our findings indicate that placing the Transformer in the final layers significantly enhances the model’s ability to capture global context. Additionally, we implement an in-context condition that allows our model to perform both class-to-image and text-to-image generation tasks within a single framework as show in Fig.1. Meanwhile, we investigate the placement of condition embeddings (Zhu et al. ()) by inserting them at differnent positions of the input sequence including head, middle, and tail. The results indicate that placing condition embedding at the middle yields optimal performance.

In the experimental section, we substantiate the generative capabilities of MaskMamba through two distinct tasks: class-conditional generation and text-conditional generation, utilizing various model sizes for each task. For class-to-image generation task, we execute training over 300 epochs on the ImageNet1k (Deng et al. ()) dataset, benchmarking our MaskMamba against Transformer-based and Mamba-based models of analogous size. The results demonstrate that our MaskMamba outperforms both counterparts with respect to generation quality and inference speed. Furthermore, we train and evaluate on CC3M (Sharma et al. ()) dataset, attaining superior performance on CC3M and MS-COCO (Lin et al. ()) valid datasets.

In summary, our contributions include:

- 1.
We redesign Bi-Mamba to improve its suitability for masked image generation tasks by replacing causal convolution with standard convolution. Additionally, we substitute multiplication with concatenation at the final stage, resulting in a significant performance boost and a $17.77\%$ increase in inference speed compared to Bi-Mamba.
- 2.
We introduce MaskMamba, a unified generative model that integrates redesign Bi-Mamba and Transformer layers, enabling class-to-image and text-to-image generation tasks to be performed in the same model through an in-context condition.
- 3.
Our MaskMamba surpasses both Transformer-based and Mamba-based models in terms of generation quality and inference speed on the ImageNet1k and CC3M datasets.

## 2 Related Work

Image Generation. The domain of image generation is witnessing significant advancements in current research. Initial autoregressive image generative models (Yu et al. (); Ding et al. ()), such as VQGAN (Esser et al. ()) and LlamaGen (Sun et al. ()), have illustrated the potential to generate high-fidelity images by transforming images into discrete tokens and applying autoregressive models to generate image tokens. The advent of text-to-image generative models, like Parti (Yu et al. (; )) and DALL-E (Ramesh et al. ()), have further propelled progress within this area. Nonetheless, these models exhibit specific inefficiencies in their generation process. To address these issues, non-autoregressive generative models such as MaskGIT (Chang et al. ()), MAGE (Li et al. ()), and MUSE (Chang et al. ()) enhance generation efficiency through Masked Image Modeling. Simultaneously, diffusion models (Song & Ermon (); Song et al. (); Ho et al. (); Dhariwal & Nichol (); Nichol et al. (); Saharia et al. ()), represented by LDM (Rombach et al. ()), excel in generation quality despite experiencing scalability constraints linked to their convolutional neural network-based architecture. To overcome these limitations, Transformer-based generative models, including DiT (Peebles & Xie ()), advance global modeling capabilities by incorporating attention mechanisms. However, these models continue to struggle with the quadratic increase in computational complexity when processing extensive sequences.

Mamba Vision. The Transformer (Vaswani ()), established as a leading network architecture, is extensively utilized across various tasks. Nonetheless, its quadratic computational complexity presents significant obstacles for the efficient handling of long sequence tasks. In recent developments, the advent of a novel State-Space Model (Gu et al. ()), denominated as Mamba (Gu & Dao (); Dao & Gu ()), has shown substantial promise in tackling long sequence tasks, capturing significant interest within the research community. The Mamba architecture has effectively supplanted conventional Transformer frameworks in multiple domains, delivering noteworthy results. The Mamba family (Gao et al. (); Hatamizadeh & Kautz (); Lieber et al. (); Pilault et al. ()) encompasses a broad spectrum of applications, including text generation, object recognition, 3D point cloud processing, recommendation systems, and image generation, with numerous implementations based on frameworks such as Vision-Mamba (Zhu et al. ()), U-Mamba (Ma et al. ()), and Rec-Mamba (Yang et al. ()). Vision-Mamba employs a bidirectional state-space model structure in conjunction with a hybrid Transformer (Hatamizadeh & Kautz ()). However, Mamba has not yet been explored in the context of non-autoregressive image generation. Presently, the majority of Mamba-based generative tasks adhere to the diffusion model paradigm (Hu et al. (); Teng et al. ()), which entails complexities related to training and the number of inference iterations. Addressing these challenges, we have designed a novel hybrid Mamba structure aimed at extending the application of Mamba in non-autoregressive image generation (Li et al. (); Chang et al. (; )) tasks, integrating it with Masked Image Modeling (He et al. ()) for both training and inference, thereby enhancing the efficiency of these processes.

## 3 Method

### 3.1 MaskMamba Model: Overview

Overview. As illustrated in Fig.2, our MaskMamba fundamentally consists of three components. Firstly, the image pixels $x\in\mathbb{R}^{H\times W\times 3}$ are quantized into discrete tokens $q\in\mathbb{Q}^{h\times w}$ via an image tokenizer (Yu et al. (); Van Den Oord et al. (); Esser et al. ()), where $h=H/r$, $w=W/r$, and $r$ represents the downsample ratio of the image tokenizer. These discrete tokens $q\in\mathbb{Q}^{h\times w}$ serve as indices of the image codebook. Then, we randomly sample the masking ratio $m_{r}$ (range from 0.55 to 1.0), and mask out $m_{r}\cdot(h\cdot w)$ tokens, replacing them with a learnable mask token $[M]$. Secondly, we transform the class id into a learnable label embedding (Peebles & Xie (); Esser et al. ()), denoted as $\{cls\}$. On the other hand, regarding the text conditions, we first extract features using a T5-Large Encoder (Colin ()) and then map the extracted features to caption embeddings (Chen et al. ()), denoted as $\{t_{1},t_{2},\ldots,t_{N}\}$. Lastly, we concat condition embeddings $\{cond\}$ with the image token embeddings $\{q_{1},q_{2},\ldots,q_{h\cdot w}\}$ at middle, where $\{cond\}$ represents $\{cls\}$ or $\{t_{1},t_{2},\ldots\}$, and add positional embedding to these $\{q_{1},q_{2},\ldots,[M],\ldots,q_{i},cond,q_{j},[M],\ldots,q_{h\cdot w}\}$. The training objective is to predict the token indices of the masked regions utilizing cross-entropy loss (Zhang & Sabuncu ()).

Figure: Figure 2: MaskMamba Pipeline Overview.
Refer to caption: x2.png

Model Configuration. We present two types of image generation models: class-conditional and text-conditional models. In accordance with the standards established by prior work (Radford et al. (); Touvron et al. ()), we adhere to the standard configurations for the Mamba. As shown in Tab.1, we provide three different versions of the class-conditional model, with parameter sizes ranging from 103M to 741M. The generated images have a resolution of $256\times 256$, and after a downsampling factor of 16, the length of the image token embeddings is set to 256. The length of the class-condition embedding is set to 1, and the length of the text-condition embedding $N$ is set to 120.

**Table 1: Model sizes and configurations of MaskMamba.**
| Type | Model | Params | Layers | Hidden dim |
| --- | --- | --- | --- | --- |
|  | MaskMamba-B | 103M | 12 | 768 |
| C2I | MaskMamba-L | 329M | 24 | 1024 |
|  | MaskMamba-XL | 741M | 36 | 1280 |
| T2I | MaskMamba-XL | 742M | 36 | 1280 |

### 3.2 MaskMamba Model: Architecture

#### 3.2.1 Bi-Mamba-V2 Layer.

Convolution Replacement. As illustrated in Fig.3 (c), we redesign the original Bi-Mamba (Zhu et al. ()) architecture to better accommodate tasks associated with masked image generation. We substitute the original causal convolution with a standard convolution. Given the non-autoregressive nature of masked image generation task, the causal convolution only permits unidirectional token mixing, which hinders the potential of non-autoregressive image generation. In contrast, the standard convolution enables tokens to interact bidirectionally across all positions in the input sequence, effectively capturing the global context.

Symmetric SSM Branch Design. We incorporate a symmetric SSM branch to better accommodate masked image generation. In the symmetric branch, we first flip the input $x$ before the Backward SSM, then flip it back after the Backward SSM to amalgamate it with the results of the Forward SSM. Additionally, compared to the right-side branch of Bi-Mamba, we employ an extra convolution to mitigate feature loss. To fully exploit the advantages of all the branches, we project the input into a feature space of size $C/2$, thereby ensuring that the final concatenated dimensions are consistent. Our output can be denoted as $X_{out}$, which is computed using the following Eq.1.

Figure: Figure 3: (a) Structure of the original Mamba (Gu & Dao ()). (b) Bi-Mamba structure proposed in VisionMamba (Zhu et al. ()), which introduces a new branch specifically designed for vision tasks. (c) Our redesigned Mamba for masked image generation tasks by using standard convolution instead of causal convolution and replacing the final-stage multiplication with concatenation to reduce computation.
Refer to caption: x3.png

$$ $\displaystyle x=\sigma\left(\operatorname{Conv}\left(\operatorname{Linear} \left(C,{C/2}\right)\left(X_{\text{in }}\right)\right)\right)$ (1) $\displaystyle x1=x,\quad x2=\operatorname{Flip}(x)$ $\displaystyle\hat{x}=\operatorname{ForwardSSM}\left(x1\right)+\operatorname{ Flip}\left(\operatorname{BackwardSSM}\left(x2\right)\right)$ $\displaystyle z=\sigma\left(\operatorname{Conv}\left(\operatorname{Linear} \left(C,{C/2}\right)\left(X_{\text{in }}\right)\right)\right)$ $\displaystyle X_{\text{out }}=\operatorname{Linear}\left({C},C\right)\left( \operatorname{Concat}\left(\hat{x},z\right)\right)$ $$

#### 3.2.2 Maskmaba Hybrid Scheme.

Group Scheme Design. As displayed in Fig.4 (a) and Fig.4 (b), we design two group mixing schemes. In group scheme v1, we divide the input into two groups along the channel dimension, which are then processed separately by our Bi-Mamba-v2 layer and Transformer layer. We then concatenate the processed results along the channel dimension and finally feed them into the Norm and Project layers. In group scheme v2, we divide the input into four groups along the channel dimension. Two of these groups are processed by our Bi-Mamba-v2 layer in the Forward SSM and the Backward SSM, while the other two groups are processed by the Transformer layer.

Serial Scheme Design. As shown in Fig.4 (c) and Fig.4 (d), we also design two serial mixing schemes. In the serial scheme v1, we alternate layer-by-layer arrangements of our Bi-Mamba-v2 and Transformer. In the serial scheme v2, we place our Bi-Mamba-v2 in the first $N/2$ layers and the Transformer in the last $N/2$ layers. Due to the attention mechanism of the Transformer, which can better enhance feature representation, we place the Transformer layer after the Mamba layer in all serial modes.

Figure: Figure 4: We design two categories of four hybrid configurations: grouped parallel and cascading serial. In parallel, the model is divided into two and four groups. In serial, we use a layer-wise interleaved structure of Bi-Mamba-v2 and Transformer, or the first $N/2$ layers are Bi-Mamba-v2 followed by $N/2$ layers of Transformer.
Refer to caption: x4.png

### 3.3 Image Generation By MaskMamba

We utilize masked image generation (Li et al. (); Chang et al. ()) methods for image synthesis. For a generation resolution of $256\times 256$ with a downsampling factor of 16, During the forward pass, we first initialize 256 masked tokens. Subsequently, we concat the condition embeddings with mask tokens at the middle position. Drawing inspiration from the iterative generation approach of MUSE (Chang et al. ()), our decoding process also adopts a cosine schedule (Chang et al. ()) that chooses a fixed proportion of the highest confidence masked tokens for prediction at each step. These tokens are then set unmasked for remaining steps and the set of masked tokens is correspondingly reduced. Through this methodology, we can infer 256 tokens utilizing merely 20 decoding steps, in contrast to the 256 steps necessitated by autoregressive methods (Touvron et al. (); Sun et al. ()).

Class-conditional image generation. The label embeddings are based on the index of each category. These label embeddings are concatenated with the masked tokens and MaskMamba gradually predicts these mask tokens through a cosine schedule.

Text-conditional image generation. We first extract text features using a T5-Large Encoder (Colin ()) and then transform the extracted features to caption embeddings. Similar to the label embeddings, we concatenate these caption embeddings with the masked token embeddings. MaskMamba gradually predicts these mask tokens through a cosine schedule.

Classifier-free guidance image generation. The classifier-free guidance (CFG) method proposed by diffusion models (Ho & Salimans ()) is a highly effective technique for enhancing the conditional generation capabilities of models, particularly in handling text and image features. Thus, we apply this approach to our model. During the training phase, to simulate the process of unconditional image generation, we randomly drop the condition embeddings with a probability of 0.1. In the inference phase, the logit $\ell_{g}$ for each token is determined by the following equation: $\ell_{g}=(1-s)\ell_{u}+s\left(\ell_{c}\right)$,where $\ell_{u}$ is uncondition logit, $\ell_{c}$ is condition logit and $s$ is scale of the CFG.

## 4 Experimental Results

### 4.1 Class-conditional Image Generation

Training Setup. All of class-to-image generation models are trained for 300 epochs on the ImageNet $256\times 256$ dataset, with consistent training parameter settings across all models. Specifically, the base learning rate is set to 1e-4 per 256 batch size, and the global batch size is 1024. Additionally, we employ the AdamW optimizer with $\beta_{1}$ = 0.9 and $\beta_{2}$ = 0.95. The dropout rate is consistently set to, including for conditions. During training, the mask rate varies from 0.5 to 1. All training and inference of the models are conducted on V100 GPUs with 32GB of memory.

Evaluation Metrics. We use FID-50K (Heusel et al. ()) as the primary evaluation metric, while employing Inception Score (Salimans et al. ()) (IS) and Inception Score standard deviation (IS-std) as assessment criteria. On the ImageNet validation dataset, we generate 50,000 images based on the CFG and evaluate all models using the aforementioned metrics.

#### 4.1.1 Qualitative Results

Figure: Figure 5: Examples of class-conditional image generation using MaskMamba-L (left) and MaskMamba-XL (right) with cfg=3.0, iterations=25.
Refer to caption: x5.png

Comparisons with Other Image Generation Methods. As shown in Tab.2, we compare our MaskMamba model with popular image generation models, including autoregressive (AR) methods (Esser et al. (); Sun et al. ()), mask-prediction models (Mask) (Li et al. (); Chang et al. ()), and Transformer-based models (Masked Image Modeling training with the same hyperparameters), focusing on the differences in their backbone networks. The MaskMamba uses the serial scheme v2 mode. Comparisons across various model sizes show that MaskMamba exhibits competitive performance. As illustrated in Fig.5, we randomly select images from MaskMamba-XL models demonstrate high-quality results even when trained only on ImageNet.

**Table 2: Model comparisons on class-conditional Generation on ImageNet $256\times 256$ benchmark. We utilized FID-50K as the primary evaluation metric, supplemented by Inception Score(IS) as an auxiliary assessment criterion. During the generation process, with cfg set to 3.0.**
| Type | Model | Parameters | FID-50k$\downarrow$ | IS$\uparrow$ | Precision$\uparrow$ | Recall$\uparrow$ | Steps |
| --- | --- | --- | --- | --- | --- | --- | --- |
| AR | VQGAN (Esser et al. ()) | 227M | 18.64 | 80.4 | 0.78 | 0.26 | 256 |
| VQGAN (Esser et al. ()) | 1.4B | 15.78 | 74.3 | - | - | 256 |  |
| LlamaGen-B (Sun et al. ()) | 111M | 8.69 | 124.33 | 0.78 | 0.46 | 256 |  |
| Mask | MAGE-B (Li et al. ()) | 200M | 11.10 | 81.17 | - | - | 20 |
| MAGE-L (Li et al. ()) | 463M | 9.10 | 105.1 | - | - | 20 |  |
| MaskGIT (Chang et al. ()) | 227M | 6.18 | 182.1 | 0.83 | 0.57 | 10 |  |
| Transformer-B | 101M | 11.72 | 90.11 | 0.73 | 0.50 | 25 |  |
| Transformer-L | 324M | 7.08 | 127.44 | 0.76 | 0.55 | 25 |  |
| Transformer-XL | 736M | 5.96 | 140.81 | 0.75 | 0.58 | 25 |  |
| MaskMamba-B | 103M | 10.88 | 89.84 | 0.70 | 0.55 | 25 |  |
| MaskMamba-L | 329M | 6.61 | 127.74 | 0.73 | 0.59 | 25 |  |
| MaskMamba-XL | 741M | 5.79 | 139.30 | 0.73 | 0.60 | 25 |  |

#### 4.1.2 Experiment Analysis

Effect of Class-Free Guidance(CFG) and generation interations. Fig.6 (a) shows FID and IS variations with the number of iterations in image generation with cfg set to 3. The model achieves best performance at 25 iterations and further increasing iterations would deteriorate FID. Fig.6 (b) shows FID and IS scores for different cfg settings, indicating that while class-free guidance enhances visual quality and while cfg=3 the model achieves best performance.

Figure: Figure 6: Fig.6 (a) shows the variation of FID and IS with respect to the number of generation iterations, while Fig.6 (b) presents the scores of FID and IS under different cfg settings.
Refer to caption: x6.png

Efficency Analysis. We conduct a series of experiments to evaluate the effectiveness of our redesigned Bi-Mamba-v2 layer, the original Bi-Mamba layer, and the Transformer layer. To assess inference experiments on higher resolution images, we primarily focus on inference speed and memory usage with a single layer. All experiments on efficiency analysis are conducted on an A100 40G device, and we compare the inference speed of these models at different resolutions, as shown in Fig. 7. The results indicate that when the resolution is less than $1024\times 1024$, our Bi-Mamba-v2 layer and the Bi-Mamba layer are slightly slower than the Transformer layer. However, when the resolution exceeds $1024\times 1024$, our Bi-Mamba-v2 layer is faster than both the Transformer and Bi-Mamba layer. Notably, at a resolution of $2048\times 2048$, our Bi-Mamba-v2 layer is 1.5 times faster than the Transformer layer. We also compare GPU memory usage at different batch sizes. The memory usage of our Bi-Mamba-v2 layer is comparable to that of the Bi-Mamba layer, while the Transformer layer, due to its quadratic complexity, exhibits a rapid increase in memory usage as the batch size increases. When the batch size reaches 6, the Transformer layer consumes 63GB of GPU memory, leading to out of memory, while our Bi-Mamba-v2 layer requires only 38GB. These experimental results demonstrate that our Bi-Mamba-v2 Layer can generate images at a faster speed and with lower memory usage.

Figure: Figure 7: The inference speed (left) and GPU memory (right) usage of the Bi-Mamba-v2 layer, original Bi-Mamba layer, and Transformer layer are assessed across various image resolutions and batch sizes.
Refer to caption: x7.png

Effect of different hybrid schemes. As indicated in Tab.3, we perform a comparative analysis of the image generation outcomes under various hybrid configurations of MaskMamba, classified into two categories: parallel and serial. As depicted in Fig.4, in the grouped parallel configurations, we examine the effects of dividing the model into two and four groups. In the layered serial configurations, we design an interleaved structure of Bi-Mamba-v2 and Transformer {MSMS…MSMS}, as well as an alternative configuration {MMMM…SSSS} where the first $N/2$ layers are Mamba and the subsequent $N/2$ layers are Transformer. The findings from these experiments elucidate the performance and efficiency of the different hybrid configurations.

**Table 3: Comparison of different hybrid schemes in MaskMamba on class-conditional Generation on ImageNet $256\times 256$ benchmark. The inference settings are configured with cfg set to 3 and iterations set to 20 for all experiments.**
| Model | Scheme | Parameters | FID-50k $\downarrow$ | IS $\uparrow$ |
| --- | --- | --- | --- | --- |
| MaskMamba-L-Group-v1 | - | 327M | 10.04 | 96.35 |
| MaskMamba-L-Group-v2 | - | 278M | 8.95 | 102.72 |
| MaskMamba-L-Serial-v1 | MSMS…MSMS | 329M | 7.45 | 115.90 |
| MaskMamba-L-Serial-v2 | MMMM…SSSS | 329M | 6.73 | 122.99 |

Effect of different backbone. We conduct ablation experiments on different backbones: the Bi-Mamba proposed in VisionMamba (Zhu et al. ()), the redesigned Bi-Mamba-V2, and the Transformer (Vaswani ()). Bi-Mamba-L uses only the original Bi-Mamba as a layer, while Bi-Mamba-V2-L uses our redesigned Bi-Mamba-v2. The Transformer uses only the Transformer architecture. In (Bi-Mamba + Transformer)-L, the first $N/2$ layers are original Bi-Mamba, followed by $N/2$ layers of the Transformer. In (Bi-Mamba-V2 + Transformer)-L, the first $N/2$ layers are Bi-Mamba-v2, followed by $N/2$ layers of the Transformer. Results show our redesigned Bi-Mamba-v2 enhances performance over the original Bi-Mamba, and combining Mamba and Transformer further improves results. Therefore, we choose (Bi-Mamba-V2 + Transformer) for MaskMamba.

**Table 4: Comparison of different backbone in MaskMamba on class-conditional Generation on ImageNet $256\times 256$ benchmark. Bi-Mamba-V2 refer to the redesigned Bi-Mamba, which reuses the original bidirectional mamba. The inference settings are configured with cfg set to 3 and iterations set to 25 for all experiments.**
| Backbone | Parameters | FID-50k $\downarrow$ | IS $\uparrow$ |
| --- | --- | --- | --- |
| Bi-Mamba-L | 377M | 12.29 | 87.39 |
| Bi-Mamba-V2-L | 333M | 8.97 | 103.97 |
| Transformer-L | 324M | 7.08 | 127.44 |
| (Bi-Mamba + Transformer)-L | 358M | 7.82 | 110.87 |
| (Bi-Mamba-V2 + Transformer)-L | 329M | 6.61 | 127.74 |

Effect of different condition embedding positions. We conduct ablation experiments to assess the impact of the placement of condition embedding $cond$ on model performance. Specifically, we examine the effects of concating the condition embedding at different positions of sequence, such as the head, middle, and tail. The experimental results indicate that optimal performance is achieved when the condition embedding is placed in the middle. This outcome is primarily attributed to the mechanism of selective scan. Given that we randomly mask a portion of the image tokens, placing the condition embedding at the head or tail results in insufficient supervision information for conditional generation control due to the increased attention distance.

**Table 5: Comparison of different condition embedding postions in sequence on class-conditional generation on ImageNet $256\times 256$ benchmark. The inference settings are configured with cfg set to 3 and iterations set to 20 for all experiments.**
| Condition embedding postions | Scheme | FID-50k $\downarrow$ | IS $\uparrow$ |
| --- | --- | --- | --- |
| Head | $\langle\textbf{cond},q_{1},q_{2},\ldots,\ldots,q_{h\cdot w}\rangle$ | 7.15 | 117.71 |
| Middle | $\langle q_{1},q_{2},\ldots,\textbf{cond},\ldots,q_{h\cdot w}\rangle$ | 6.73 | 122.99 |
| Tail | $\langle q_{1},q_{2},\ldots,\ldots,q_{h\cdot w},\textbf{cond}\rangle$ | 7.04 | 119.78 |

### 4.2 Text-conditional Image Generation

Training setup. Similar to the class-conditional training strategy, we adapt a Masked Generative non-autoregressive training strategy for the text. We train the model for 30 epochs on the CC3M (Sharma et al. ()) dataset with the image resolution $256\times 256$. The training parameters are consistent with those of the previous experiments, the base learning rate is set to 1e-4 per 256 batch size, and the global batch size is 1024. Additionally, we employ the AdamW optimizer with $\beta_{1}$ = 0.9 and $\beta_{2}$ = 0.95.

**Table 6: Comparison of different backbone in MaskMamba on Text-conditional Generation on CC3M (Sharma et al. ()) and MS-COCO (Lin et al. ()) $256\times 256$ benchmark. The inference settings are configured with cfg set to 3 and iterations set to 20 for all experiments.**
| Backbone | Parameters | FID-CC3M-10K $\downarrow$ | IS $\uparrow$ | FID-COCO-30K $\downarrow$ | IS $\uparrow$ |
| --- | --- | --- | --- | --- | --- |
| Transformer-XL | 736M | 19.20 | 14.98 | 43.21 | 14.44 |
| MaskMamba-L | 741M | 18.11 | 16.90 | 25.93 | 18.34 |

Figure: Figure 8: Visualizing the MaskMamba-XL model’s performance in text-conditional image generation. We randomly select the image generation structure from the validation set of 10,000 images in CC3M, using the text prompts directly from the CC3M dataset.
Refer to caption: x8.png

Model trained on CC3M. As shown in Table 6, we compare the performance of Transformer-XL and our MaskMamba-XL in text-to-image generation, evaluating the FID and IS on the validation sets of CC3M and MS-COCO. Our results consistently outperform the Transformer-based model. As displayed in Figure 8, we utilize text from CC3M as prompts to generate images. MaskMamba-XL is capable of producing high-quality images. However, due to the limited training data and the imprecision of the text descriptions in the CC3M dataset, some of the generated images exhibit limitations.

## 5 Conclusion.

In this work, we propose MaskMamba, a novel hybrid model that combines Mamba and Transformer architectures, utilizing Masked Image Modeling for non-autoregressive image synthesis. We not only redesign a new Bi-Mamba structure to make it more suitable for image generation but also investigate the effects of different model mixing strategies and the placement of condition embeddings, ultimately identifying the optimal settings. Additionally, we provide a series of class-conditional image generation models and text-conditional image generation models in a single framework with an in-context condition. Our experiment results indicate that our MaskMamba surpasses both Transformer-based and Mamba-based models in terms of generation quality and inference speed. We hope our Masked Image Modeling for non-autoregressive image synthesis in MaskMamba can inspire further exploration in Mamba image generation tasks.

## Appendix A Appendix

### A.1 Base Model Configurations

Our MaskMamba Training configurations is given in Tab.7.

**Table 7: Training hyperparameters for MaskMamba.**
| Configuration | Value |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Optimizer | AdamW |  |  |  |  |
| Optimizer momentum | $\beta_{1}$ = 0.9,$\beta_{2}$ = 0.95 |  |  |  |  |
| Base learning rate | 1e-4 |  |  |  |  |
| Learning rate schedule | cosine decay |  |  |  |  |
| Training epochs | 300 |  |  |  |  |
| Warmup epochs | 3 0 |  |  |  |  |
| Weight decay | 0.05 |  |  |  |  |
| EMA | 0.999 |  |  |  |  |
| Mask ratio min | 0.5 |  |  |  |  |
| Mask ratio max | 1.0 |  |  |  |  |

### A.2 Examples of class-conditional image generation using MaskMamba

Figure: Figure 9: Examples of class-conditional image generation using MaskMamba-B with cfg=3.0, iterations=25.
Refer to caption: x9.png

Figure: Figure 10: Examples of class-conditional image generation using MaskMamba-L with cfg=3.0, iterations=25.
Refer to caption: x10.png

Figure: Figure 11: Examples of class-conditional image generation using MaskMamba-XL with cfg=3.0, iterations=25.
Refer to caption: x11.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: -
soundness: 2
strengths: This paper is generally in good shape, and the reported results are promising. For example, the structure is well-defined and smoothly covers all changes (albeit with various degrees of detail). Multiple experiments and comparisons are conducted to evaluate the proposed model and visual results and code of the proposed block are provided in the appendix.
summary: This paper introduces MaskMamba, hybrid model combining Mamba and Transformer architectures, designed for non-autoregressive image synthesis through Masked Image Modeling. It adapts bidirectional Mamba by (1) replacing causal with standard convolutions to improve global context capture and (2) using concatenation instead of multiplication to boost performance and speed up inference. The model also explores different hybrid configurations, including serial and grouped parallel arrangements. Experiments are conducted on class conditional and text conditional image generation.
weaknesses: There are several issues with the paper as currently presented to be considered for a top tier conference:

- Text-to-image models cannot be accurately evaluated with image metrics alone—CLIP score, ImageReward, and additional benchmarks like TIFA and T2I-CompBench are necessary for a comprehensive assessment.
- The proposed method consistently has more parameters than the transformer baselines, yet these baselines don’t receive any benefit of the doubt for this disparity.
- The baseline models are notably weak, limited to outdated versions, and lack diffusion models, resulting in an incomplete evaluation framework. 
- The rationale for the group scheme and serial scheme designs is missing—they appear abruptly with no background on the issues they address, the process that led to their design, or previous attempts that were unsuccessful.
- The description of the 20-step approach as "non-autoregressive" is misleading; while it may not be causal, it’s also not a 1-step approach, which creates confusion.
- The paper mentions that convolutions "impose constraints that hinder scalability," but lacks clarification, details, or citations to support this claim. In particular, the paper still puts an emphasis on using convolutions.
- There’s insufficient explanation on the distinctions between forward and backward SSM, as well as on the transition from bi to bi-v2, leaving the motivations for these changes unclear.
- It’s unclear if a consistent dropout value is applied; if so, the specific value should be specified.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please answer the question I mentioned in "weakness".
soundness: 2
strengths: 1. The writing logic of this paper is very clear and easy to understand, from the rationality of raising questions to making improvements.
2. The referencing speed has been greatly improved.
summary: In this paper, a new bidirectional Mamba structure is designed to improve the inference speed, and several hybrid Mamba schemes are explored.
weaknesses: 1. The contribution of the paper is relatively limited, and replacing multiplication with concatenation would obviously improve inference speed, but its rationality and motivation should be properly explained, and the change should also be compared in detail with Bi-Mamba.
2. The paper mentions that it can complete the text-to-image task simultaneously, but there is no comparison of its performance on this task throughout the paper, which raises questions about its performance.
3. The comparison shown in Figure 5 reveals huge gaps in certain parameters, including IS and Precision, for this model.
4. A more detailed comparison with Bi-Mamba should be conducted, and Precision and Recall

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: Please address my concerns in the "Weaknesses" Section.
soundness: 3
strengths: - It is very reasonable to introduce Mamba or SSM layers in high-resolution, MaskGIT-based image generation. SSM models naturally have better capability handling long sequences compared with transformer models.
- The ablation study provides valuable insights on the combination of Mamba blocks and transformer blocks (MMM..SSS is better than MS...MS). 
- The authors also show promising signals that the proposed method can scale up to text-to-image generation.
summary: This paper introduced a new non-autoregressive model for image generation. It replaces the transformer architecture in MaskGIT with MaskMamba, where some of the transformer layers are replaced with a new Mamba layer with forward and backward SSM. The authors show reasonable image generation results on ImageNet, as well as proof-of-concept speedup results at 2048x2048.
weaknesses: 1. Experiment results on ImageNet are not strong enough.
- 5.79 FID at 741M parameters is far from the state-of-the-art (for example, MAGVITv2). 
- The cited baselines are also a bit weak. MaskGIT can achieve much better results using a better tokenizer and training recipe. See [here](https://github.com/baaivision/MUSE-Pytorch/tree/master) and [here](https://github.com/bytedance/1d-tokenizer/tree/main) for more details. It is also recommended to have a look at Open-MAGVITv2's tokenizer, which should be orthogonal to this work and could push the results closer to SOTA.

2. There is a little bit of overclaim in terms of the speedup over transformer in the abstract. "54% faster" refers to the comparison between a single Mamba layer and a single transformer layer, not the end-to-end speed comparison. In fact, MaskMamba is a hybrid model and there are many transformer layers inside the entire model as well. It would be misleading for the readers when reading the last sentence in the abstract. 

3. Although introducing SSMs in non-autoregressive image generation models is interesting, the paper lacks technical novelty. It seems to me that this work is combining MaskGIT with Jamba, a hybrid transformer-Mamba model for language.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: n/a
soundness: 2
strengths: * Interesting architecture, which seems to bring interesting improvements on inference speed.
* Detailed analysis of architectural changes
* Tested on both class conditioned and text condition image generation
* Well written paper, easy to follow
summary: The authors build upon the Mamba architecture to tackle the subject of image generation, class conditions and text conditioned, using a masking approach inspired from MaskGIT. The model is trained to predict mask tokens and at inference, only a fraction of the tokens are unmasked, iteratively, until all tokens are predicted.

Specifically, the authors propose:
* architecture changes to the Mamba architecture (causal convolutions are replaced with full convolutions, flipping the backward SSM branch)
* combining the Mamba blocks for the first N/2 layers with transformer blocks at the last N/2 layers

The new architecture shows:
* Similar or slightly better results than a full transformer of similar size
* Faster inference, up to 1.5X faster on larger image size and larger models
weaknesses: * The architecture seems a bit cumbersome:
  * Having to care about the position of the conditioning (the authors show that it is important) is a major downside compared to transformer based architectures
  * Tuning this architecture is likely harder, because of the mix of convolutions and transformers, each with different optimization patterns (simplicity is a quality)
* The boost of performance is not that important, -0.17 FID on XL architecture, and highly depends on the position of the conditioning, Tail or Head conditioning in Table 5 do lead to similar or worse performance than transformer in Table 4. The main benefit is mostly inference speed, which is something that can be tuned for transformers as well and more of a problem for video generation (it could be super interesting to test it there)
* The baseline approach is far from SOTA (for example, MAR gets much lower FID on ImageNet 256, and much faster inference speed) and so it would be more impactful to apply this to MAR-like approach to see if that still brings benefits.
* If anything, the paper shows that transformers are needed for image generation, as shown in Table 4, defeating a bit the interest in Mamba.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Q1: The text condition length is set to 120. Can this be variable?

Q2: There seems to be an inconsistency between Figure 4(a) and the description of bi-mamba-v2 on line 249. Could you please align these?

Q3: The reason CFG works is due to the matching of two probability distributions. However, in AR, VAR, and MAR, the cross-entropy loss is not about matching two probability distributions. Why does CFG still work in this case? Could there be a possible explanation?

Q4: In Table 2, the comparison lacks similar works involving VAR [1] and MAR [2], as well as comparisons with CFG-based diffusion models. Additionally, several diffusion models [3, 4, 5] based on the Mamba architecture are not discussed in this paper. Could you elaborate further on the advantages of MaskMamba in terms of performance and speed? As I review the paper, it appears that MaskMamba still requires 20+ steps to generate.

Q5: Within 1024 resolution, single-layer bi-mamba-v2 is slower than the transformer. This time difference could accumulate with additional layers, which is not a positive indicator. The claim is that bi-mamba-v2 is faster beyond 1024px, but the experiments were only conducted on ImageNet at 256px. It remains unclear how performance scales from 256px to 512px, and then to higher resolutions such as 1024px and 2048px.

[1]. Tian K, Jiang Y, Yuan Z, et al. Visual autoregressive modeling: Scalable image generation via next-scale prediction[J]. arXiv preprint arXiv:2404.02905, 2024.
[2]. Li T, Tian Y, Li H, et al. Autoregressive Image Generation without Vector Quantization[J]. arXiv preprint arXiv:2406.11838, 2024.
[3]. Teng Y, Wu Y, Shi H, et al. DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis[J]. arXiv preprint arXiv:2405.14224, 2024.
[4]. Fei Z, Fan M, Yu C, et al. Dimba: Transformer-Mamba Diffusion Models[J]. arXiv preprint arXiv:2406.01159, 2024.
[5]. Hu V T, Baumann S A, Gui M, et al. Zigma: Zigzag mamba diffusion model[J]. arXiv preprint arXiv:2403.13802, 2024.
soundness: 2
strengths: The writing is clear and easy to understand. Building on MaskGit, the paper explores the combination of the Mamba architecture with various Mamba module designs and its integration with transformers. Experiments demonstrate that MaskMamba achieves solid performance on both ImageNet and CC3M.
summary: The writing is clear and easy to understand. Building on MaskGit, the paper explores the combination of the Mamba architecture with various Mamba module designs and its integration with transformers. Experiments demonstrate that MaskMamba achieves solid performance on both ImageNet and CC3M.

However, the paper lacks a more comprehensive comparison, including related methods such as VAR, MAR, and diffusion models. Additionally, comparisons with other Mamba-based approaches in the field of image generation are missing. In terms of experimental design, evaluations are primarily conducted at 256 resolution, where MaskMamba is slower than transformers. However, results are not presented at higher resolutions (above 1024), where MaskMamba is claimed to have an advantage.
weaknesses: The paper lacks a more comprehensive comparison, including related methods such as VAR, MAR, and diffusion models. Additionally, comparisons with other Mamba-based approaches in the field of image generation are missing. In terms of experimental design, evaluations are primarily conducted at 256 resolution, where MaskMamba is slower than transformers. However, results are not presented at higher resolutions (above 1024), where MaskMamba is claimed to have an advantage.

Meta-review: The goal of this paper is to introduce a new non-autoregressive masked generative model approach that leverages advantages of both Mamba and Transformer architectures, while reducing their limitations in generation quality and inference efficiency. In order to build this new hybrid model, the authors first addressed inefficiencies with the previously introduced Bi-Mamba model, including replacing causal convolutions with standard convolutions to improve global context, and multiplcation with concatenation for the final stage of Mamba. Then, to incorporate class- and text-to-image synthesis, the MaskMamba approach combines Mamba and Transformer layers, exploring a range of different design choices. Experiments indicate that MaskMamba outperforms Transformer and Mamba models on standard benchmarks.

Strengths:
The paper is well-written and comprehensive in methodological detail and experiments. The use of Mamba layers and making them more efficient and compatible with Transformers is a useful direction to explore for efficiency reasons. It’s also exciting to see that there is potential in scaling this kind of hybrid (non-diffusion/non-autoregressive) method to text-image generation.


Weaknesses:
Experiments are not conducted with a full suite of state of the art baselines, which makes it difficult to determine how the model performs more comprehensively. Also, reviewers feel that some of the design choices could have been better motivated. In addition, the modeling contributions are fairly narrow, which places more emphasis on the need for strong empirical results and analysis - reviewers felt the empirical results in particular on imagenet are not strong enough to be impactful on their own. Results without further comparisons against state of the art may not move the needle enough to get much interest since current SoTA (including standard Transformer-based approaches) are still considerably stronger.

Decision reasoning:
While the paper has been improved with feedback from reviewers, who engaged with the authors in back and forth during the discussion period, reviewers unanimously felt that the paper in it’s current form is not quite there in terms of suitability for ICLR publication. This is mostly due to lack of results against strong baselines, and lack of results at high resolution, where one might expect to see advantages of the Mamba style approach more clearly. Given these concerns, I feel it is important for the authors to continue to improve the paper with more comparisons with state of the art methods, or to include deeper analysis into the method choices (or both). I do think the method is in good shape for a workshop contribution to get visibility and feedback from the community but needs more work for a full conference publication.

Decision: Reject