请加小助理
加入AIGC技术交流群
备注公司/学校+昵称+研究方向，
广告勿扰
在本文中，作者探讨了一种策略，该策略使用专家混合（Mixture-of-Experts, MoE）来简化而非增强视觉 Transformer 。MoE层中的每个专家都是一个SwiGLU前馈网络，不采用复杂的注意力或卷积机制。
逐深度缩放被应用于逐步减少隐藏层的大小，并且分阶段增加专家的数量。使用了分组 Query 注意力。作者研究了在小数据集上进行预训练和不进行预训练的所提方法，并探讨了在这种规模下迁移学习是否有效。作者发现，即使参数量仅有0.67M，该架构也具有竞争力。
1 Introduction
在计算机视觉的实际应用中，例如边缘智能，小型且性能高效的模型仍然被优先考虑以克服计算挑战[1]。视觉 Transformer （ViTs）[2]取得了显著成果，但当模型大小和数据集较小时，其性能会大幅下降[3]。因此，有研究开始探究在中等规模数据集上表现良好的轻量级视觉 Transformer 。几乎所有这些研究都使用新型注意力块[4, 5]或将卷积机制[6]整合到其架构中。
另一方面，Tan[7]通过使用 Mask 自动编码器（MAE）[8]作为预训练策略，证明了可以使ViT有效地从小数据集中学习。在该研究中，ViT由12个 Transformer 编码层组成，每个编码层包含一个多头注意力组件和一个前馈网络。前馈网络由两个线性层组成：第一个将输出扩展到嵌入大小的
两倍
，而不是四倍，第二个将输出减小回嵌入大小。为了进一步减轻模型负担，减少前馈网络中间的扩展输出大小可能会有帮助，但过度减少可能会负面影响模型性能。
考虑到这些因素，作者设计了一种使用专家混合（MoE）[9]来简化视觉 Transformer 的架构。在作者的架构中，每个MoE层中的每个专家由一个SwiGLU[10]前馈网络形成。从设计上讲，SwiGLU在参数方面比典型的多层感知器要重。然而，通过在MoE层中设置多个专家，作者能够使SwiGLU中的隐藏大小小于嵌入大小，而不会对模型性能产生负面影响。此外，作者在每个SwiGLU中将3个线性变换中的2个在层间共享。这有助于显著降低参数数量，同时保持MoE的优势。除此之外，为了进一步减少参数数量，作者按照深度方向逐步增加MoE中的专家数量，同时线性地减小隐藏大小，借鉴了深度方向缩放[11]的思想。最后，作者使用分组 Query 注意力[12]以保持参数数量低。源代码将在不久的将来提供。
2 Method
作者提出的方法包括两个部分：
mLIT
和
mmLIT
。
mLIT
是一种基于专家混合的轻量级视觉 Transformer ，而
mmLIT
是通过 Mask 自动编码器预训练策略进行预训练和微调的
mLIT
。
mLIT
与原始的ViT[2]类似，作者首先将每个图像划分为不重叠的块。这些块通过线性变换转化为一系列嵌入向量，并通过可学习的位置嵌入进行增强。处理后的嵌入向量经过一系列基于MoE的 Transformer 编码器层。图1展示了作者 Transformer 编码器层的整体结构。
2.1.1 Grouped Query Attention (GQA)
分组 Query 注意力将 Query 头分为$G$个组，每个组共享一个键头和一个值头。例如，GQA-1，只有一个组，等同于多 Query 注意力。
2.1.2 专家混合（MoE）
MoE的损失函数。
为了鼓励所有专家具有同等的重要性，引入了两种损失函数：负载均衡损失和重要性损失。为了计算负载均衡损失，作者首先通过以下概率确定
在
mLIT
中，每个专家中的$\mathbf{V}$和$\mathbf{W_{2}}$及其对应偏置在MoE层中是共享的。
2.1.3 Depth-wise Scaling
作者的视觉 Transformer 在 Transformer 编码器层之间并没有固定的隐藏尺寸。相反，隐藏尺寸是从第一层到最后一层线性减少的。
此外，
mLIT
中的专家数量在不同阶段会增加，具体来说，在每3、4或5层之后，从3增加到5，这取决于模型的大小。更多详细信息请参见图4。
mmLIT
在
mmLIT
中，作者将在
mLIT
上应用 Mask 自动编码器。作者严格遵循文献[7]中概述的流程，该流程基于原始MAE论文[8]。Mask 比例设定为0.75。然而，与文献[7]不同的是，作者没有在解码器的输入中合并可分离的学习位置嵌入（见图5）。与文献[7]类似，作者计算自动编码器的损失为 Mask 块上的损失与未 Mask 块上的额外折减损失的之和。总损失由以下公式给出：
3 Experimental Setup
作者研究了
mLIT
和
mmLIT
在三种不同规模：S、XS和XXS上的性能表现。表1展示了每个模型的编码器细节，表2展示了解码器细节。对于
mmLIT
，作者在Cifar100 [14] 训练数据集上进行了自监督预训练，并在Cifar100、Cifar10[14]、Flowers102 [15]和Svhn [16]上进行微调。此外，作者还对上述四个数据集进行了有监督学习（不进行任何预训练）。
图3：MoE的实际实现。与输入仅是一个向量时的情况不同，当$b\cdot n$个嵌入进入层时，几乎所有的专家网络都会接收到一些由分配器发送的嵌入。
输入模型的图像大小为36 x 36，略大于这些数据集原始的32 x 32尺寸（除了Flowers102，其尺寸更大且多样）。这导致每个图像被划分为144个 Patch ，给定3 x 3的 Patch 大小。与[7]类似，在预训练和微调阶段，每张图像都添加了一个辅助的虚拟 Patch 。这个虚拟 Patch 对所有元素的值均为零，作为第145个 Patch 用于分类目的。
基于MoE的 Transformer 编码层中的所有线性层均包含偏置。然而，作者从编码器和解码器中的其他线性投影层中排除了偏置。对于所有层类型的权重和偏置初始化，作者依赖于Pytorch提供的默认方法。对于层归一化[17]的处理方式也同样如此，作者使用Pytorch的默认设置。
Pre-training
mmLIT-S
、
mmLIT-XS
和
mmLIT-XS
分别在 Cifar100 数据集上预训练了4000、6000和8000个周期。作者使用了 AdamW 优化器 [18]，并将权重衰减设置为 0.05。总周期的前5%用于 Warm up [19]。随后，作者采用了余弦衰减计划 [20] 来调整学习率，并遵循线性学习率缩放规则，基础学习率为：
微调
对于每个预训练模型，作者进行了两组微调实验。首先，作者对在 Cifar100 上预训练的模型进行了针对 Cifar100 分类任务的微调，共300个周期。其次，作者通过在 Cifar100、Cifar10、Flowers102 和 SVHN 上对每个模型进行100个周期的微调，来评估模型的迁移学习能力。表4展示了第一组微调的配置，而表5则展示了第二组微调相对于表4的配置差异。
此外，作者还评估了不同尺寸的
mLIT
