Autoregressive Pretraining with Mamba in Vision
Authors: Sucheng Ren
Xianhang Li
Haoqin Tu
Feng Wang
Fangxun Shu
Lei Zhang
Jieru Mei
Linjie Yang
Peng Wang
Heng Wang
Alan Yuille
Cihang Xie
Venue: ICLR 2025 Poster (2025)
Keywords: Auto regressive Pretraining
Primary area: foundation or frontier models, including LLMs

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\times384$ inputs), notably surpassing all other Mamba variants in vision. The code is available at \url{https://github.com/OliverRensu/ARM}.

Published as a conference paper at ICLR 2025
AUTOREGRESSIVE PRETRAINING WITH MAMBA IN VI- SION
Sucheng Ren'
Xianhang Li
Haoqin Tu Feng Wang'
Fangxun Shu
Lei Zhang Cihang Xie²
Jieru Mei Linjie Yang"
Peng Wang²
Heng Wang²
Alan Yuille'
Johns Hopkins University
UC Santa Cruz
3 Alibaba Group
+UCSD
ByteDance
ABSTRACT
The vision community has started to build with the recently developed state space model, Mamba, the new backbone for range of tasks. This paper shows that Mamba's visual capability can be significantly enhanced through autoregressive pretraining, direction not previously explored. Efficiency-wiss the autoregres- sive nature can well capitalize on the Mamba's 's unidirectional recurrent structure, enabling faster overall training speed compared other training strategies like mask modeling. Performance-win wiss,, autoregressive pretraining equips the Mamba architecture with markedly higher accuracy over its supervised-traime parts and, more importantly, successfully unlocks its scaling potential to large and even huge model sizes. For example, with autoregressive pretraining, base-size Mamba attains 83.2% ImageNet accuracy, outperforming its supervised counter- part by 2.0%; our huge-size Mamba, the largest Vision Mamba to date, attains 85.0% ImageNet accuracy (85.5% when finetuned with 384 x 384 inputs), no- tably surpassing all other Mamba variants in vision. The code is available at
counter-
INTRODUCTION
In natural language processing (NLP), state space models (SSMs) (Gu et al. 2021a b. Mehta et al. 2022| Gu et al.) (2022) demonstrate strong potential for modeling long sequences with linear complex- ity. Among these, recent variant, Mamba (Gu & Dao, (023), 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 (Liu et al. 2024b Zhu et al. 2024
Huang et al. 2024 Pei et al. 2024 Wang et al. 2024a), applications to segmentation (Liu et al. 2024a Wang et al. 2024b
Xing et al. 2024) and image synthesis (Guo et al. 2024).
However, these prior studies 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 He et al. (2022): Bao et al. (2022); He et al. (2020); Chen et al.] (2020b). 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, method that has demonstrated notable successes in helping models secure strong and scalable visual
representations (He al. 2022
Bao al. 2022 He et
2020 Chen al. 20206).
In this paper, we primarily focus on the autoregressive pretraining paradigm for self-supervise 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 (Dosovitskiy jet al. 2020 Radford & Narasimhan 2018 Gu & Dao 2023). The recent literature has also success- fully, albeit preliminarily, confirmed its efficacy in the computer vision domain, e.g., helping Vision Transformer (ViT) develop strong and scalable feature representations (El-Nouby et al. 20244; Ren al. 2023a). Secondly, Mamba architectures inherently well-suited for autoregressive modeling

---

Published as a conference paper at ICLR 2025
Top-1 Accuracy ImageNet
ARM
ARM-H
84
Vim
ARM
ARM6B
Vimi-
Vim-Hl
Collegne
Small
Base
Large
Hugo
Model Size
Figure 1: Compared to Vim, our ARM considerably boosts the ImageNet accuracy and, more critically, offers a stronger pathway for scaling up.
due their uniquely designed linear attention nature, which methodically constructs rosen-wise relationships in 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 effi- ciency. For example, under the setting of training the base-size Mamba for 300 epochs, autoregressive training requires only ~34 hours (measured by 8 x A5000), a 42x x to 100 improvement in training speed compared other pretraining strategies (see Table in Sec. 4.5
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 x 16 patches as basic units of prediction, we opt for more strategic approach by grouping spatially neighboring patches to form larger clusters; empirically, we find the cluster size of 64 64 reaches the best performance. Secondly, in our ablation of mapping 2D images into ID 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 effective choice. We term this method ARM.
Extensive results are provided showing our proposed ARM achieves substantially stronger perfor- mance. As shown in Figure [ ARM helps our base-size model attain 83.2% ImageNet accuracy, outperforming the supervised counterpart by 2.0% and achieves 85.2% Top-1 accuracy with the input resolution of 448 448. 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 55.0% ImageNet accuracy, significantly outperforming all previous Mamba variants. Additionally, ARM also improves the performance on out-of-domain datasets by 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) (Gu et al. 202 a stands as a novel alterna- tive Transformers for long-range dependency modeling with linear complexity. Linear atten- tion (Katharopollos et al., Choromankk et al. 2020 Peng et al. 2021) recurrently approxi- mates 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 (Gu et al. (021a) computes more efficiently than prior approaches while preserving their theoretical strengths based on parameterization. S5 (Smith et al.) 2022) extends S4 by adding multi-input multi-output (MIMO) SSM and efficient parallel scan. RWKV (Peng et al. 2023) is a recent RNN with its key "WKV" components that operate similarly
system with two SSMs. Its updated version (Peng ct al. 2024) incorporates state expansion and input-dependent gating for more flexible sequence modeling. Following this, Mamba (Gu & Dao 2023) proposes a data-dependent SSM layer with
2

---

Published as a conference paper at ICLR 2025
hidden state expansion and builds generie 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) (Zhu et al., 2024) utilizes Vim blocks composed of pure Mamba layers: each Vim block leverages both forward and backward scans to model bidirectional represen- tations and mitigate the direction-sensitive problem in Mamba. Alternatively, Vmamba (Liu et al.. 2024b) employs Visual State Space (VSS) blocks that integrate both Mamba and 2D convolution layers, supported by a pyramid architecture akin to the Swin Transformer (Liu et al. 202T): each VSS block first models 2D local information via 2D depth-wiss wise convolution as the token mixer, followed by a CrossScan Module that processes 2D global information both horizontally and ver- tically. Mamba-ND (Li et al. 2024) further expands Mamba's 's capabilities to multi-dimensionn data, including images and videos. LocalMamba (Huang et al. (024) splits the input image into several local windows and performs SSM in various directions within these windows, enhancing local processing. Efficient 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.
(Pei et al. 2024) introduces an efficient 2D scanning technique
Self-supervised visual representation learning. Self-supervised visual representation learn- ing (Chen et al. 2020c; He et al.] 2020 Chen al. 2021 2020b Ren et al. 2022a 2024 Zhai et al. 2022 He et al. 2022 Bao et al. 2022) Ren et al. 2023b) aims to learn strong and transferable
representations without labels, including contrastive learning Chen al.) 2020c He et al. 2020
Chen et al. 2021 2020b), position prediction (Zhai et al. 20222, masked image modeling (He et al.
2022 Bao et al. 2022, Ren et al., 2023b), etc. This paper focuses on autoregressive pretraining.
which is highly successful in NLP but still less explored in computer vision. iGPT (Chen et al., 2020ap is the first work introduce Generative Pretrained Transformer vision and highlights the
potential of autoregressive pretraining as a general self-supervised visual representation learning strategy. SAIM (Qi et al. 2023) and RandS AC (Hua et al.] 2022) further enhance autoregressive pretraining, achieving performance on par with MAE (He et al. 2022) by utilizing the ViT architec- ture and a stochastic sequence permutation strategy. D-iGPT Ren et al. 2023a) slightly modifies the learning objective predict not only the next token but also visible tokens. AIM (El-Nouby eet al. 2024) demonstrates that, with autoregressive pretraining. ViT scales effectively with increased model capacity and data quantity. Different from these prior works, which focus Transformer architecture, we provide the first study of exploring autoregressive visual pretraining with Mamba architectures.
METHOD
3.1
MAMBA PRELIMINARIES
The Mamba architecture inherits from state space sequence models (Gu et al., 2021a), which models 1-D function or sequence xx() E IR - y(t) E R at time ł via expanded hidden states h; E R.v. The hidden state is evolved through time driven by parameters A, B, C following linear ordinary differential equations (ODEs):
h' ((()) Ah(t) + Bz(t),
(1)
y(t) Ch(t). To discretize parameters in this continuous system, a common solution is to introduce time scale parameter A to transform continuous A, B discrete A, B using zero-order hold (ZOH) model (Oppenheim et al. 1997):
Ā = exp AA);
(2)
B = (AA) (exp( (A) - I). AB.
By applying such transformation, we can rewrite Eq. aas
h1: Aht- BrG
(3)
y/ Che.
3

---

Published as a conference paper at ICLR 2025
(a) laput Image
(b) Pixel-based
(c) Patch-based
(d) Cluster-hased
Figure 2: Different prediction units in the autoregressive modeling.
We then employ matrix K for fast computation:
K = (CB, CAB
CA B...).
(4)
x= K,
where [0. L) and L is the input sequence length. We also have fy....., "u......9..... v2), (01, (..,.... æ while K E TR. can be regarded 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 vision, including the prediction unit and prediction order design. Lastly, present the model variants.
3.2.1 AUTOREGRESSIVE PRETRAINING IN NLP
Autoregressive pretraining models the probability of the next word one by one given corpus (u..... ). This can be formulated as:
p(u)
II M1 M U1-1,
(5)
Here, autoregressive pertaining computes the likelihood of each word u, based on the context of all preceding words from U1 to Us-1 and minimizes the negative likelihhooo::
log p(u)
(6)
This strategy plays fundamental role in training large language models like ChatGPT (Brown et 2020) and GPT-4 (OpenAl 2023) in NLP.
3.2.2
AUTOREGRESSIVE PRETRAINING WITH MAMBA IN VISION
Prediction unit. Transitioning from ID sentences to 2D images introduces the challenge of defining suitable autoregressive prediction unit. We start with the vanilla strategy presented in iGPT (Chen et al. 2020a) which considers each individual pixel as the prediction unit, as illustrated Figure {p1,-.. ..., P.), ), objective is to minimize the loss function:
(b). For image X
n-1
c Σ ((f(pp....... pil).P++1).
( 9) 10 2
Here f(-) denotes the Mamba model, and p; represents the its 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 (Chen et al.) (2020a), this constraint necessitates the use of low-resse resolution images for computationally feasible autoregressive pretraining.
Patchifying Dosoviiskiy et al. 2020) images into non-overlapped regions and then mapping them into visual tokens address this computation challenge. For example, with image size of
4

---

Published as a conference paper at ICLR 2025
Row-lirst Forward
Row-first buckwward
(c)
Ibt Column fins and Forward
(d) Column first and Backoardd
Figure 3: Different prediction orderings of visual sentence.
224x224, the sequence length would reduce significantly from 50,176 in the iGPT framework to
just 196 patches with the x 16 patchifying operation. Intuitively, shifting the prediction unit from
pixels (Chen et al.) 2020a) patches (Dosovitskiy al., [020 Zhu al. 2024: El-Nouby et al.
2024). as shown in Figure 2(c), adjusts the autoregressive input to [P. P
R-1
Σ 1=1
P.),
109 v) 1 11
Here P: E R16x16 is the in patch. Moreover, to encapsulate the 2D spatial information at the token level, we propose grouping spatially adjacent patches into larger clusters unit, illustrated in Figure d). d). The clustered input = fci, ...,...., Go) aims to be optimized by:
as the prediction
CARM = Σ (ff([e).... c.), ci+),,
100.0)= 10 10
He W. patches, and our model scans Here, each E R. is cluster formed by grouping 16 cluster by cluster. each cluster, our model sean patch by patch. The ablation studies (Section 4.5] Table 5) show that using clusters prediction targets significantly enhances performance compared to the use individual pixels or patches. Next, we explore the strategies for sequencing these clusters into a coherent visual sentence.
Prediction order. Unlike the ID sentences in NLP, which inherently have clear sequence order for autoregressive modeling, we hereby explore four different prediction orders when projecting 2D images into ID visual sentences, e.g how these clusters should be arranged given cluster size of s. with W. clusters per row and H. clusters per column. We hereby explore four primary prediction orders: i) 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 (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 3b.. 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-mot cluster, moving upwards, seen in Figure (... To consider an approach free from pre-defined sequential biases, we also experimented with a Random permutation (Yang et al.] 2019) of cluster order, visualized in Figure (()). Note this randomness is running the fly, i..,, the same image will be processed in different orders at different training steps. We add position embedding for all prediction order designs.
Detailed empirical comparisons of these four predefined orders alongside the random order are presented Section 4.5 Our findings reveal that while the predefined orders exhibit minimal differences in performance, employing random order leads severe performance degradation. Consequently, the straightforward and effective row-first and forward order (Figure B(a) is adopted standard ordering strategy for autoregressive modeling.
5

---

Published as a conference paper at ICLR 2025
24
x.12
SwiGL.U
SwiGLU
Scan Expand-2
Activation
Activation
Activation
Scan
Scan Expund-1
Som
Sear
Expand-2
san
Sace Exparel
Norm
Norm
(a) Vim
(b) Mam CambbuMLP-rretraiingg prettraining
(c) HambaMi.P-fintumm
Figure 4: The comparison of block architectures between Vim, and MambaMLP in pretraining and in finetuning.
3.3 MAMBAMLP
We hereby introduce our newly developed MambaMLP blocks. Specifically, drawing inspiration from the self-attention block in Transformer (Dosovitskiy et al.) 2020; Vaswani al.] 2017), our MambaMLP block uses Mamba the token mixer and keep the multi-layer perceptron (MLP) the same as that in the self-attention block. Note that the configuration of the MambaMLP block varies between pretraining and fine-tuning phases to cater their different requirements. During pretraining, as illustrated in Figure (b)) the MambaMLP block contains the Mamba layer with only scan (Liu et al. 2024b). We stress that using scan is hard requirement in pretraining autoregressive modeling requires a strict causal structure, where each token is predicted based only on past tokens in the defined sequence order; multiple scans are used, this autoregressive nature will break and information leakage will happen, resulting in the learned feature representation collapsing. In finetuning (shown in Figure c)). the block is then adapted to contains the Mamba layer with the standard 4 scans as in V Vmabba (Liu et al. 2024b), which enables bi-directional modeling of global information strengthen the overall performance. Note that in these scans, the parameters are all initialized using the values of the parameters from the single scan during pretraining, The other architectural components in the pretraining and the finetuning stay the same: the block utilizes SwiGLU (Touvron et al. 2023) as the MLP layer, and the expand refers to the order of the sequence that the Mamba processes and is set to to enhance scanning efficiency. Additionally, we provide visual comparison between our MambaMLP block and the Vim block in Figure We can see that the Vim block contains Mamba layers with scans (Liu et al. 2024b) for bi-directional global information processing and has MLP layer, and the expand of each sean is set to 2. Practically, this larger expand in each scan results in higher performance but slower inference speeds.
this is because
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 [ 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
(El-Nouby et al. 2024), which is wider but less deep compared to ViT-H. balancing performance and computational efficiency. The decoder follows exactly the same design as in MAE (He et al.) 2022), which uses Transformer blocks (by default, we choose depth=4 and width=512). Note that the decoder only used for pretraining and we remove the decoder in downstream tasks. In the next section, we will extensively validate the efficacy of ARM.
4 EXPERIMENT
4.1 IMPLEMENTATION DETAILS
Pretraining. We pretrain ARM using the ImageNet- IK dataset (Deng et al. 2009). Specifically, ARM-B and ARM-L are pre-trained for 1600 epochs, and ARM-H is pre-trained for 800 epochs. for ARM-B/L/H, respectively, and a learning rate of = 1.5e-4x batchsize We adopt cosine decay schedule with warm-up for 5 epochs. We adopt the 25g AdamW (Loshchilov & Hutter 2019) optimizer with a weight decay 0.05. We use random resized cropping and random horizontal flipping. The pretraining input size set to 192 x 192.
We use a batch size of 2048/1024/51

---

Published as a conference paper at ICLR 2025
Table I: The configuration of different architecture variants.
Model
Block
Width
Depth
(aram.(M)
VIT-B
(Attention+MLP)
768
86
Vim-B
Mamba
768
24
98
ARM-B
(Mamba+MLP)
768
12
85
VIT-L
(Attention+MLP)
307
Vim-L.
Mamba
1024
48
340
ARM-L
(Mamba+MLP)
1024
24
297
VIT-H
(Attention+MIPP)
1280
32
632
Vim-H
Mamba
1536
755
ARM-H
(Mamba+MI.P)
1536
24
662
Table 2: Performance comparison on ImageNet- K. Throughputs are measured with an A5000 GPU. denotes extend the training of Vim to the large-size model, using its original GitHub repo. indicates the stride is 8. Hybrid architectures are in Gray.
Param. (M)
Throughputs (imgs/s)
Top-I (%)
Model
Token Mixer
Image Size
Base-size models
RegNetY-166
2D Conv.
224
84
870
82.9
DeiT-B
Attention
224
21
1073
81.2
Vim-Bf
Mamba
224
98
890
81.2
MambaMLP-B
Mamba
224
85
1301
81.2
VMamba-B
Mamba+2D Conv.
224
89
315
83.9
ARM-B
Mamba
224
85
1301
83.2
ARM-B
Mamba
3842
ARM-B $
Mamba
4482
85
86
85.2
Large-size models
Vim-L.
Mamba
224
340
345
81.0
MambaMIP
Mamba
224
297
445
81.4
ARM-L
Mamba
224
297
445
84.5
ARM-L
Mamba
3842
297
154
85.1
Huge-size models
Vim-Ht
Mamba
224
755
211
collapsed
ARM-H
Mamba
224
662
85.0
ARM-H
Mamba
3842
662
94
85.5
Finetuning. Following pretraining, we finetune the ARM models on the ImageNet classification task. Specifically, we finetune all models for 100 epochs with batch size of 1024, with the input size set at 224 x 224. We use the same data augmentation MAE (He al. 2022 We adopt AdamW as an optimizer, and the peak learning rate is Ir=5e-4x x batchsize with cosine decay schedule and 256
warm-up for 5 epochs. Additionally, employ the exponential moving average (EMA) (Izmailov et al. 2018) for stronger performance.
In addition to testing the ImageNet evaluation set, evaluate model robustness without finetuning on various out-of-domain ImageNet variants, including natural adversarial examples (ImageNet- A (Hendrycks et al., (021b)), semantic shifts (ImageNet-R (Hendrycks al.) 2021a)) image sketches (ImageNet-S (Wang et al. (2019), ImageNet-V2 (Recht et al. (001)). and ImageNet-Real (Beyer) al. 2020). Moreover, we finetune the pretrained model different downstream tasks including object detection and instance segmentation on COCO (Lin et al. 2014). and semantic segmentation on ADE20K (Zhou
2019
4.2 MAIN RESULTS
In Table 2 we compare ARM with convolution-based RegNet (Radosavovic al. 2020). caantion-bassed VIT, and different Mamba architectures in vision. For the sase-size model, our ARM achieves 83.2% accuracy, making substantial 2.0% improvement over its supervised MambaMLP counterpart. Additionally, we note that ARM outperforms Vim by 2.0%, and the only Mamba architecture that attains stronger performance than convolution-based RegNerY-16Gi (i.e., by 0.3%).

---

Published as a conference paper at ICLR 2025
Table 3: Robustness and Generalization evaluation out-of-domain datasets.
Method
IN-IK f
IN-V2
IN-Real f IN-Adv.f
IN IN-Ren.t
IN IN-Ske.1
Vim-S
Zhu et al.
2024
80.6
69.4
86.0
20.3
45.8
33.4
Vim-B
Zhu
2024
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 (huu et al.
2024
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: Object detection and instance segmentation use Mask R-CNN on COCO dataset
resolution of 1024.1024. Semantic segmentation uses the UperNet framework on the ADE20K dataset resolution of 512x512.
Object detection
Instance segmentation Apmisk
Semantic segmentation
Method
Apho
mloU
Vim-S Zhu el al.
2024
43.2
40.0
44.9
Vim-B
N/A
N/A
45.2
DeiT-B TTuvron et al. MambaMLP
2020
46.8
41.5
45.5
N/A
N/A
45.3
ARM
49.2
43,9
47.7
Further enhancements are observed when ARM-B is finetuned with increased input sizes of x and 448 x44 with the patchify stride of 8, where performance improves to 84.2% and 85.2%, respectively. We also report the comparison to VMamba-B, which takes a hybrid architecture: When configured with inputs of 224x224, ARM-B slightly underperforms VMamba-B by 0.7% but enjoys much faster throughput, i.e., ~4 faster; ARM-B with the inputs of 384x388 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 3.5% improvement over Vim-L, and ARM-H sets new benchmark for the largest Mamba architecture in vision to date by reaching 85.0% accuracy. Moreover, by tuning ARM at larger resolution of 384 x 384, further leveraging the model's capacity to handle long sequences at a linear complexity, we observe additional gains: 0.6% increase with ARM-L and 0.5% increase with ARM-H. Notably, ARM-H attains the best Mamba accuracy of 85.5% % on ImageNet classification.
ARM-L achieves accuracy of 84.5%, marking
4.3 ROBUSTNESS
We report the "zero-shot" robustness evaluation (i.e., no finetuning on these ImageNet variants) of Mamba architectures in 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, largest model variant, not only continues this trend but also shows an average performance superiority of 1. 1.1% over ARM-L, reaffirming the efficacy of scaling up the model size enhancing robustness.
4.4 DOWNSTREAM GENERALIZATION
Object Detection and Instance Segmentation We benchmark object detection and instance seg- mentation performance on COCO 2017 (Lin et al. 2014) by integrating our ARM as backbone into Mask R-CNN Following Swin Transformer's protocol (Liu et al., (021), we finetuning on COCO with 36 epochs, multi-scale training with shorter edges randomly resized between 480-800 pixels. Optimization employs AdamW (weight decay 0.05, batch size 16, initial learning rate le-4), with consistent augmentation strategies across all backbones to ensure comparability. As shown in Table 4 ARM achieves 49.2 Aphox object segmentation and outperforms DeiT-B and Vim-S by 2.4 and 6.0 Aph Similarly, on instance segmentation, ARM achieves 49.2 A pmask on object segementation and outperform DeiT-B and Vim-S by 2.4 and 6.0 A pmask
8

---

Published as a conference paper at ICLR 2025
Table 5: Ablation the number of predictions units.
Table 6: Ablation on prediction orders.
Num of Prediction unit
Cluster size
Top-1 (%)
Order
Direction
Top-I (%)
0 (Supervised)
N/A
81.2
Row-first
Forward
82.5
144 (iGPT)
1 x 1 (Pixel)
79.8
Row-first
Backward
82.3
Column-finn
Forward
82.5
4
96x96
82.0
Column-fine
Backward
82.4
9
x4
82.5
16
48x48
82.2
Random
Random
81.5
36
32x32
81.9
144
16x16
81.7
Semantic Segmentation We hereby report ARM's performance on ADE20K semantic segmenta- tion (Zhou et (019). We adopt our ARM-B the backbone and UperNet as the framework. The
implementation details align with established practices from (Wang al. 2021; Ren et al. 2022b) 2023c) and the mmsegmentation toolkit (Contributors 2020). We employ the Adam optimizer with a weight decay coefficient of 0.01 over 160,000 iterations. The learning rate is initialized at 6e-5, preceded by ..500-iteration warm-up phase, followed by linear decay schedule. We incorporate standard data augmentation techniques like random horizontal flipping and multi-scale resizing. Training is conducted 512 x 512 resolution crops, while inference is performed both single-scale and multi-scale resolutions to evaluate generalization. As shown in Table ARM achieves 47.7 mIoU on semantic segmentation and outperforms DeiT-B and Vim-S by 2.2 and 2.8. Under identical architecture, ARM outperform MambaMLP by 2.4 mloU.
4.5 ABLATION STUDY
Unless otherwise specified, all ablations are performed on ARM-B under 300 epochs pretraining.
Number of prediction units. Table 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 total of 144 prediction units. We note that, even with this vamilla 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 64x64, 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 cluster size of 144). We also report the comparison MambaMLP trained under the iGPT-style autoregressive pretraining - with the input image size at 144 x 144 and setting per pixel as the prediction unit, it underperforms
best setup by 2.7% (i.e., 79.8% vs. 82.5%).
Prediction Order: Tshle/ that different pre-defined orders only lead to minor performance variances. For example, both row-first and column-first forward prediction orders achieve identical performance of 82.5%; even the least favorable case, where the prediction order was row-fir first and backward, only underperforms the best case by 0.2%. Nonetheless, if we do not predefine the prediction order and pick 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 decoder depth. finding that increasing the depth up to 4 progressively enhanced performance up 82.5%; further increasing the decoder depth to 8 sees 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 width 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 (Bao et al. 2022) with cross-entropy loss. The results, presented Table g show that employing normalized pixels as the target with MSE loss yields the best performance. achieving 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%.
9

---

Published as a conference paper at ICLR 2025
Table 7: Ablation decoder designs.
Table 8: Ablation prediction targets.
Dec. Depth
Dec. Width
Top-1 (%)
Targets
Top-1 (%)
512
82.1
dVAE Bao et al 2022
82.2
512
82.4
Pixel He al 2022
81.9
512
82.5
Normed PixeI HH 2022
82.5
512
82.5
384
82.3
512
82.5
1024
82.2
Table Comparison of architecture and pretraining paradigms. FPS represents the inference speed after supervised finetuning of the model. The symbol indicates that Vim, when subjected to contrastive learning, experiences poor performance, potentially due to mode collapse.
Architecture
Pretraining paradigm Training Cost (h)
FPS (imgs/s)
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
Vim
MAE
106
923
81.4
Vim
ARM
.57
923
82.2
Pretraining paradigm. As shown in Table we evaluate different pretraining paradigms, including contrastive learning (Chen al. (022), MAE (He et al. [202)) and our ARM. Firstly, we note that all pretraining methods result in performance gains over the supervised counterpart, demonstrating
benefits 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 bascline. In contrast, our ARM achieves significant improvements of 1.3% over the supervised baseline, well 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 MAE, which is already noted for its relatively low pretraining demands.
Architecture design. Exploring further into architectural impacts, Table 9 (from the 5th row 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 side note, it is important to highlight that although Vim's performance improves with ARM pretraining, it operates ~45% slower during inference compared MambaMLP. Additionally, MambaMLP incurs only ~666 of the training cost required for pretraining Vim under the ARM framework. These points underscore the superior efficiency of our default ARM framework.
CONCLUSION
This study introduces novel autoregressive visual pretraining strategy tailored for Mamba ar- chitectures. Through this method, we have significantly improved the scalability and benchmark performance of Mamba-based models, setting new standards their operational functionality. We hope this work can lay strong foundation for future explorations and potential expansions in the usage of autoregressive pretraining strategics for Mamba architectures within the vision community.
Acknowledge This work is supported by ONR with N00014-23-1-266 We thank the AWS Cloud Credit for Research program for partially supporting the computation of this project.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. BEÎT: BERT pre-training of image transformers. In International Conference on Learning Representations, 2022.
Lucas Beyer, Olivier Hénaff, Alexander Kolesnikov. Xiaohua Zhai, and Aäron van den Oord. Are we done with imagenet? arXiv preprint arXiv:2006.0713 200..0016... 2020.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, T. J. Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners. ArXiv, abs/2005. 14165, 2020.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and lyy Sutskever. Generative pretraining from pixels. In ICML, 2020a.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICLR, 2020b.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. preprint arXiv:2003.0422) 2020c.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. ArXiv, abs/2104.02057, 2021.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 2020.
MMSegmentation Contributors. MMScgmentation:: Openmmlab semantic segmentation toolbox and benchmark.
wwwp:/ttwcccc.oooooeceeeeeee http://///.pppmmmmmmmmmam/eeeeee htt.mm///mmmmmmggttta...... 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020.
Alaacldin El-Nouby, Michal Klein, Shuangfei Zhai, Miguel Angel Bautista, Alexander Toshev, Vaishaal Shankar, Joshua M Susskind, and Armand Joulin. Scalable pre-training of large autore- gressive image models. arXiv preprint arXiv:2401.08544 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00755 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00300 2021a.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems, 34:572-585. 2021b.
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Re. How to train your hippo: State space models with generalized orthogonal basis projections. In ICLR, 2022.
Hang Guo, Jinmin Li, Tao Dai, Zhihao Ouyang, Xudong Ren, and Shu-Tao Xia. Mambair: A simple baseline for image restoration with state-space model. arXiv preprint 15545 2024.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In CVPR, 2020.
11

---

Published as a conference paper at ICLR 2025
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. CVPR, 2022.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer. The many faces of robustness: A critical analysis of out-of-distribution generalization. ICCV, 2021a.
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. CVPR, 2021b.
Vincent Tao Hu, Stefan Andreas Baumann, Ming Gui. Olga Grebenkova, Pingchuan Ma, Johannes Fischer, and Björn Ommer. Zigma: A dit-style zigzag mamba diffusion model. In European Conference on Computer Vision, pp. 148-166. Springer, 2024.
Tianyu Hua, Yonglong Tian, Sucheng Ren, Michalis Raptis, Hang Zhao, and Leonid Sigal. Self- supervision through random segments with autoregressive coding (randsac). In ICLR, 2022.
Tao Huang, Xiaohuan Pei, Shan You, Fei Wang, Chen Qian, and Chang Xu. Localmamba: Visual state space model with windowed selective sean, arXiv preprint arXiv:2403.09338, 2024.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Av- eraging weights leads to wider optima and better generalization. arXiv preprint 2018.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In ICML, 2020.
Shufan Li, Harkanwar Singh, and Aditya Grover. Mamba-nd: Selective state space modeling for multi multi-dimensional data. arXiv preprint arXiv:2402.0555 242209 2024.
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In European conference on computer vision, pp. 740-755. Springer, 2014.
Jiarun Liu, Hao Yang, Hong- Yu Zhou, Yan Xi, Lequan Yu, Yizhou Yu, Yong Liang, Guangming Shi, Shaoting Zhang, Hairong Zheng, al. Swin-umamba: Mamba-based unet with imagenet-based pretraining, arXiv preprint arXiv:2402.03322 2024a.
Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, and Yunfan Liu. Vmamba: Visual state space model. arXiv preprint 2rXiv:2401.10166 2024b.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR, 2019.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprint arXiv:2206. 13447, 2022.
OpenAl. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023. URL https://ap
semant
icschola.......r wwwpshholarrrrgoooooooo CorpCorpus ID: ID:255532855
Alan V Oppenheim, Alan S Willsky, Syed Hamid Nawab, and Jian-Jiun Ding. Signals and systems. Prentice hall Upper Saddle River, NJ, 1997.
Xiaohuan Pei, Tao Huang, and Chang Xu. Efficientvmamba: Atrous selective scan for light weight visual mamba. arXiv preprint arXiv:2403.09997, 2024.
Bo Peng, Eric Alcaide, Quentin Gregory Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman. Huanqi Cao, Xin Cheng, Michael Nguyen Chung, Leon Derezynski, et al. Rwkv: Reinventing mas for the transformer In EMNLP. 2023.
12

---

Published as a conference paper at ICLR 2025
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemystaw Kazienko, et al. Eagle and finch: Rwkv with valued states and dynamic recurrence. arXiv preprint 20Xhv:2404.05892 2024.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong.
Random feature attention. arXiv preprini arXiv:2/03.021443 2021.
Yu Qi, Fan Yang, Yousong Zhu, Yufei Liu, Liwei Wu, Rui Zhao, and Wei Li. Exploring stochastic autoregressive image modeling for visual representation. In AAAI, 2023.
Alec Radford and Karthik Narasimhan. Improving language understanding by generative pre-training. 2018.
Ilija Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, and Piotr Dollár. Designing network design spaces. In CVPR, 2020.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In ICML, 2019.
Jingjin Ren, Qingqing Zheng, Yuanyuan Zhao, Xuemiao Xu, and Chen Li. Diformer: Discrete latent transformer for video inpainting. In Proceedings of the IEFE/CVE conference on computer vision and pattern recognition, pp. 3511-3520, 2022a.
Jingjing Ren, Wenbo Li, Haoyu Chen, Renjing Pei, Bin Shao, Yong Guo, Long Peng, Fenglong Song, and Lei Zhu. Ultrapixel: Advancing ultra-high-resolution image synthesis to new peaks. arXiv preprint arXiv:2407.02158, 2024.
Sucheng Ren, Daquan Zhou, Shengfeng He, Jiashi Feng, and Xinchao Wang. Shunted self-attention via multi-scale token aggregation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10853-10862 June 2022b.
Sucheng Ren, Zeyu Wang, Hongru Zhu, Junfei Xiao, Alan Yuille, and Cihang Xie. Rejuvenating image-gpt strong visual representation learners. In Forty-first International Conference Machine Learning, 2023a.
Sucheng Ren, Fangyun Wei, Zheng Zhang, and Han Hu. Tinymim: An empirical study of distilling mim pre-trained models. In CVPR, 2023b.
Sucheng Ren, Xingyi Yang, Songhua Liu, and Xinchao Wang. Sg-former: Self-guided transformer with evolving token reallocation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 6003-6014, 2023c.
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208. 049933, 2022.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. preprint 2rXiv::012.12277 2020.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint urXiv:2307.0922 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In NeurIPS, 2017.
Feng Wang, Jiahao Wang, Sucheng Ren, Guoyizhe Wei, Jieru Mei, Wei Shao, Yuyin Zhou, Alan Yuille, and Cihang Xie. Mamba-r: Vision mamba also needs registers. arXiv preprint arXiw:2405.14858, 2024a.
Haohan Wang. songwei Ge, Zachary Lipton, and Eric P Xing. Learning robust global representations by penalizing local predictive power. In NearIPS, 2019.
13

---

Published as a conference paper at ICLR 2025
Wenhai Wang. Enze Xie, Xiang Li, Deng- Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer. A versatile backbone for dense prediction without convolutions. In Proceedings of the IEEE/CVF International Conference Computer Viston, pp. 568-578, 2021.
Ziyang Wang, Jian-Qing Zheng, Yichi Zhang, Ge Cui, and Lei Li. Mamba-unet: Unet-like pure visual mamba for medical image segmentation. arXiv preprint arXiv:2402.05079, 2024b.
Zhaohu Xing, Tian Ye, Yijun Yang, Guang Liu, and Lei Zhu. Segmamba: Long-range sequential modeling mamba for 3d medical image segmentation. arXiv preprint arXiv:2401.13560 2024.
Zhilin Yang, Zihang Dai, Yiming Yang. Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. XLNet: Generalized autoregressive pretraining for language understanding. In NeurlPS, 2019.
Shuangfei Zhai, Navdeep Jaitly. Jason Ramapuram, Dan Busbridge. Tatiana Likhomanenko, Joseph Yitan Cheng, Walter Talbott, Chen Huang. Hanlin Goh, and Joshua Susskind. Position prediction as an effective pretraining strategy. arXiv preprint arXiv:2207.0761, 2022.
Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ADE20K dataset. Int. J. Comput. Vis., 127(3): 302-321, 2019.
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint arXiv:2401.09417, 2024.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 COMPARED WITH OTHER AUTOREGRESSIVE MODELING FOR SELF-SUPERVISED
LEARNING
We compare our ARM with AIM (El-Nouby et al. 2024) which naively flatten patches in raster order (row by row), ours introduces more strategic approach by grouping spatially neighboring patches into larger clusters and using these clusters as prediction units. As shown both in our ablation and recent work (Hu et al. 2024), selecting the appropriate sequence order crucial for achieving high performance in autoregressive models e.g., all these works find better scan strategy than the raster scan (used AIM) in achieving higher performance. Moreover, we finetune AIM-0.6B on ImageNet- IK. Specifically, we used the AdamW optimizer for 50 epochs, peak learning rate of 2e-3 with cosine decay, a warm-up period of 5 epochs, and a batch size of 1024. Data augmentation and regularization techniques included drop path rate of 0.3, RandAugment, label smoothing of 0.1, Mixup at 0.8, and CutMix at 1.0.
We can observe that ARM reaches 85, 0% Top-1 accuracy and makes 0.8% improvement over AIM-0.6B and a 1.2% improvement over AIM-0.6B when causal masking is used.
Table 10: Comparison with AIM.
Model
Top-1 Accuracy (%)
AIM-0.6B (causal)
83.8
AIM-0.6B
84.2
ARM-H (ours)
85.0
A.2 STATISTICAL SIGNIFICANCE
The default seed used in experiments is 0 previous work (He al. 2022 Zhu al. 2024 Liu et al., 20246). We run additional experiments with seeds = 1024, 9718, 8, 87144, and report the results in Table Cluster deign is stable and make consistent improvements by 0.8%.
Table 11:
Statistical significance with different seeds.
Prediction Unit
Top-1 Accuracy (%)
cluster
82.5 0.1
token
81.7 + 0.2
15


Reviews:
[{"review_id": "81FCKxRMD1", "reviewer": "Reviewer_xhQF", "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.", "questions": "- As shown in MAE, pre-training has another benefit to improve downstream tasks (like object detection, segmentation). Does this conclusion hold for ARM?\n- 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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well written, including all technical details. It should be straightforward to implement.\n- 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.\n- It is good to know that cluster size plays an important role in pre-training.", "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.\n- 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?\n- Section 4.4, is it possible to improve random-order prediction by incorporating position encoding?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UmshkR7BmM", "reviewer": "Reviewer_BJtj", "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.", "questions": "See weaknesses", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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. \nThe 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.", "weaknesses": "1.My main concern is the experimental setting:   \n(1)\tWhen 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?  \n(2)\tBesides, 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. \n\n\n2.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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ud3aiTAk0r", "reviewer": "Reviewer_mXuW", "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.", "questions": "See above.  Additional, more minor questions/comments below.\n\n* 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.\"\n\n* 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\n\n* Sec 4.3:  When applied to other datasets, is autoregressive pretraining performed only on imagenet1k, the target dataset, or both?\n\n* Eq (3) says $h'(t)$ but this is the descretized system, so should this be $h_t$ on the first line ?\n\n* 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)?\n\n* typo on l.470, should say width is 512 no depth", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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).", "weaknesses": "There are several important pieces of the method that lacked detailed explanation:\n\n* 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.\n\n* 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.\n\n* 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)\n\nA 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9G3MSNO2Rt", "reviewer": "Reviewer_Xz8m", "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.", "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.\n\nDo 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.\n\n[3] Masked Autoencoders Are Scalable Vision Learners (http://arxiv.org/abs/2111.06377)", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 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.", "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.\n\nGiven 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.\n\n“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.\n\nFigure 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.\n\nTable 4 has clarity issues. It would be helpful to directly specify the cluster size instead of having it be $\\frac{cluster\\\\_size}{patch\\\\_size}$.\n\n[1] Scalable Pre-training of Large Autoregressive Image Models (http://arxiv.org/abs/2401.08541)\n\n[2] Mamba: Linear-Time Sequence Modeling with Selective State Spaces (https://arxiv.org/abs/2312.00752)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)