Autoregressive Pretraining with Mamba in Vision

## Abstract

Abstract The vision community has started to build with the recently developed state space model, Mamba, as the new backbone for a range of tasks. This paper shows that Mamba’s visual capability can be significantly enhanced through autoregressive pretraining, a direction not previously explored. Efficiency-wise, the autoregressive nature can well capitalize on the Mamba’s unidirectional recurrent structure, enabling faster overall training speed compared to other training strategies like mask modeling. Performance-wise, autoregressive pretraining equips the Mamba architecture with markedly higher accuracy over its supervised-trained counterparts and, more importantly, successfully unlocks its scaling potential to large and even huge model sizes. For example, with autoregressive pretraining, a base-size Mamba attains 83.2% ImageNet accuracy, outperforming its supervised counterpart by 2.0%; our huge-size Mamba, the largest Vision Mamba to date, attains 85.0% ImageNet accuracy (85.5% when finetuned with 384 × 384 384 384 384\times 384 384 × 384 inputs), notably surpassing all other Mamba variants in vision. The code is available at https://github.com/OliverRensu/ARM .

## 1 Introduction

In natural language processing (NLP), state space models (SSMs) [, , , ] demonstrate strong potential for modeling long sequences with linear complexity. Among these, a recent variant, Mamba [], has substantially advanced beyond traditional SSMs by synthesizing the best attributes of selective scanning.
This innovation has also catalyzed its rapid adoption within the vision community, leading to its application across diverse visual tasks. These include the design of novel architectures [, , , , ], applications to segmentation [, , ] and image synthesis [].

However, these prior studies are mostly in the setting of supervised visual representation learning. While such trained models exhibit promising results in different visual tasks, they generally suffer from limited transferability and encounter notable difficulties in scaling [, , , ].
For example, as illustrated in Figure 1, attempts to scale the Vision Mamba (Vim) under supervised conditions often lead to either performance plateauing or even training collapse when pushed to very large sizes. These issues, therefore, motivate us to alternatively explore self-supervised visual representation learning with Mamba architectures, a method that has demonstrated notable successes in helping models secure strong and scalable visual representations [, , , ].

In this paper, we primarily focus on the autoregressive pretraining paradigm for self-supervised visual representation learning, which predicts the next token unidirectionally and autoregressively from the start to the end of the input sequence. This focus is driven by two reasons. First, autoregressive pretraining has already established itself as the de-facto standard in training large language models, with successful applications in various architectures including Transformers and Mamba [, , ]. The recent literature has also successfully, albeit preliminarily, confirmed its efficacy in the computer vision domain, *e.g*., helping Vision Transformer (ViT) develop strong and scalable feature representations [, ].
Secondly, Mamba architectures are inherently well-suited for autoregressive modeling due to their uniquely designed linear attention nature, which methodically constructs token-wise relationships in a strictly progressive and unidirectional manner. This configuration ensures that each token can only attend to its preceding tokens, aligning perfectly with the underlying principles of autoregressive modeling. Additionally, this synergy practically leads to higher overall training efficiency. For example, under the setting of training the base-size Mamba for 300 epochs, autoregressive training requires only $\scriptstyle\sim$34 hours (measured by 8$\times$A5000), a $\scriptstyle\sim$$2\times$ to $\scriptstyle\sim$$10\times$ improvement in training speed compared to other pretraining strategies (see Table 8 in Sec. 4.4).

Figure: Figure 1: Compared to Vim, our ARM considerably boosts the ImageNet accuracy and, more critically, offers a stronger pathway for scaling up.
Refer to caption: x1.png

Importantly, to further unleash the power of autoregressive visual representation learning with Mamba architectures, we highlight two key recipes for forming input sequences.
First, instead of naively taking $16\times 16$ patches as basic units of prediction, we opt for a more strategic approach by grouping spatially neighboring patches to form larger clusters; empirically, we find the cluster size of $64\times 64$ reaches the best performance. Secondly, in our ablation of mapping 2D images into 1D visual sentences with various orderings, we note that vanilla ordering, which simply orders clusters with the row-by-row and forward scan approach, is already an effective choice. We term this method ARM.

Extensive results are provided showing our proposed ARM achieves substantially stronger performance. As shown in Figure 1, ARM helps our base-size model attain 83.2% ImageNet accuracy, outperforming the supervised counterpart by 2.0%. Moreover,
ARM enables the training of the first successful huge-size model (ARM-H), marking it as the largest vision Mamba model to date. Specifically, ARM-H achieves an impressive 85.0% ImageNet accuracy, significantly outperforming all previous Mamba variants. Additionally, ARM also improves the performance on out-of-domain datasets by a large margin: ARM-B outperforms supervised Vim-B by 4.4% on ImageNet-A, 2.9% on ImageNet-R, and 3.3% on ImageNet-S.

## 2 Related Work

##### State space model.

The state space model (SSM) [] stands as a novel alternative to Transformers for long-range dependency modeling with linear complexity.
Linear attention [, , ] recurrently approximates self-attention via a softmax-free attention matrix with linear complexity, which can be viewed as a degenerate linear SSM.
The Structured State-Space Sequence (S4) model [] computes more efficiently than prior approaches while preserving their theoretical strengths based on a new parameterization.
S5 [] extends S4 by adding multi-input multi-output (MIMO) SSM and efficient parallel scan.
RWKV [] is a recent RNN with its key “WKV” components that operate similarly to a system with two SSMs. Its updated version [] incorporates state expansion and input-dependent gating for more flexible sequence modeling.
Following this, Mamba [] proposes a data-dependent SSM layer with hidden state expansion and builds a generic language model backbone, which performs comparably to transformers at various sizes and enjoys linear scaling in sequence length.
This work focuses on Mamba in vision, aiming to enhance it via autoregressive visual pretraining.

##### Mamba in vision.

The successful application of Mamba in NLP has inspired its adoption in vision applications. Vision Mamba (Vim) [] utilizes Vim blocks composed of pure Mamba layers: each Vim block leverages both forward and backward scans to model bidirectional representations and mitigate the direction-sensitive problem in Mamba.
Alternatively, Vmamba [] employs Visual State Space (VSS) blocks that integrate both Mamba and 2D convolution layers, supported by a pyramid architecture akin to the Swin Transformer []: each VSS block first models 2D local information via 2D depth-wise convolution as the token mixer, followed by a CrossScan Module that processes 2D global information both horizontally and vertically.
Mamba-ND [] further expands Mamba’s capabilities to multi-dimensional data, including images and videos.
LocalMamba [] splits the input image into several local windows and performs SSM in various directions within these windows, enhancing local processing. EfficientVMamba [] introduces an efficient 2D scanning technique using atrous sampling on feature map patches to reduce computational demands.
Compared to these newly designed Mamba architectures, ours is less novel, which closely follows the design of ViT, but substituting the self-attention with the Mamba module. With this naive Mamba architecture, our main focus is to show autoregressive pretraining can enhance its visual capabilities.

##### Self-supervised visual representation learning.

Self-supervised visual representation learning aims to learn strong and transferable representations without labels, including contrastive learning [, , , ], position prediction [], masked image modeling [, , ], *etc*. This paper focuses on autoregressive pretraining,
which is highly successful in NLP but still less explored in computer vision. iGPT [] is the first work to introduce Generative Pretrained Transformer to vision and highlights the potential of autoregressive pretraining as a general self-supervised visual representation learning strategy.
SAIM [] and RandSAC [] further enhance autoregressive pretraining, achieving performance on par with MAE [] by utilizing the ViT architecture and a stochastic sequence permutation strategy. D-iGPT [] slightly modifies the learning objective to predict not only the next token but also visible tokens. AIM [] demonstrates that, with autoregressive pretraining, ViT scales effectively with increased model capacity and data quantity. Different from these prior works, which focus on Transformer architecture, we provide the first study of exploring autoregressive visual pretraining with Mamba architectures.

## 3 Method

### 3.1 Mamba Preliminaries

The Mamba architecture inherits from state space sequence models [], which models a $1$-D function or sequence $x(t)\in\mathbb{R}\to y(t)\in\mathbb{R}$ at time $t$ via expanded hidden states $h_{t}\in\mathbb{R}^{N}$. The hidden state is evolved through time driven by parameters $\mathbf{A},\mathbf{B},\mathbf{C}$ following linear ordinary differential equations (ODEs):

$$ $\begin{split}h^{\prime}(t)&=\textbf{A}h(t)+\textbf{B}x(t),\\ y(t)&=\textbf{C}h(t).\end{split}$ (1) $$

To discretize parameters in this continuous system, a common solution is to introduce a time scale parameter $\mathbf{\Delta}$ to transform continuous $\textbf{A},\textbf{B}$ to discrete $\overline{\textbf{A}},\overline{\textbf{B}}$ using zero-order hold (ZOH) model []:

$$ $\displaystyle\overline{\textbf{A}}$ $\displaystyle=\exp(\mathbf{\Delta}\textbf{A}),$ (2) $\displaystyle\overline{\textbf{B}}$ $\displaystyle=(\mathbf{\Delta}\textbf{A})^{-1}(\exp(\mathbf{\Delta}\textbf{A}) -\textbf{I})\cdot\mathbf{\Delta}\textbf{B}.$ $$

By applying such transformation, we can rewrite Eq. 1 as:

$$ $\displaystyle h^{\prime}(t)$ $\displaystyle=\overline{\textbf{A}}h_{t-1}+\overline{\textbf{B}}x_{t},$ (3) $\displaystyle y_{t}$ $\displaystyle=\textbf{C}h_{t}.$ $$

We then employ a matrix $\overline{\textbf{K}}$ for fast computation:

$$ $\displaystyle\overline{\textbf{K}}$ $\displaystyle=(\textbf{C}\overline{\textbf{B}},\textbf{C}\overline{\textbf{AB} },...,\textbf{C}\overline{\textbf{A}}^{k}\overline{\textbf{B}},...),$ (4) y $\displaystyle=\textbf{x}*\overline{\textbf{K}},$ $$

where $k\in[0,L)$ and $L$ is the input sequence length. We also have $\textbf{y}=\{y_{1},...,y_{L}\}$, $\textbf{x}=\{x_{1},...,x_{L}\}$, while $\overline{\textbf{K}}\in\mathbb{R}^{L}$ can be regarded as the convolutional kernel. Note this computing structure allows Mamba to model the input sequence that perfectly matches the unidirectional, next-word prediction in autoregressive modeling.

### 3.2 Autoregressive pretraining

We first briefly revisit autoregressive pretraining in NLP. Then, we shift our attention to autoregressive pretraining with mamba in vision, including the prediction unit and prediction order design. Lastly, we present the model variants.

#### 3.2.1 Autoregressive Pretraining in NLP

Figure: Figure 2: Different prediction units in the autoregressive modeling.
Refer to caption: x2.png

Autoregressive pretraining models the probability of the next word one by one given a corpus $\mathcal{U}=\{u_{1},...,u_{n}\}$. This can be formulated as:

$$ $p(u)=\prod\limits_{i=1}^{n}p(u_{i}|u_{1},...,u_{i-1},\Theta)$ (5) $$

Here, autoregressive pertaining computes the likelihood of each word $u_{i}$ based on the context of all preceding words from $u_{1}$ to $u_{i-1}$ and minimizes the negative log-likelihood:

$$ $\mathcal{L}=-log~{}p(u)$ (6) $$

This strategy plays a fundamental role in training large language models like ChatGPT [] and GPT-4 [] in NLP.

#### 3.2.2 Autoregressive Pretraining with Mamba in Vision

##### Prediction unit.

Transitioning from 1D sentences to 2D images introduces the challenge of defining a suitable autoregressive prediction unit. We start with the vanilla strategy presented in iGPT [] which considers each individual pixel as the prediction unit, as illustrated in Figure 2(b). For an image $X=\{p_{1},...,p_{n}\}$, our objective is to minimize the loss function:

$$ $\centering\begin{split}&\mathcal{L}=\sum_{i=1}^{n-1}l(f([p_{1},...,p_{i}]),p_{ i+1}),\\ &l(\hat{y},y)=|\hat{y}-y|^{2}.\end{split}\@add@centering$ (7) $$

Here $f(\cdot)$ denotes the Mamba model, and $p_{i}$ represents the $i_{th}$ pixel of the image. This pixel-based approach, while straightforward, imposes significant computational demands, particularly for high-resolution images. Therefore, as shown in the original iGPT paper [], this constraint necessitates the use of low-resolution images for computationally feasible autoregressive pretraining.

Patchifying [] images into non-overlapped regions and then mapping them into visual tokens can address this computation challenge. For example, with an image size of 224$\times$224, the sequence length would reduce significantly from 50,176 in the iGPT framework to just 196 patches with the $16\times 16$ patchifying operation. Intuitively, shifting the prediction unit from pixels [] to patches [, , ], as shown in Figure 2(c), adjusts the autoregressive input to $X=\{P_{1},...,P_{n}\}$:

$$ $\begin{split}&\mathcal{L}=\sum_{i=1}^{n-1}l(f([P_{1},...,P_{i}]),P_{i+1}),\\ &l(\hat{y},y)=|\hat{y}-y|^{2}.\end{split}$ (8) $$

Here $P_{i}\in\mathcal{R}^{16\times 16}$ is the $i_{th}$ patch. Moreover, to encapsulate the 2D spatial information at the token level, we propose grouping spatially adjacent patches into larger clusters to serve as the prediction unit, illustrated in Figure 2(d). The clustered input $X=\{c_{1},...,c_{n}\}$ aims to be optimized by:

$$ $\begin{split}&\mathcal{L}_{\text{ARM}}=\sum_{i=1}^{n-1}l(f([c_{1},...,c_{i}]), c_{i+1}),\\ &l(\hat{y},y)=|\hat{y}-y|^{2}.\end{split}$ (9) $$

Here, each $c_{i}\in\mathcal{R}^{H_{c}\times W_{c}}$ is a cluster formed by grouping $\frac{H_{c}}{16}\times\frac{W_{c}}{16}$ patches. Our ablation studies (Section 4.4, Table 5) show that using clusters as prediction targets significantly enhances performance compared to the use of individual pixels or patches. Next, we explore the strategies for sequencing these clusters into a coherent visual sentence.

Figure: Figure 3: Different prediction orderings of a visual sentence.
Refer to caption: x3.png

##### Prediction order.

Unlike the 1D sentences in NLP, which inherently have a clear sequence order for autoregressive modeling, we hereby explore four different prediction orders when projecting 2D images into 1D visual sentences, *e.g*., how these clusters should be arranged given a cluster size of $s$, with $\frac{W}{s}$ clusters per row and $\frac{H}{s}$ clusters per column.
We hereby explore four primary prediction orders: 1) *Row-first and forward* orders the clusters row by row, processing from the first to the last cluster within each row sequentially, as depicted in Figure 3(a). 2) *Row-first and backward* similarly orders the clusters row by row but inverts the processing direction, starting with the last cluster and moving to the first within each row, illustrated in Figure 3(b). 3) *Column-first and forward* organizes the clusters column by column, processing sequentially within each column from top to bottom, shown in Figure 3(c). 4) *Column-first and backward* similarly sequences the clusters column by column but starts with the bottom-most cluster, moving upwards, as seen in Figure 3(c).
To consider an approach free from pre-defined sequential biases, we also experimented with a *Random* permutation [] of cluster order, visualized in Figure 3(e).

Detailed empirical comparisons of these four predefined orders alongside the random order are presented in Section 4.4. Our findings reveal that while the predefined orders exhibit minimal differences in performance, employing a random order leads to severe performance degradation. Consequently, the straightforward and effective *row-first and forward order* (Figure 3(a)) is adopted as our standard ordering strategy for autoregressive modeling.

Figure: Figure 4: The comparison of block architectures between Vim, and MambaMLP in pretraining and in finetuning.
Refer to caption: x4.png

### 3.3 MambaMLP

We hereby introduce our newly developed MambaMLP blocks. Specifically, our MambaMLP block uses Mamba as the token mixer and the multi-layer perceptron (MLP) as the channel mixer, drawing inspiration from the self-attention block in Transformer [, ].
Note that the configuration of the MambaMLP block varies between pretraining and fine-tuning phases to cater to their different requirements. During pretraining, as illustrated in Figure 4(b), the MambaMLP block contains the Mamba layer with only 1 scan [] to match the uni-directional modeling manner in autoregressive pertaining; while in finetuning (displayed in Figure 4(c)), the block is then adapted to contains the Mamba layer with 4 scans, thus enabling bi-directional modeling of global information analogous to that in Vmamba []. The other architectural components in the pretraining and the finetuning stay the same: the block utilizes SwiGLU [] as the MLP layer, and the $expand$ is set to 1 to enhance scanning efficiency.
Additionally, we provide a visual comparison between our MambaMLP block and the Vim block in Figure 4. We can see that the Vim block contains Mamba layers with 2 scans [] for bi-directional global information processing and has no MLP layer, and the $expand$ of each scan is set to 2. Practically, this larger $expand$ in each scan results in higher performance but slower inference speeds.

By stacking multiple MambaMLP blocks and training with our autoregressive strategy developed in Section 3.2.2, we name the resulting model ARM. As detailed in Table 1, ARM is designed to match the depth and width of ViT in its base and large configurations. For the huge model size, ARM adopts the structure of AIM-600M [], which is wider but less deep compared to ViT-H, balancing performance and computational efficiency. In the next section, we will extensively validate the efficacy of ARM.

**Table 1: The configuration of different architecture variants.**
| Model | Block | Width | Depth | Param.(M) |
| --- | --- | --- | --- | --- |
| ViT-B | (Attention+MLP) | 768 | 12 | 86 |
| Vim-B | Mamba | 768 | 24 | 98 |
| ARM-B | (Mamba+MLP) | 768 | 12 | 85 |
| ViT-L | (Attention+MLP) | 1024 | 24 | 307 |
| Vim-L | Mamba | 1024 | 48 | 340 |
| ARM-L | (Mamba+MLP) | 1024 | 24 | 297 |
| ViT-H | (Attention+MLP) | 1280 | 32 | 632 |
| Vim-H | Mamba | 1536 | 48 | 755 |
| ARM-H | (Mamba+MLP) | 1536 | 24 | 662 |

## 4 Experiment

### 4.1 Implementation Details

##### Pretraining.

We pretrain ARM using the ImageNet-1K dataset []. Specifically, ARM-B and ARM-L are pre-trained for 1600 epochs, and ARM-H is pre-trained for 800 epochs. We
use a batch size of 2048/1024/512 for ARM-B/L/H, respectively, and a learning rate of lr = 1.5e-4$\times\frac{\mathrm{batchsize}}{256}$. We adopt a cosine decay schedule with
a warm-up for 5 epochs. We adopt the AdamW [] optimizer with a weight decay of 0.05. We use random resized cropping and random horizontal flipping. The pretraining input size is set to $192\times 192$.

##### Finetuning.

Following pretraining, we finetune the ARM models on the ImageNet classification task. Specifically, we finetune all models for 100 epochs with a batch size of 1024, with the input size set at $224\times 224$. We use the same data augmentation as MAE []. We adopt AdamW as an optimizer, and the peak learning rate is lr=5e-4$\times\frac{\mathrm{batchsize}}{256}$ with a cosine decay schedule and a warm-up for 5 epochs. Additionally, we employ the exponential moving average (EMA) [] for stronger performance.

Further, we evaluate model robustness on various out-of-domain ImageNet variants, including natural adversarial examples (ImageNet-A []), semantic shifts (ImageNet-R []), image sketches (ImageNet-S []), ImageNet-V2 [], and ImageNet-Real [].

**Table 2: Performance comparison on ImageNet-1K. Throughputs are measured with an A5000 GPU. $\dagger$ denotes we extend the training of Vim to the large-size model, using its original GitHub repo. $\dagger$ indicates the stride is 8. Hybrid architectures are in Gray.**
| Model | Token Mixer | Image Size | Param. | Throughputs | Top-1 |
| --- | --- | --- | --- | --- | --- |
| (M) | (imgs/s) | (%) |  |  |  |
| Base-size models |  |  |  |  |  |
| RegNetY-16G | 2D Conv. | 224^2 | 84 | 870 | 82.9 |
| DeiT-B | Attention | 224^2 | 21 | 1073 | 81.2 |
| Vim-B$\dagger$ | Mamba | 224^2 | 98 | 890 | 81.2 |
| MambaMLP-B | Mamba | 224^2 | 85 | 1301 | 81.2 |
| VMamba-B | Mamba+2D Conv. | 224^2 | 89 | 315 | 83.9 |
| ARM-B | Mamba | 224^2 | 85 | 1301 | 83.2 |
| ARM-B | Mamba | 384^2 | 85 | 440 | 84.2 |
| ARM-B $\ddagger$ | Mamba | 448^2 | 85 | 86 | 84.8 |
| Large-size models |  |  |  |  |  |
| Vim-L$\dagger$ | Mamba | 224^2 | 340 | 345 | 81.0 |
| MambaMLP | Mamba | 224^2 | 297 | 445 | 81.4 |
| ARM-L | Mamba | 224^2 | 297 | 445 | 84.5 |
| ARM-L | Mamba | 384^2 | 297 | 154 | 85.1 |
| Huge-size models |  |  |  |  |  |
| Vim-H$\dagger$ | Mamba | 224^2 | 755 | 211 | collapsed |
| ARM-H | Mamba | 224^2 | 662 | 275 | 85.0 |
| ARM-H | Mamba | 384^2 | 662 | 94 | 85.5 |

### 4.2 Main Results

In Table 2, we compare our ARM with convolution-based RegNet [], Attention-based ViT, and different Mamba architectures in vision. For the base-size model, our ARM achieves 83.2% accuracy, making a substantial 2.0% improvement over its supervised MambaMLP counterpart. Additionally, we note that ARM outperforms Vim by 2.0%, and is the only Mamba architecture that attains stronger performance than convolution-based RegNetY-16G (*i.e*., by 0.3%). Further enhancements are observed when ARM-B is finetuned with increased input sizes of 384$\times$384 and 448$\times$448 with the patchify stride of 8, where performance improves to 84.2% and 84.8%, respectively. We also report the comparison to VMamba-B, which takes a hybrid architecture: When configured with inputs of 224×224, ARM-B slightly underperforms VMamba-B by 0.7% but enjoys a much faster throughput, *i.e*., $\scriptstyle\sim$$4\times$ faster; ARM-B with the inputs of 384$\times$384 outperforms Vmamba-B by 0.3% and still maintains a faster throughput, *i.e*., 440 imgs/s *vs*. 315 imgs/s.

Next, we scale the Mamba architectures to much larger model sizes. First, we observe that Mamba-based Vim sees a performance dip with the large size and fails to train stably at the huge size. This observation suggests that these prior Mamba-based architectures grapple with scaling challenges. Contrarily, ARM models excel in scalability — ARM-L achieves an accuracy of 84.5%, marking a 3.5% improvement over Vim-L, and ARM-H sets a new benchmark for the largest Mamba architecture in vision to date by reaching 85.0% accuracy. Moreover, by tuning ARM at a larger resolution of $384\times 384$, further leveraging the model’s capacity to handle long sequences at a linear complexity, we observe additional gains: a 0.6% increase with ARM-L and a 0.5% increase with ARM-H. Notably, ARM-H attains the best Mamba accuracy of 85.5% on ImageNet classification.

### 4.3 Robustness and Generalization

We report the robustness evaluation of Mamba architectures in Table 3. We can observe that ARM consistently shows much stronger robustness than the supervised Vim by, *e.g*., ARM-B exhibits
improvements ranging from 1.8% to 4.4% over supervised Vim-B across these robustness benchmarks. More impressively, ARM-L extends these gains even further, showing enhancements ranging between 2.6% and 7.4% when compared to supervised Vim-L. In addition, ARM-H, our largest model variant, not only continues this trend but also shows an average performance superiority of 1.1% over ARM-L, reaffirming the efficacy of scaling up the model size on enhancing robustness.

**Table 3: Robustness and Generalization evaluation on out-of-domain datasets.**
| Method | IN-1K $\uparrow$ | IN-V2 $\uparrow$ | IN-Real $\uparrow$ | IN-Adv.$\uparrow$ | IN-Ren.$\uparrow$ | IN-Ske.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| Vim-S [] | 80.6 | 69.4 | 86.0 | 20.3 | 45.8 | 33.4 |
| Vim-B [] | 81.2 | 70.0 | 86.2 | 27.5 | 46.0 | 33.9 |
| ARM-B | 83.2 | 72.3 | 88.0 | 31.9 | 48.9 | 37.2 |
| Vim-L [] | 81.0 | 69.8 | 86.0 | 27.9 | 44.7 | 31.8 |
| ARM-L | 84.5 | 74.0 | 88.6 | 41.4 | 52.1 | 39.2 |
| ARM-H | 85.0 | 75.6 | 89.2 | 42.3 | 53.2 | 40.5 |

**Table 4: Ablation on the number of predictions units.**
| Num of Prediction unit | Cluster size | Top-1 (%) |
| --- | --- | --- |
| 0 (Supervised) | N/A | 81.2 |
| 144 (iGPT) | $1\times 1$ (Pixel) | 79.8 |
| 4 | 96$\times$96 | 82.0 |
| 9 | 64$\times$64 | 82.5 |
| 16 | 48$\times$48 | 82.2 |
| 36 | 32$\times$32 | 81.9 |
| 144 | 16$\times$16 | 81.7 |

### 4.4 Ablation Study

This section provides different ablations on ARM. Unless otherwise specified, all ablation studies are performed on ARM-B under 300 epochs pretraining.

##### Number of prediction units.

Table 5 reports the ablation on the number of prediction units. We start from the cluster size equal to the patch size (*i.e*., each cluster contains only one patch), resulting in a total of 144 prediction units. We note that, even with this vanilla setup, autoregressive pretraining successfully helps MambaMLP improve performance from 81.2% (via supervised training) to 81.7%. Then, we gradually group multiple patches into one cluster, thereby reducing the total number of prediction units. We note that the performance first increases and then decreases — the best performance is achieved when the number of the prediction units is set to 9, corresponding to a cluster size of 64$\times$64. Specifically, this setup provides a performance improvement of 1.3% over the supervised counterpart and 0.8% over the vanilla autoregressive pretrained counterpart (*i.e*., with a cluster size of 144). We also report the comparison to MambaMLP trained under the iGPT-style autoregressive pretraining — with the input image size at $144\times 144$ and setting per pixel as the prediction unit, it underperforms our best setup by 2.7% (*i.e*., 79.8% *vs*. 82.5%).

##### Prediction Order.

As shown in Table 5, we find different pre-defined orders only lead to minor performance variances. For example, both row-first and column-first forward prediction orders achieve an identical performance of 82.5%; even the least favorable case, where the prediction order was row-first and backward, only underperforms the best case by 0.2%. Nonetheless, interestingly, if we do not predefine the prediction order and pick a random permutation, the performance significantly drops to 81.5%.

##### Decoder Design.

Our exploration into decoder design is summarized in Table 7. We first focus on the design of decoder depth, finding that increasing the depth up to 4 progressively enhanced performance up to 82.5%; further increasing the decoder depth to 8 sees a performance saturation.
With this 4-layer decoder setup, we next study the width of the decoder. By ablating these three options {384, 512, 1024}, we empirically observe that setting the decoder depth to 512 yields optimal accuracy.

##### Prediction targets.

We hereby explore different prediction targets for our ARM. By default, we use per-patch normalized pixels with mean square error (MSE) loss. For comparison, we ablate it against two setups: 1) unnormed pixels with MSE loss, and 2) discretized tokens of the patches derived from dVAE [] with cross-entropy loss. The results, presented in Table 7, show that employing normalized pixels as the target with MSE loss yields the best performance, achieving an accuracy of 82.5%. Comparatively, this configuration outperforms the model using discrete tokens from dVAE by 0.3% and the model leveraging unnormed pixels which trailed by 0.6%.

**Table 6: Ablation on decoder designs.**
| Dec. Depth | Dec. Width | Top-1 (%) |
| --- | --- | --- |
| 1 | 512 | 82.1 |
| 2 | 512 | 82.4 |
| 4 | 512 | 82.5 |
| 8 | 512 | 82.5 |
| 4 | 384 | 82.3 |
| 4 | 512 | 82.5 |
| 4 | 1024 | 82.2 |

##### Pretraining paradigm.

As shown in Table 8, we evaluate different pretraining paradigms, including contrastive learning [], MAE [], and our ARM. Firstly, we note that all pretraining methods result in performance gains over the supervised counterpart, demonstrating the benefits of self-supervised visual pretraining on Mamba architectures. However, using MAE or contrastive learning, the performance is only moderately improved by 0.4% and 0.2%, respectively, over the supervised baseline. In contrast, our ARM achieves significant improvements of 1.3% over the supervised baseline, as well as achieves higher accuracy than both contrastive learning and MAE. Additionally, in terms of efficiency, ARM requires just 34 hours of pretraining, cutting the training duration in half compared to MAE, which is already noted for its relatively low pretraining demands.

**Table 8: Comparison of architecture and pretraining paradigms. FPS represents the inference speed after supervised finetuning of the model. The $\dagger$ symbol indicates that Vim, when subjected to contrastive learning, experiences poor performance, potentially due to mode collapse.**
| Architecture | Pretraining paradigm | Training Cost (h) $\downarrow$ | FPS (imgs/s) $\uparrow$ | Top-1 (%) |
| --- | --- | --- | --- | --- |
| MambaMLP | Supervised | 110 | 1330 | 81.2 |
| MambaMLP | Contrastive | 330 | 1330 | 81.4 |
| MambaMLP | MAE | 70 | 1330 | 81.6 |
| MambaMLP | ARM | 34 | 1330 | 82.5 |
| Vim | Supervised | 165 | 923 | 81.2 |
| Vim | Contrastive | 510 | 923 | 80.2$\dagger$ |
| Vim | MAE | 106 | 923 | 81.4 |
| Vim | ARM | 57 | 923 | 82.2 |

##### Architecture design.

Exploring further into architectural impacts, Table 8 (from the 5th row to the 8th row) presents our investigation into whether Vim, another variant within the Mamba architecture, benefits from autoregressive pretraining. Results indicate a positive response as ARM-trained Vim reaches an 82.2% accuracy on ImageNet, marking a 1.0% improvement over its supervised-only counterpart. Contrastingly, other pretraining paradigms did not fare as well for Vim: when subjected to contrastive learning, Vim experiences training instability, falling below the supervised baseline; MAE pretraining on Vim only slightly improved over the supervised method, with a marginal gain of 0.2%. These results further support the effectiveness of ARM in pretraining Mamba in Vision.

As a side note, it is important to highlight that although Vim’s performance improves with ARM pretraining, it operates $\scriptstyle\sim$45% slower during inference compared to MambaMLP. Additionally, MambaMLP incurs only $\scriptstyle\sim$66% of the training cost required for pretraining Vim under the ARM framework. These points underscore the superior efficiency of our default ARM framework, particularly when paired with our newly developed MambaMLP architecture.

## 5 Conclusion

In this study, we introduced a novel autoregressive visual pretraining strategy tailored for Mamba architectures, known as ARM. This approach enhances pretraining efficiency and effectiveness by strategically treating groups of spatially neighboring image patches as prediction units. Through our method, we have significantly improved the scalability and benchmark performance of Mamba-based models, setting new standards in their operational functionality. We hope this work can lay a strong foundation for future explorations and potential expansions in the usage of autoregressive pretraining strategies for Mamba architectures within the vision community.

## Acknowledge

This work is partially supported by the AWS Cloud Credit for Research program.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: - As shown in MAE, pre-training has another benefit to improve downstream tasks (like object detection, segmentation). Does this conclusion hold for ARM?
- It would be better to discuss a comparison with a prior work, “Image as First-Order Norm+Linear Autoregression: Unveiling Mathematical Invariance” (https://arxiv.org/pdf/2305.16319), which uses first-order autoregression to pre-train CNN architectures.
soundness: 3
strengths: - The paper is well written, including all technical details. It should be straightforward to implement.
- It is good to see Mamba follow similar pre-training conclusion with other vision architecture (CNN, Transformer): masked prediction in pre-training followed by fine-tuning boosts performance.
- It is good to know that cluster size plays an important role in pre-training.
summary: This paper shows that Vision Mamba can be enhanced through autoregressive pre-training, similar to using MAE pre-training to improve ViT. It uncovers two key recipes: (a) grouping spatially neighboring patches, and (b) vanilla ordering (row-by-row scan) is effective enough. It achieves significant improvement over supervised baselines.
weaknesses: - Figure 3 is a bit confusing. It seems to illustrate a case of 4 clusters, each has 4 patches. I assume that the arrows within a cluster show the order to flat patches, while the arrows across clusters show the order of prediction. If so, please use different colors and add explanation in the caption.
- In Table 2, please distinguish methods between supervised and pre-training+supervised fine-tuning. In addition, what is the performance for MambaMLP in huge size?
- Section 4.4, is it possible to improve random-order prediction by incorporating position encoding?

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: See weaknesses
soundness: 3
strengths: This paper is overall writing well. This work to investigates the mamba architecture using autoregressive pre-training in visual data is important for the community. 
The results are inspiring for the community, e.g., ARM can scale somewhat well, compared to the Vim counterparts; and the ablation on the prediction unit shows the effectiveness of the cluster-based prediction.
summary: This paper investigates the Mamba architecture for visual data using autoregressive pre-training. It shows the Mamba’s visual capability can be enhanced through autoregressive pretraining. It finds autoregressive pretraining equips the Mamba architecture can obtain better performance compared to its supervised-trained counterparts, and also shows the potential for scaling, based on the experiments on ImageNet-1K.
weaknesses: 1.My main concern is the experimental setting:   
(1)	When the results of the proposed method in Table 2 is good, it is not clear for the setting up of other supervised baselines (RegNetY-16G, DeiT-B, Vim-B), e.g, how many epochs these model training? Is it fair to compare these baselines without the same training epochs?  
(2)	Besides, the protocol for evaluating self-supervised method is usually pre-trained on large-scale dataset, then evaluated on the downstream tasks (e.g., linear probe or finetuned with few-shot examples). It is weird that this paper pre-trained on ImageNet using self-supervised learning and fine-tuned with all the ImageNet examples. For a stronger result, I think this paper should pretrain on larger unlabeled dataset, then finetuned on ImageNet. In current experiments, the results are not strong, since Vits using self-supervised learning has shows its potential in performance and scaling ability. Another point, I think this paper should compare the state-of-the-art Vits with the same parameters in Table.2 (the parameters matters for performance), and I donot find the proposed method have better performance than Vits. 


2.It is not clear why this paper uses different setting for pre-training (Fig.4 b) and fine-tuning (Fig.4 c) for the MambaMLP architecture. Is it from the empirical observation? If yes, it is better to show the results, when unifying the setting (e.g., using the same setting for pre-training and finetuning). It is weird that MambaMLP needs two configurations for pre-training and fine-tuning, especially in the situation that this paper only conducts experiments on one dataset (not generable but only tailored for one dataset)

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See above.  Additional, more minor questions/comments below.

* Additional related work on autoregression in vision might be mentioned and compared, e.g.  Bai et al 2023 "Sequential Modeling Enables Scalable Learning for Large Vision Models"; Yu et al 2021 "Diverse image inpainting with bidirectional and autoregressive transformers"; Esser et al 2021 "Taming transformers for high-resolution image synthesis."

* Could use better background summary for what is meant by "scan" and "expand" in Fig 4 / Sec 3.3.  While the basic SSM is summarized, as are the orderings in sec 3.2, a little more explanation of the "scans" here would be useful

* Sec 4.3:  When applied to other datasets, is autoregressive pretraining performed only on imagenet1k, the target dataset, or both?

* Eq (3) says $h'(t)$ but this is the descretized system, so should this be $h_t$ on the first line ?

* Random order l.294:  To confirm, the order is a single random permutation order that is fixed for each model instance (so every time the model is called on an image, the same order is always used)?

* typo on l.470, should say width is 512 no depth
soundness: 3
strengths: The experiments clearly show the benefit of autoregressive pretraining for this task, on both MambaMLP (the proposed architecture here) and Vim.  Ablations on token cluster size (a nice idea), decoder and patch representations, and sequence order describe the system behavior well.  Overall, the presentation is quite clear and the system well described, though there are key points that I think could be explained better (see weaknesses below).
summary: This paper applies autoregressive pretraining to Mamba for visual classification tasks.  The Mamba-based architecture uses a block similar to Visual Mamba (VIM) and VMamba, MambaMLP, that applies Mamba over a patch sequence ordered by a 2-level "clustered" raster order and MLP over channels (though, I may have misunderstood the channel vs sequence decomposition, see questions below).  Multiple ablations are performed on Imagenet1k, profiling the system's behavior and performance; substantial improvements are found from using autoregressive pretraining on Imagenet1K and several of its variants.
weaknesses: There are several important pieces of the method that lacked detailed explanation:

* The MambaMLP block is described only in a sketch diagram (Fig 4) and a sentence in sec 3.3 ("uses Mamba as the token mixer and MLP as the channel mixer").  This is a good summary, but didn't explain the block architecture well enough for me to understand the details:  Does this mean Mamba is applied independently to each channel and alternates with MLP?  And how are the Delta, B and C set (which inputs and which channels)?  A more explicit explanation, making use of the summary of Mamba in Sec 3.1 and showing where and how this is used in MambaMLP in addition to MLP would be useful here.

* Likewise, there is no explanation of the patch decoder --- there is an ablation on its design, but I didn't find anything saying what architecture was used.

* Another key point, is how can only one scan order be used in pretraining but four in fine-tuning?  If the other 3 scan orders aren't pretrained, how are they initialized?  And what is the impact of using only one scan order (the one that was pretrained) in fine-tuning (and conversely, is it possible somehow to pretrain all four orders)?  (The terms "scan" and "expand" are assumed and not summarized or used in an explanatory context, either)

A second, more minor weakness is that only image classification is considered, and no other tasks.  Additional task applications would strengthen the results, though the existing ones with classification are already convincing to show the merits of the method.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 4.4.4 Which normalization method are you referring to for pixel targets? I’m guessing you’re referring to MAE’s [3] method, but this needs to be explicitly cited.

Do you have any experiments where you just vary the random seed to get a sense of the variance in model quality just due to randomness? For example, there are 6 examples of the word “significant”, but I assume this is not statistical significance. Taking table 4 as an example, there is a spread of 0.8% between a 1x1 cluster and a 4x4 cluster. Knowing the variance of your method on Top-1% would be really helpful for determining whether these difference are actually significant.

[3] Masked Autoencoders Are Scalable Vision Learners (http://arxiv.org/abs/2111.06377)
soundness: 2
strengths: The paper finds an effective way to apply AIM to the space of vision-mamba variants, demonstrating that it improves the pareto frontier for mamba-based vision models. The pretraining unlocking the use of larger vision-mamba variants is a useful outcome.
summary: The paper introduces autoregressive pretraining to the domain of vision mamba models. Drawing inspiration from "Scalable Pre-training of Large Autoregressive Image Models" (AIM), it shows how the method naturally adapts to Mamba models, which are inherently autoregressive, at least in their simplest form. This pretraining method shows positive scaling traits, improving downstream classification accuracy, and also allowing for larger variants to be trained than previously possible.
weaknesses: At a high level, the novelty in this paper is quite limited. It appears to mostly be “AIM [1], but with Mamba”. Similarly, I think table 2 needs to include AIM; I know this is difficult because they don’t provide finetuned accuracy numbers (only probes), but they do provide the model checkpoints, so even a casual attempt at finetuning their 0.6B model on imagenet would be helpful to contextualize on the huge models.

Given that ARM isn’t bringing anything new to the theory of SSSMs or autoregressive pretraining, it doesn’t feel like equations 1-6 are necessary, or at least could be greatly condensed [1, 2], . Particularly given the degree to which this paper builds on AIM, leveraging their section 4.1 would be better. Similarly, equations 7-9 are highly redundant. Distilling all of this down, all of the equations presented in this paper don’t seem to aid in clarity.

“Prediction units”. This is defined in 3.2.2 as generally being patches, aside from iGPT which uses pixels. This seems to be reinforced by figure 2d and figure 3. However, section 4.4.1 confusingly switches to clusters being prediction units. Based on the figures, it seems like a patch is always a prediction unit. But a good amount of the paper is spent on processing order, with a distinction between tokens and clusters. Equation 9 seems to suggest that a full cluster is predicted simultaneously. Presumably ARM is patch based though, so it seems like clusters are mostly just a way to define a particular order of patches, and due to scan being an enforced sequential algorithm, a cluster isn’t actually predicted simultaneously. This would make figure 3 make sense.

Figure 4: I’m struggling to find anywhere where the switch from 1 scan to 4 scans is justified. It would be good to at least demonstrate the difference in accuracy between 1 scan and 4 in your setting.

Table 4 has clarity issues. It would be helpful to directly specify the cluster size instead of having it be $\frac{cluster\\_size}{patch\\_size}$.

[1] Scalable Pre-training of Large Autoregressive Image Models (http://arxiv.org/abs/2401.08541)

[2] Mamba: Linear-Time Sequence Modeling with Selective State Spaces (https://arxiv.org/abs/2312.00752)

Meta-review: The paper focuses on pre-training SSL strategies for MAMBA, showing that it is an effective strategy similar to ViT architectures. Most reviewers are positive about the paper, although there are some concerns regarding whether the comparisons with the best SSL competitors are sufficient. The authors acknowledge that the method is competitive but not necessarily better than the existing baselines, and they claim that what is important is to present an alternative architectural strategy compared to the standard ViT. I completely agree with this, thus I recommend acceptance.

Decision: Accept (Poster)