阿里 & 字节提出 ARM 自回归预训练，解锁 Mamba 架构在视觉任务中的强大潜力 ！

视觉社区已经开始采用最近开发的态空间模型Mamba作为一系列任务的新 Backbone 。本文展示了通过自回归预训练可以显著提升Mamba的视觉能力，这是一个之前未曾探索的方向。 在效率方面，自回归的特性能够很好地利用Mamba的单向循环结构，使得整体训练速度相比其他训练策略（如 Mask 建模）更快。 在性能方面，自回归预训练使Mamba架构在准确性上显著高于其监督训练的对应版本，更重要的是，它成功解锁了Mamba扩展到大型甚至巨大模型尺寸的潜力。 例如，通过自回归预训练，基础尺寸的Mamba达到了83.2%的ImageNet准确率，比其监督训练的对应版本高出2.0%； 作者的大型Mamba，迄今为止最大的视觉Mamba，达到了85.0%的ImageNet准确率（在使用输入进行微调后达到85.5%），显著超越了所有其他视觉Mamba变体。 代码可在https://github.com/OliverRensu/ARM获取。

1 Introduction

在自然语言处理（NLP）中，状态空间模型（SSMs）在以线性复杂度建模长序列方面显示出巨大潜力。其中，一个最近的变体，Mamba [12]，通过综合选择性扫描的最佳属性，显著超越了传统的SSMs。这一创新也促进了其在视觉社区的快速采用，被应用于各种视觉任务中。包括新型架构的设计，以及在分割[26; 48; 49]和图像合成[16]中的应用。

然而，这些先前的研究大多是在监督视觉表示学习的背景下进行的。尽管这些训练过的模型在不同的视觉任务中表现出有希望的结果，但它们通常在可迁移性方面受限，并且在扩展时遇到显著困难。例如，如图1所示，在监督条件下尝试扩展Vision Mamba（Vim）通常会导致在推向非常大尺寸时出现性能停滞甚至训练崩溃。这些问题因此激励作者探索使用Mamba架构进行自监督视觉表示学习，这种方法在帮助模型获得强大且可扩展的视觉表示方面已经取得了显著的成功。

在本文中，作者主要关注自回归预训练范式的自监督视觉表示学习，该方法从输入序列的开始到结束单向且自回归地预测下一个标记。这一关注由两个原因驱动。首先，自回归预训练已经在训练大型语言模型中确立了事实上的标准，成功应用于各种架构，包括Transformers和Mamba。最近的文献也初步确认了其在计算机视觉领域的有效性，例如，帮助Vision Transformer（ViT）发展强大且可扩展的特征表示[11; 41]。

其次，Mamba架构由于其独特设计的线性注意力本质，本质上非常适合自回归建模，该方法系统地以严格递进和单向的方式构建标记间的关系。这种配置确保每个标记只能关注其前面的标记，与自回归建模的底层原则完美对齐。

此外，这种协同实际上也导致了整体训练效率的提高。例如，在训练基础尺寸Mamba 300个周期的设置下，自回归训练只需约34小时（以8×A5000测量），与其他预训练策略相比，训练速度提高了约2倍到10倍（见表4.4中的表8）。

重要的是，为了进一步释放使用Mamba架构进行自回归视觉表示学习的潜力，作者强调形成输入序列的两个关键方法。首先，作者不是简单地以的块作为预测的基本单位，而是选择了一种更策略性的方法，通过将空间上相邻的块组合成更大的簇；从经验上，作者发现簇大小为时达到最佳性能。其次，在作者将2D图像以不同顺序映射到1D视觉句子消融研究中，作者注意到，简单的顺序（仅按行逐行和正向扫描的方式排列簇）已经是一个有效的选择。作者将这种方法称为ARM。

作者提供了大量结果显示，作者提出的ARM实现了实质上更强大的性能。如图1所示，ARM帮助作者的基础尺寸模型达到83.2%的ImageNet准确率，比监督模型高出2.0%。此外，ARM还使得首次成功训练了大型模型（ARM-H），使其成为迄今为止最大的视觉Mamba模型。具体来说，ARM-H实现了令人印象深刻的85.0%的ImageNet准确率，显著优于所有先前的Mamba变体。此外，ARM还在域外数据集上的性能有了大幅提升：ARM-B在ImageNet-A上比监督Vim-B高出4.4%，在ImageNet-R上高出2.9%，在ImageNet-S上高出3.3%。

2 Related Work

状态空间模型。状态空间模型（SSM）[13]作为一种新颖的替代Transformers的方法，用于具有线性复杂度的长距离依赖建模。线性注意力[24; 8; 36]通过一个无需softmax的注意力矩阵以线性复杂度递归地近似自注意力，这可以被视为一种退化的线性SSM。结构化状态空间序列（S4）模型[13]在基于新参数化的同时，比先前的方法计算得更有效率，同时保留了它们的理论优势。S5[43]通过增加多输入多输出（MIMO）SSM和高效的并行扫描来扩展S4。RWKV[34]是一种最近的RNN，其关键“WKV”组件与具有两个SSM的系统类似地操作。其更新版本[35]通过状态扩展和输入依赖的门控来实现更灵活的序列建模。接着，Mamba[12]提出了一种具有隐藏状态扩展的数据相关SSM层，并构建了一个通用的语言模型 Backbone ，该模型在各种尺寸下与Transformers的表现相当，并且在序列长度上具有线性扩展。本文关注于视觉中的Mamba，旨在通过自回归视觉预训练来增强它。

视觉中的Mamba。

Mamba在NLP领域的成功应用激发其在视觉应用中的采用。视觉Mamba（Vim）[52]利用由纯Mamba层组成的Vim块：每个Vim块利用前向和后向扫描来建模双向表示，并减轻Mamba中的方向敏感问题。

另外，Vmamba[27]采用视觉状态空间（VSS）块，整合了Mamba和2D卷积层，并支持类似于Swin Transformer[28]的金字塔架构：每个VSS块首先通过2D深度卷积作为标记混合器来建模2D局部信息，然后通过CrossScan模块水平垂直处理2D全局信息。Mamba-ND[25]进一步将Mamba的能力扩展到多维数据，包括图像和视频。

LocalMamba[22]将输入图像分割成几个局部窗口，并在这些窗口内进行各种方向的SSM，增强局部处理。EfficientVMamba[33]引入了一种高效的2D扫描技术，使用特征图块上的孔洞采样来减少计算需求。

与这些新设计的Mamba架构相比，作者的设计_不那么新颖_，它紧随ViT的设计，但是用Mamba模块替换了自注意力。采用这种_简单_的Mamba架构，作者主要关注的是展示自回归预训练可以增强其视觉能力。

Self-supervised visual representation learning.

自监督视觉表示学习旨在不使用标签的情况下学习强大且可迁移的表示，包括对比学习，位置预测[51]， Mask 图像建模[17, 1, 42]等。本文关注自回归预训练，这一方法在NLP中非常成功，但在计算机视觉中仍较少探索。

iGPT[4]是第一个将生成性预训练 Transformer 引入视觉领域的工作，并强调了自回归预训练作为一种通用的自监督视觉表示学习策略的潜力。SAIM[37]和RandSAC[21]进一步增强了自回归预训练，通过利用ViT架构和随机序列排列策略，实现了与MAE[17]相当的性能。D-iGPT[41]略微修改了学习目标，不仅预测下一个标记，还预测可见标记。AIM[11]展示了，通过自回归预训练，ViT可以有效地随着模型容量和数据量的增加而扩展。与这些先前关注Transformer架构的工作不同，作者提供了第一个研究，探索使用Mamba架构进行自回归视觉预训练。

3 Method

Mamba Preliminaries

Mamba架构继承了状态空间序列模型[13]，该模型在时间t通过扩展的隐藏状态对1-D函数或序列进行建模。隐藏状态通过时间演变，由参数驱动的线性常微分方程（ODEs）:

为了离散化这个连续系统中的参数，一个常见的解决方案是引入时间尺度参数，使用零阶保持（ZOH）模型[32]将连续的转换为离散的，：

通过应用这种转换，作者可以将方程式1重写为：

然后作者使用一个矩阵进行快速计算：

其中且是输入序列的长度。作者也有，，而可以被视为卷积核。请注意，这种计算结构使Mamba能够建模与自回归建模中单向、下一个词预测完全匹配的输入序列。

Autoregressive pretraining

作者首先简要回顾了NLP中的自回归预训练。然后，作者将注意力转向视觉中带有mamba的自回归预训练，包括预测单元和预测顺序的设计。最后，作者展示了模型的变体。

3.2.1 Autoregressive Pretraining in NLP

自回归预训练模型根据给定语料库 逐个建模下一个词的概率。这可以表述为：

在这里，自回归预训练计算每个词 的可能性，基于从 到 所有前导词的上下文，并最小化负对数似然：

这种策略在训练像 ChatGPT [3] 和 GPT-4 [31] 这样的大型语言模型中发挥着基础性作用。

3.2.2 Autoregressive Pretraining with Mamba in Vision

预测单元。从1D句子过渡到2D图像带来了定义合适自回归预测单元的挑战。作者从iGPT [4]中提出的原始策略开始，该策略将每个单独的像素视为预测单元，如图2(b)所示。对于一个图像 ，作者的目标是最小化损失函数：

这里 表示Mamba模型， 表示图像的第 像素。这种基于像素的方法虽然直接，但对计算的要求非常高，特别是对于高分辨率图像。因此，如iGPT原始论文[4]所示，这种限制使得必须使用低分辨率图像进行计算上可行的自回归预训练。

将图像[10]划分为非重叠区域，并将它们映射成视觉标记，可以解决这个计算挑战。例如，对于224224大小的图像，序列长度将显著地从iGPT框架中的50,176减少到使用16 16划分操作后的仅196个区块。直观地说，将预测单元从像素[4]转移到区块[10; 52; 11]，如图2(c)所示，调整自回归输入为 ：

这里 是第 个区块。此外，为了在标记 Level 封装2D空间信息，作者 Proposal 将空间上相邻的区块组合成更大的簇作为预测单元，如图2(d)所示。簇状输入 旨在通过以下方式进行优化：

在这里，每个 是由 图块分组的集群。

作者的消融研究（第4.4节，表4）显示，将集群作为预测目标显著提高了与使用单个像素或图块相比的性能。接下来，作者探讨了将这些集群序列化为一个连贯视觉句子的策略。

预测顺序。 与自然语言处理中固有的具有明确序列顺序的1D句子不同，在将2D图像投射到1D视觉句子时，作者在此探索了四种不同的预测顺序，例如，给定集群大小 ，每行有 个集群，每列有 个集群，这些集群应该如何排列。

作者在此探索了四种主要的预测顺序：

1. 先行排行 顺序地对集群逐行排序，在每个行内从第一个集群顺序处理到最后一个，如图3(a)所示。
2. 先行排行（反向） 同样逐行对集群排序，但改变了处理方向，从每个行的最后一个集群开始，向前移动到第一个，如图3(b)所示。
3. 先行列排 逐列组织集群，在每个列内从上到下顺序处理，如图3(c)所示。
4. 先行列排（反向） 类似地，逐列对集群排序，但从最底部的集群开始，向上移动，如图3(d)所示。为了考虑一种不受预定义顺序偏差的方法，作者还尝试了集群顺序的 随机 排列[50]，如图3(e)所示。

在4.4节中，详细比较了这四种预定义顺序与随机顺序的实证比较。作者的研究发现，尽管预定义顺序在性能上表现出最小的差异，但采用随机顺序会导致性能严重下降。

因此，简单有效的_先行后列顺序_（图3(a)）被采纳为作者进行自回归建模的标准排序策略。

MambaMLP

作者在此介绍作者新开发的MambaMLP模块。具体来说，作者的MambaMLP模块使用Mamba作为标记混合器，使用多层感知机（MLP）作为通道混合器，灵感来自于Transformer中的自注意力块[10; 45]。请注意，MambaMLP模块的配置在预训练和微调阶段之间会有所不同，以适应它们不同的需求。

在预训练期间，如图4(b)所示，MambaMLP模块包含只有一个扫描[27]的Mamba层，以匹配自回归预训练中的单向建模方式；而在微调阶段（如图4(c)所示），该模块被适配为包含4个扫描的Mamba层，从而实现类似于Vmamba[27]的双向全局信息建模。预训练和微调阶段的其他架构组件保持不变：该模块使用SwiGLU[44]作为MLP层，并将设置为1以提高扫描效率。

此外，作者在图4中提供了作者的MambaMLP模块与Vim模块的视觉比较。作者可以看到，Vim模块包含有2个扫描[27]的Mamba层以进行双向全局信息处理，并且没有MLP层，每个扫描的设置为2。实际上，每个扫描中较大的会导致性能更高但推理速度更慢。通过堆叠多个MambaMLP模块，并使用作者在3.2.2节中开发的自动回归策略进行训练，作者将得到的模型命名为ARM。

如表1所示，ARM被设计为在其基础和大型配置中与ViT的深度和宽度相匹配。对于巨大的模型尺寸，ARM采用了AIM-600M[11]的结构，该结构比ViT-H更宽但更浅，平衡性能和计算效率。在下一节中，作者将广泛验证ARM的有效性。

4 Experiment

Implementation Details

预训练. 作者使用ImageNet-1K数据集[9]对ARM进行预训练。具体来说，ARM-B和ARM-L预训练1600个周期，ARM-H预训练800个周期。作者分别为ARM-B/L/H使用2048/1024/512的批量大小，学习率为lr = 1.5e-4。作者采用带有5个周期 Warm up 的余弦衰减调度。作者采用AdamW[29]优化器，权重衰减为0.05。作者使用随机调整大小的裁剪和随机水平翻转。预训练输入大小设置为。

微调. 预训练之后，作者在ImageNet分类任务上对ARM模型进行微调。具体来说，作者以1024的批量大小对所有模型进行100个周期的微调，输入大小设置为。作者使用与MAE[17]相同的数据增强。作者采用AdamW作为优化器，峰值学习率为lr=5e-4，带有余弦衰减调度和5个周期的 Warm up 。此外，作者还采用指数移动平均(EMA)[23]以获得更强的性能。

表1：不同架构变体的配置。

进一步地，作者在各种域外ImageNet变体上评估模型的鲁棒性，包括自然对抗样本（ImageNet-A[20]）、语义偏移（ImageNet-R[19]）、图像草图（ImageNet-S[47]）、ImageNet-V2[40]和ImageNet-Real[2]。### 主要结果

在表2中，作者将ARM与基于卷积的RegNet[39]、基于注意力的ViT以及视觉中不同的Mamba架构进行了比较。对于基础尺寸模型，作者的ARM达到了83.2%的准确率，比其监督式MambaMLP对应模型提高了2.0%。此外，作者注意到ARM比Vim提高了2.0%，并且是唯一一个性能优于基于卷积的RegNetY-16G（_即_提高了0.3%）的Mamba架构。当ARM-B以增加的输入大小384384和448448以及patchify步长为8进行微调时，进一步观察到性能提升至84.2%和84.8%。作者还报告了与采用混合架构的V Mamba-B的比较：当配置为224224的输入时，ARM-B的性能比V Mamba-B低0.7%，但享有更快的吞吐量，_即_快约4倍；ARM-B使用384384的输入超过了Vmamba-B 0.3%，并且仍然保持更快的吞吐量，_即_440张/秒 对比 315张/秒。

接下来，作者将Mamba架构扩展到更大的模型尺寸。首先，作者观察到基于Mamba的Vim在大尺寸上性能下降，并且在超大尺寸上无法稳定训练。这一观察表明，这些先前的基于Mamba的架构面临扩展挑战。相反，ARM模型在可扩展性方面表现出色——ARM-L达到了84.5%的准确率，比Vim-L提高了3.5%，ARM-H达到了迄今为止视觉中最大Mamba架构的新基准，准确率为85.0%。此外，通过在更大的分辨率调整ARM，进一步利用模型处理长序列的能力，以线性复杂度观察到了额外的增益：ARM-L提高了0.6%，ARM-H提高了0.5%。值得注意的是，ARM-H在ImageNet分类上取得了最佳的Mamba准确率85.5%。

Robustness and Generalization

作者在表3中报告了Mamba架构的鲁棒性评估。可以观察到，ARM（自适应鲁棒建模）一致性地展现出比监督式Vim更强的鲁棒性，例如，ARM-B在多种鲁棒性基准测试中，相较于监督式Vim-B的改进幅度在1.8%到4.4%之间。更令人印象深刻的是，ARM-L（大型模型）在这些增益上更进一步，与监督式Vim-L相比，改进幅度在2.6%到7.4%之间。此外，作者最大的模型变种ARM-H不仅延续了这一趋势，而且平均性能还比ARM-L高出1.1%，这进一步证实了通过增大模型规模来提升鲁棒性的有效性。### 消融研究

本节提供了关于ARM的不同消融研究。除非另有说明，所有消融研究都是在300个周期的预训练下的ARM-B上进行的。

预测单元的数量。 表4报告了关于预测单元数量的消融研究。作者从与 Patch 大小相等的簇大小开始（即，每个簇只包含一个 Patch ），总共有144个预测单元。作者注意到，即使在这种默认设置下，自回归预训练也成功帮助MambaMLP将性能从81.2%（通过监督训练）提升到81.7%。然后，作者逐渐将多个 Patch 组合成一个簇，从而减少预测单元的总数。作者注意到，性能先是上升然后下降——当预测单元的数量设置为9，对应于64×64的簇大小时，取得了最佳性能。具体来说，这种设置相较于监督式对应模型性能提升了1.3%，相较于默认自回归预训练对应模型（即簇大小为144）性能提升了0.8%。作者还报告了与iGPT风格的自回归预训练下的MambaMLP训练对比——输入图像大小为，将每个像素设置为预测单元，这种设置的性能比作者最佳设置差了2.7%（即79.8% 对比 82.5%）。

预测顺序。 如表5所示，作者发现不同的预定义顺序只会导致性能上的微小差异。例如，行优先和列优先的前向预测顺序都实现了82.5%的相同性能；即使是预测顺序为行优先且后向的最不利情况下，也只比最佳情况差0.2%。然而，有趣的是，如果作者不预定义预测顺序，而是选择一个随机排列，性能会显著下降到81.5%。

解码器设计。作者在表6中总结了关于解码器设计的探索。首先关注解码器深度设计，发现将解码器深度增加到4可以逐步提高性能，直至82.5%；进一步将解码器深度增加到8，性能趋于饱和。在这个4层解码器设置下，作者接下来研究解码器的宽度。通过消融这三个选项{384, 512, 1024}，作者实证观察到将解码器深度设置为512会产生最佳的准确度。

预测目标。作者在此探索不同预测目标用于作者的ARM。默认情况下，作者使用每块归一化像素，并以均方误差（MSE）损失为目标。

为了进行比较，作者将其与以下两个设置进行消融：

1)未归一化像素与MSE损失；

2)dVAE [1]得出的块离散标记，以交叉熵损失为目标。表7的结果表明，使用归一化像素作为目标，配合MSE损失，可以获得最佳性能，达到82.5%的准确度。相比之下，这个配置比使用dVAE的离散标记的模型性能高出0.3%，比使用未归一化像素的模型领先0.6%。

预训练范式。如表8所示，作者评估了不同的预训练范式，包括对比学习[7]，MAE [17]和作者的ARM。首先，作者注意到所有预训练方法相对于监督对应方法的性能都有所提升，这证明了自监督视觉预训练在Mamba架构上的好处。然而，使用MAE或对比学习，相对于监督 Baseline ，性能仅适度提高了0.4%和0.2%。相比之下，作者的ARM在监督 Baseline 上实现了1.3%的显著改进，并且比对比学习和MAE获得更高的准确度。此外，在效率方面，ARM仅需要34小时的预训练，相比于已经以相对较低的预训练需求著称的MAE，将训练时间缩短了一半。

架构设计。进一步探索架构的影响，表8（从第5行到第8行）呈现了作者对Mamba架构中的另一个变体Vim是否从自回归预训练中受益的探究。结果表明是积极的，因为经过ARM训练的Vim在ImageNet上达到了82.2%的准确度，比仅监督训练的对应模型提高了1.0%。相比之下，其他预训练范式对Vim的效果并不好：当Vim受到对比学习的影响时，训练不稳定，低于监督 Baseline ；MAE在Vim上的预训练仅比监督方法略有改善，只提高了0.2%。这些结果进一步支持了ARM在视觉Mamba预训练中的有效性。

顺便提一下，尽管Vim在ARM预训练下的性能有所提升，但在推理过程中，其运行速度大约比MambaMLP慢45%。此外，在ARM框架下，预训练MambaMLP所需的训练成本仅为Vim的大约66%。这些事实突显了作者默认的ARM框架在效率上的优势，特别是与作者新开发的MambaMLP架构结合使用时。

5 Conclusion

在本研究中，作者引入了一种针对Mamba架构量身定制的全新自回归视觉预训练策略，称为ARM。这种方法通过策略性地将空间上相邻的图像块组视为预测单元，提高了预训练的效率和有效性。

通过作者的方法，作者显著提高了基于Mamba模型的可扩展性和基准性能，为其操作功能设定了新的标准。

作者希望这项工作能为未来在视觉社区中探索和扩展Mamba架构的自回归预训练策略的使用奠定坚实的基础。

参考

[1].Autoregressive Pretraining with Mamba in Vision.