Autoregressive Pretraining with Mamba in Vision

Q1: Compare with AIM

A1: Thanks for raising this concern. We would like to clarify that, while these two works are highly related, our ARM is significantly different from AIM. While AIM naively flatten patches in a raster order (row by row), ours introduces a 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 like [1,2], selecting the appropriate sequence order is crucial for achieving high performance in autoregressive models — e.g., all these works find a better scan strategy than the raster scan (used in AIM) in achieving higher performance. Moreover, we follow your suggestion to finetune AIM-0.6B on ImageNet-1K. Specifically, we used the AdamW optimizer for 50 epochs, a 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 a drop path rate of 0.3, RandAugment, label smoothing of 0.1, Mixup at 0.8, and CutMix at 1.0. We report the experiment results below:

We can observe that ARM achieves a 0.8% improvement over AIM-0.6B and a 1.2% improvement over AIM-0.6B when causal masking is used.

We will add these discussions and results in the next version to clearly highlight the differences between our work and AIM.

Q2: all of the equations presented in this paper don’t seem to aid in clarity.

A2: Thanks for your suggestion. As Mamba is still a relatively new architecture, we believe including Equations 1 to 4 can provide useful background for readers to understand them. However, we are open to moving them into the supplementary material given the slight redundancy.

Regarding Equations 5 - 9, we believe these are important for illustrating our step-by-step adaptation of autoregressive modeling to our proposed ARM method. But we are also open to revisiting these equations to see if we can re-deliver them in a more concise format.

Q3: Confusion of “Prediction units”.

A3: Thanks for bringing up this confusion. Yes, your understanding is correct — the image patch is our basic prediction unit, and clusters are mostly just a way to define a particular scan order of patches. We will revise the corresponding equations, figures, and text descriptions to more consistently and clearly deliver this point.

Q4: Switch from 1 scan to 4 scan.

A4: Sorry for the confusion. Firstly, we would like to clarify that using one scan is a hard requirement in pretraining. This is because autoregressive modeling requires a strict causal structure, where each token is predicted based only on past tokens in the defined sequence order. If multiple scans are used, this autoregressive nature will break and information leakage will happen, resulting in the learned feature representation collapsing. In the finetuning stage, four scans are used mainly to enhance performance, following the setup in Vim. In contrast, if only one scan is used, the performance will drop by 0.4%, from 82.5% to 82.1%.

Q5: It would be helpful to directly specify the cluster size instead of having it be $\frac{cluster\\_ size}{pathch\\_size}$.

A5: Yes, in Table 4, our second column directly reports the setup of cluster size.

Q6: normalization method

A6: Thanks for the suggestion. Yes, your understanding is correct and we will add the reference for better clarity.

Q7: Do you have any experiments where you just vary the random seed to get a sense of the variance in model quality?

A7: Thanks for your suggestion. We first clarify that the default seed used in our experiments is 0. Following your suggestion, we run 3 additional experiments with seeds = {1024, 9718, 87144}, and report the results below

These experiments demonstrate that our results are stable and that the performance improvements are not due to random chance. We will include these results and a discussion on statistical significance in the next version

[1] Hu, V. T., Baumann, S. A., Gui, M., Grebenkova, O., Ma, P., Fischer, J., & Ommer, B. (2025). Zigma: A dit-style zigzag mamba diffusion model. In European Conference on Computer Vision (pp. 148-166). Springer, Cham.

[2] Yu, Q., He, J., Deng, X., Shen, X., & Chen, L. C. (2024). Randomized Autoregressive Visual Generation. arXiv preprint arXiv:2411.00776.

[3] VMamba: Visual State Space Model

Q1: 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. Does this mean Mamba is applied independently to each channel and alternates with MLP?

A1: Sorry for the confusion. We would like to clarify that, in our architecture, we employ Mamba as the token mixer, which mixes information among different tokens to capture their relationships. Or in simpler terms, by replacing self-attention with Mamba in the ViT, we obtain our MambaMLP. We will make this point clearer in the next version.

Q2: The explanation of the patch decoder.

A2: Sorry for missing this detail. We clarify that our decoder follows exactly the same design as in MAE [2], which uses Transformer blocks (by default, we choose depth=4 and width=512); Also in Table 6, we ablate different choices of decoder depth and width. Please note this decoder is dropped when finetuning to different downstream tasks.

Q3: 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?

A3: Sorry for the confusion. Firstly, we would like to clarify that using one scan is a hard requirement in pretraining. This is because autoregressive modeling requires a strict causal structure, where each token is predicted based only on past tokens in the defined sequence order. If multiple scans are used, this autoregressive nature will break and information leakage will happen, resulting in the learned feature representation collapsing. In the finetuning stage, four scans are used mainly to enhance performance, following the setup in Vim. In contrast, if only one scan is used, the performance will drop by 0.4%, from 82.5% to 82.1%.

Regarding initialization, we clarify that the parameters of all four scans in finetuning are all initiated from the one used in pretraining.

Q4: Additional task applications would strengthen the results.

A4: Thank you for your suggestion! We have fine-tuned ARM-B on the semantic segmentation task and achieved a mIoU of 47.3%, which is a notable improvement over Vim-B (45.2% mIoU) and MambaMLP-B (45.3% mIoU). We are also experimenting with object detection setups, which require more training resources, and will provide those results as soon as they are available. If we cannot include them before the discussion deadline, we will ensure they are incorporated in the final revision.

Q5: more related work.

A5: Thanks for bringing in these related works which focus on leveraging autoregressive modeling for image generation. We will discuss them in the related work section of the revision.

Q6:the scan and expand

A6: Sorry for the confusion. The scan refers to the order of the sequence that the Mamba processes; and the expand refers to expanding the hidden states with more channels. We will clarify these concepts in the revision.

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

A7: No, our models are only pretrained on ImageNet-1k, and the accuracy on Tab 3 reflects the models’ zero-shot ability on these datasets. We will make it clearer in the next version.

Q8: ht for the discretized system.

A8: Yes, you are correct. We will fix this typo in the revision.

Q9: Random order

A9: This randomness is running on the fly, i.e., the same image will be processed in different orders at different training steps. We will make it clear in the next version.

Q10: typo

A9: Thanks! We will fix it in the next version.

[1] Yu, W., Luo, M., Zhou, P., Si, C., Zhou, Y., Wang, X., ... & Yan, S. (2022). Metaformer is actually what you need for vision. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 10819-10829). [2] He, K., Chen, X., Xie, S., Li, Y., Dollár, P., & Girshick, R. (2022). Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 16000-16009).

Q1: it is not clear for the setting up of other supervised baselines. Is it fair to compare these baselines without the same training epochs?

A1: Sorry for the confusion. We can add one more column in Table 2 to report the training epochs. But we would like to point out that comparing training epochs alone may not provide a fair basis for evaluating different methods, especially when comparing self-supervised learning to supervised learning approaches. This is because, in self-supervised learning strategies like MAE, the inputs are typically masked at a very high ratio, resulting in more efficient training per epoch compared to supervised methods where the full input is processed.

To provide a more equitable comparison, we believe that reporting the actual training cost in terms of computational resources (e.g., A5000 GPU hours) offers a more direct measure. The updated Table 2 with the training cost for the base-size Mamba is shown below:

From the updated table, we observe that in the original setup (i.e., 1600 epoch pretraining + 100 epoch finetuning, which follows MAE’s setup), ARM-B achieves a higher accuracy (83.2%) but with a higher training cost (214 hours on 8XA5000). To better demonstrate the efficiency of our method, we conducted an additional experiment with a shorter pre-training schedule of 300 epochs (totaling 300 pre-training + 100 fine-tuning epochs). This setup reduces the total training cost to 71 hours on 8XA5000, which is less than that of the supervised baselines, while still achieving a superior accuracy of 82.5%. This indicates that ARM-B can attain better performance even with reduced computational resources compared to the supervised methods. We will include these details and the updated Table 2 in the next version to enhance the clarity and fairness of our comparisons.

Q2: 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. I think this paper should pretrain on larger unlabeled dataset, then finetuned on ImageNet.

A2: We appreciate your concern regarding the evaluation protocol. While early self-supervised learning methods like SimCLR [4] and MoCo [5] focused on linear evaluation or few-shot learning on downstream tasks, the evaluation protocols in recent research have evolved. Specifically, starting from works like BEiT [1], it has become a standard practice to pre-train models on ImageNet in a self-supervised manner and then fine-tune on the same ImageNet dataset using all available labels. This approach allows for a direct and fair comparison with supervised baselines under the same data constraints. Additionally, as shown in many recent works like [7] – the strong performance on ImageNet often positively correlates with generalization abilities when scaled to larger datasets.

Therefore, we believe that our evaluation protocol is consistent with current research practices and sufficient to demonstrate the effectiveness of our method. We will add these clarifications in the next version to avoid future confusion.

We first thank the reviewer for the detailed comments and the appreciation of our work. We address the concerns below:

Q1: Figure 3 is a bit confusing.

A1: Thank you for pointing this out. You are correct—the arrows within each cluster indicate the order in which patches are flattened, and the arrows across clusters represent the order of cluster-wise prediction. We will follow your suggestion to use different colors for arrows and add corresponding explanations in the captions in the next version to improve the overall clarity of Figure 3.

Q2: 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?

A2: Thanks for the suggestion! we will add one more column to mark whether the method is supervised or pre-training+supervised fine-tuning. Regarding the performance of MambaMLP at the huge size, it achieves an accuracy of 82.0%, which is 3.0% lower than our ARM-H. We will include this result in the next version.

Q3: Section 4.4, is it possible to improve random-order prediction by incorporating position encoding?

A3: Thanks for the suggestion! Actually, incorporating position encoding before the random permutation is our default setting for the random permutation scenario. Despite this, as presented in the paper, the random permutation still performs worse than our proposed row-first order. We will clarify this point more explicitly in the next version to avoid any confusion.

Q4: pre-training has another benefit to improve downstream tasks (like object detection, segmentation). Does this conclusion hold for ARM?

A4: Yes, this conclusion holds for ARM as well. For example, by finetuning ARM-B on semantic segmentation achieves 47.3 mIoU. This result is 2.0 higher than finetuning MambaMLP-B(45.3 mIoU) and 2.1 higher than finetuning Vim-B (45.2 mIoU). We are currently conducting experiments on object detection, which require more computational resources. We will share these results as soon as they are available; but if not being able to catch the discussion deadline, we will ensure they are incorporated in the final revision.

Q5: Discuss with one related work.

A5: Thank you for bringing this relevant work to our attention. Both works aim to pretrain models using autoregressive modeling techniques — while ARM pre-trains Mamba by predicting the next cluster, the method in [1] pre-trains CNN architectures by predicting surrounding masked regions. We will elaborate more on this discussion in the next version.

[1] Chen, Y., Dai, X., Chen, D., Liu, M., Yuan, L., Liu, Z., & Lin, Y. (2023). Image as First-Order Norm+ Linear Autoregression: Unveiling Mathematical Invariance. arXiv preprint arXiv:2305.16319.