Vision MLP | ActiveMLP: An MLP-like Architecture with Active Token Mixer

ActiveMLP: An MLP-like Architecture with Active Token Mixer

• 原始文档：https://www.yuque.com/lart/papers/bbn4i2
• 论文：https://arxiv.org/abs/2203.06108
• 代码：https://github.com/microsoft/ActiveMLP/blob/main/models/activemlp.py

与CycleMLP的思路和实现都非常类似的一篇工作。直观上来看，本文将偏移量的约束放宽，使用了可学习的形式，这一点正如我之前的文章《Pytorch中Spatial-Shift-Operation的5种实现策略》（Pytorch中Spatial-Shift-Operation的5种实现策略）中提到的那样，都可以看做是NIPS 2018的Constructing Fast Network through Deconstruction of Convolution一种特化版本。

现有方法

现有的主流视觉任务的骨干架构主要包括三类，CNN、Transformer和MLP类的架构。彼此之间的主要差异在于融合上下文信息到给定token的方式。现有方法将更加有效的token-mixing机制作为架构设计的核心。
如果按照token-mixing的角度来看：

• CNN：在一个固定形状的偏移窗口中局部混合token。
• Transformer：基于计算得到的成对注意力权重，将其他token的信息传递到query中。这依赖于token之间在嵌入空间中的相似关系。
• MLP：按照人为制定的规则以确定的形式部分token或者是全部token构建全连接操作，以确保信息交互。

所以可以说，当前的模型架构设计中token-mixing是非常重要的组件，它极大地影响着模型的有效性和效率。
由于本文主要改进MLP架构，所以作者们指出了现有MLP的问题：

• 横跨全部token的全连接操作限制了模型处理不同分辨率输入的能力，局限了下游任务的使用。
• 手动设计的确定规则也限制了空间信息交互仅发生在固定的区域，这缓解了分辨率的约束，但是也使得模型缺少了对于信息交互中有着多样特征模式的视觉内容的适应能力。

主要改进

基于现有架构的token-mixing行为，作者们有这样两点观察：

1. 空间维度上，视觉目标有着多样的形状和形变，使得固定范围内的信息交互结构低效且难以胜任。这就需要信息传播函数具有一定的自适应性。
2. 通道维度上，token的不同的语义属性分布在不同通道中。这使得不同token之间的关系建模起来很复杂。所以跨通道共享的token信息交互策略不够有效。就需要更加细粒度的信息传播方式的设计。

这篇文章作者们通过引入可学习的空间偏移参数，构建了一种内容自适应的非局部的token-mixing，即Active Token Mixer (ATM)。其主要的特点为：

• 通过自适应的方式预测将会按照通道级别集成到目标token的源token的位置。不同通道的采样偏移位置不同。
• 延续了之前许多方法的分解思路（VIP、DynaMixer等等），使用水平和垂直两个空间方向的分支并行处理。这可以降低偏移量学习的复杂性。两个轴向的分支和原始的token会被加权融合得到最终的token-mixing结果。

通过这样的结构，实现了：

• 内容自适应：不再受限于手工设计。
• 灵活性：从全局范围中，按照通道层级动态选择上下文token，可以良好的兼容具有不同尺度和形变的视觉内容。
• 有效性：ATM模块的计算复杂度只有O(HWC^2)，关于输入分辨率的线性复杂度，和灵活且不可知的感受野，使得其对于大尺寸图像任务更加友好。

计算流程

抛开符号化复杂的表达形式，整个模块的计算过程可以简单表示如下：

1. 对输入x分别送入两个轴向分支中进行处理。
2. 单个轴向分支中，输入会被送到一个独立的全连接层，无约束地预测偏移向量，其中每个元素对应一个独立的输入通道。
3. 预测的偏移量按照双线性插值的方式以浮点坐标的形式获得指定轴向上被采样的token。
4. 对原始token使用恒等连接来保留原始信息。
5. 三个分支的结果各自使用独立的全连接层进行变换得到hatxw,hatxh,hatxi。然后使用学习的权重加权求和。这里加权用的权重是对hatxw+hatxh+hatxi使用三个独立的全连接层变换后三个结果之间逐通道softmax得到。

与现有方法的差异

• 相较于传统卷积，ATM与1x1卷积具有相同的计算复杂度（实际上就是基于1x1的可变形卷积实现的）。但是ATM具有扩展感受野到全局范围的灵活性。
• 相较于MHSA，有着可学习偏移参数的ATM同样具有逐通道/头混合全局信息的能力，但却可以避免过多的计算量。
• 相较于容易让人联想到的可变形卷积，ATM主要有两点不同：
  • ATM嵌入在MLP风格的骨干网络中，可以将各通道内上下文信息更灵活的集成。
  • ATM解耦了不同轴向上的偏移参数的学习，使得其更容易被优化。

模型细节

• 模型仍然延续PreNorm的范式，并使用LN。
• 延续分层金字塔架构形式。
• 作者们发现没必要为每个ATM生成新的偏移，所以会在每个stage中每K层使用一组参数，不同模型变体使用的K也不同。另外注意，offset生成结构在每个ATM分支中每个token之间都是共享的。
• 对于每个stage，会在ATM之前引入一个位置编码生成器PEG（Conditional positional encodings for vision transformers），位置信息的引入有助于加速偏移参数的学习。

实验效果

消融实验

• 这里的baseline只包含channel-mixing MLP，即所有的offset都设置为0。
• 通过引入位置信息，COCO上得到了进一步的提升。可以看到密集任务可以受益于位置信息，而分类任务不会受其明显影响。
• 第四个模型中，作者们尝试将生成的偏移参数的数量设置为1，即所有通道共享同一组参数，从而构建了token-level的可变形算子。可以看到，虽然减少了一些参数，但是性能下降还是比较明显的。

这里进一步探究了偏移参数的分布情况。从COCO数据集中随机采样1000张图，将计算得到的不同层的中心token的水平偏移参数平均后，汇总得到直方图。可以看到，随着层的加深，感受野逐渐增大。这也验证了提出的结构可以灵活地捕获全局信息。

ImageNet1K

ADE20K

COCO