[论文笔记] ARM

信息

Title: Autoregressive Pretraining with Mamba in Vision

Author: Sucheng Ren, Xianhang Li, Haoqin Tu, Feng Wang, Fangxun Shu, Lei Zhang, Jieru Mei, Linjie Yang, Peng Wang, Heng Wang, Alan Yuille, Cihang Xie

Year: 2024

Publish: arxiv

Code: https://github.com/OliverRensu/ARM

Keyword: Mamba, 自回归, 扫描策略

背景

SSM在NLP和视觉领域由于可以用线性复杂度处理长序列关系，Mamba利用选择扫描机制进行了实现

经典SSM模型：

$$\begin{aligned} h^{\prime}(t) & =\mathbf{A} h(t)+\mathbf{B} x(t) \\ y(t) & =\mathbf{C} h(t) \end{aligned}$$

零阶保持离散化：

$$\begin{aligned} & \overline{\mathbf{A}}=\exp (\boldsymbol{\Delta} \mathbf{A}) \\ & \overline{\mathbf{B}}=(\boldsymbol{\Delta} \mathbf{A})^{-1}(\exp (\boldsymbol{\Delta} \mathbf{A})-\mathbf{I}) \cdot \boldsymbol{\Delta} \mathbf{B} \end{aligned}$$

定义快速计算：

$$\begin{aligned} \overline{\mathbf{K}} & =\left(\mathbf{C} \overline{\mathbf{B}}, \mathbf{C} \overline{\mathbf{A B}}, \ldots, \mathbf{C} \overline{\mathbf{A}}^k \overline{\mathbf{B}}, \ldots\right) \\ \mathbf{y} & =\mathbf{x} * \overline{\mathbf{K}} \end{aligned}$$

但之前主要集中于有监督学习（泛化性和scaling）

相较于其它方法，自回归

SSM：

• 没有Softmax的线性attention可以视为一种退化的SSM

• S5在S4的基础上增添了多入多出SSM和并行扫描策略

• RWKV是一种类似于像有两个SSM一样运行的RNN，融合了状态拓展和输入无关的门控机制

• Mamba是使用了隐藏状态拓展的数据无关的SSM层

自回归：

• iGPT第一个在视觉领域使用了自回归

• SAIM和RandSAC继续提升了自回归，使用了ViT和随机序列排列

• D-iGPT在自回归时不仅预测下一个token，并且预测可见token

• AIM进行scale

Vision Mamba：

• Vim：Vim层使用了堆叠的纯Mamba层，使用了双向扫描
• Vmamba：使用了VSS层，结合了Mamba层和2D卷积，使用了与类似Swin的金字塔结构，每个VSS层通过2D深度卷积-CrossScan进行处理
• Mamba-ND：增强Mamba到多维度
• LocalMamba：对图像切分，然后对于这些window使用SSM
• EfficientMamba使用空洞采样来减少计算量

Mamba背景

状态空间模型使用拓展的隐藏状态$h_t$对将1D的序列进行建模，其中，隐藏状态通过参数A，B，C随着时间变化，变化符合线性常微分方程：

$$\begin{aligned} h^{\prime}(t) & =\mathbf{A} h(t)+\mathbf{B} x(t) \\ y(t) & =\mathbf{C} h(t) \end{aligned}$$

引入一个时间缩放参数$\Delta$，使用零阶保持来对连续的A，B进行离散化：

对第一个线性常微分方程两边同乘$e^{-A t}$可以得到：

$$e^{-A t} h^{\prime}(t)=e^{-A t} A h(t)+e^{-A t} B x(t)$$

移项化简：

$$\frac{d}{d t}\left(e^{-A t} h(t)\right)=e^{-A t} B x(t)$$

同时取积分：

$$e^{-A t} h(t)-h(t_0)=\int_{t_0}^t e^{-A \tau} B x(\tau) d \tau$$

左乘$e^{A t}$并移项得$h(t)$：

$$h(t)=e^{A t} h(t_0)+\int_{t_0}^t e^{A(t-\tau)} B x(\tau) d \tau$$

进行离散化，代入$k=(k+1)T$：

$$h((k+1) T)=e^{A(k+1) T} h(kT)+\int_{kT}^{(k+1) T} e^{A((k+1) T-\tau)} B x(\tau) d \tau$$

零阶保持的假设是：

$$x(t)=x_k, \quad t \in\left[kT, (k+1)T \right),$$

那么可以令$\alpha=k T+T-\tau$对上式进行化简：

$$h((k+1) T)=e^{A T} h(k T)+\left(\int_0^T e^{A \alpha} d \alpha\right) B x(k T)$$

$$\begin{aligned} & \overline{\mathbf{A}}=\exp (\boldsymbol{\Delta} \mathbf{A}) \\ & \overline{\mathbf{B}}=(\boldsymbol{\Delta} \mathbf{A})^{-1}(\exp (\boldsymbol{\Delta} \mathbf{A})-\mathbf{I}) \cdot \boldsymbol{\Delta} \mathbf{B} \end{aligned}$$

最终变为：

$$\begin{aligned} h^{\prime}(t) & =\overline{\mathbf{A}} h_{t-1}+\overline{\mathbf{B}} x_t, \\ y_t & =\mathbf{C} h_t . \end{aligned}$$

$$h_t=\bar{A} h_{t-1}+\bar{B} x_t$$

$$h_{t-1}=\bar{A} h_{t-2}+\bar{B} x_{t-1}$$

$$h_t=\bar{A}^2 h_{t-2}+\bar{A} \bar{B} x_{t-1}+\bar{B} x_t=\sum_{k=0}^t \bar{A}^k \bar{B} x_{t-k}$$

$$y_t=C h_t=C\left(\sum_{k=0}^t \bar{A}^k \bar{B} x_{t-k}\right)=\sum_{k=0}^t C \bar{A}^k \bar{B} x_{t-k}$$

而离散卷积：

$$(x * h)[t]=\sum_{k=0}^{N-1} x[k] h[t-k]$$

定义一个矩阵$\overline{\mathbf{K}}$进行快速计算：

$$\begin{aligned} \overline{\mathbf{K}} & =\left(\mathbf{C} \overline{\mathbf{B}}, \mathbf{C} \overline{\mathbf{A B}}, \ldots, \mathbf{C} \overline{\mathbf{A}}^k \overline{\mathbf{B}}, \ldots\right) \\ \mathbf{y} & =\mathbf{x} * \overline{\mathbf{K}} \end{aligned}$$

其中，$k \in[0, L)$，$L$为输入序列长度，$\overline{\mathbf{K}} \in \mathbb{R}^L$可以被视为卷积核，这可以使得Mamba像自回归一样对输入序列建模

贡献点（对于输入特征进行改进）：

• 将自回归预训练与Mamba模型结合，探索视觉Mamba架构的自监督范式
• 利用聚类将临近的patch聚为簇作为输入，替代基于像素或patch的策略
• 对于将2D图像映射为1D输入，证明逐行扫描已经足够有效

方法

视觉中自监督

图像和文本信息存在本质上的差异，没有很强的因果关系：

• 文本的自监督研究主要集中于自回归，可以做到scaling law
• 图像的自监督研究可以分为对比学习，MAE，自回归，位置预测，Jigsaw…，对于哪种方法更有效和如何scale模型还没有统一的结论，主要难点在于：
  • 实现有效且稳定的自监督学习（模型，代理任务…）
  • 预训练方法在所有下游任务上都有效（对比偏好分类，MAE偏好生成任务…）
  • 如何简单有效地拓展模型（ViT大规模自监督训练易崩溃，Mamba拓展性问题）

自回归预训练

NLP中的自回归预训练

NLP中的自回归是对语料库（句子）$\mathcal{U}=\left\{u_1, \ldots, u_n\right\}$中下一个单词的概率的预测：

$$p(u)=\prod_{i=1}^n p\left(u_i \mid u_1, \ldots, u_{i-1}, \Theta\right)$$

通过最小化负对数概率（最大化预测出句子的概率）：

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

视觉中带Mamba的自回归预训练

预测单位：

• 基于像素：首先需要定义2D中自回归的预测单位，iGPT定义为像素

$$\begin{aligned} & \mathcal{L}=\sum_{i=1}^{n-1} l\left(f\left(\left[p_1, \ldots, p_i\right]\right), p_{i+1}\right) \\ & l(\hat{y}, y)=|\hat{y}-y|^2 \end{aligned}$$

​ 其中，$f(\cdot)$代表模型

​ 这种基于像素的方法，引入了大量计算，在iGPT中，就使用了低分辨率的图像

• 基于patch：将图像切分为不重叠的patch，并转换为视觉token

  $$\begin{aligned} & \mathcal{L}=\sum_{i=1}^{n-1} l\left(f\left(\left[P_1, \ldots, P_i\right]\right), P_{i+1}\right), \\ & l(\hat{y}, y)=|\hat{y}-y|^2 . \end{aligned}$$

• 基于patch簇的：将空间上邻近的patch组成更长的token

  $$\begin{aligned} \mathcal{L}_{\mathrm{ARM}} & =\sum_{i=1}^{n-1} l\left(f\left(\left[c_1, \ldots, c_i\right]\right), c_{i+1}\right) \\ l(\hat{y}, y) & =|\hat{y}-y|^2 \end{aligned}$$

预测顺序：

对于将2D图像变换为1D序列后如何决定预测顺序，但更多是以实验为主，缺乏相关的理论依据

MambaMLP

由Transformer Block启发而来，在预训练和微调阶段不同，将Mamba块用作于token混合，MLP用作于通道混合：

在预训练阶段，只有一次扫描来符合自回归的单方向，在微调阶段，使用4个方向的扫描，类似于VMamba中的双向扫描来获取全局信息

扫描的expand参数设定为1，Vim使用expand参数为2

expand参数越大，性能更好，但推理速度更慢

实验

实验设定

在ImageNet-1k上预训练，ARM-B和ARM-L为1600epoch，ARM-H为800epoch，batch分别为2048，1024，512，lr为$1.5 \mathrm{e}-4 \times \frac{\text { batchsize }}{256}$，余弦退火策略，warm-up为5epoch，AdamW优化器输入为$192\times192$，使用随机裁剪和flipping

在ImageNet分类上进行微调，100epoch，batch1024，输入为$224\times224$，增强与MAE相同，AdamW，lr=$5 \mathrm{e}-4 \times \frac{\text { batchsize }}{256}$，余弦退火策略，warm-up为5epoch，使用了EMA

在ImageNet的变体上进行了测试，包括自然增强（ImageNet-A），语义变化（ImageNet-R），骨架（ImageNet-S），ImageNet-V2，ImageNet-Real

对比实验

其它模型没有采用了自监督预训练

消融实验

簇中patch的消融：

预测顺序：

Decoder设计：

目标设计：

不同的代理任务：

总结

文章探索了自回归预训练与Mamba模型结合的训练，探索视觉Mamba架构的自监督范式，然而：

• 视觉自回归方法缺乏理论基础或者直觉经验
• 性能和拓展性仍比不上最新的ViT结构