在本文中，作者探讨了一种策略，该策略使用专家混合（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 头分为
个组，每个组共享一个键头和一个值头。例如，GQA-1，只有一个组，等同于多 Query 注意力。
2.1.2 专家混合（MoE）
MoE的损失函数。为了鼓励所有专家具有同等的重要性，引入了两种损失函数：负载均衡损失和重要性损失。为了计算负载均衡损失，作者首先通过以下概率确定
在mLIT中，每个专家中的
和
及其对应偏置在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的实际实现。与输入仅是一个向量时的情况不同，当
个嵌入进入层时，几乎所有的专家网络都会接收到一些由分配器发送的嵌入。
输入模型的图像大小为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 模型（仅使用监督学习从头开始训练）的性能。表6展示了在四个数据集上对 mLIT 进行监督学习时的配置差异。作者遵循线性学习率缩放规则（方程16）进行所有微调和监督学习。所有预训练和微调都使用混合精度（torch.float16）进行。
4 Result
本研究的结果部分旨在详尽地展示实验数据和模拟结果的详细信息。以下各节将分别阐述所采用的方法在不同测试集上的表现，并对所得结果进行分析讨论。
On Pre-Training
图6展示了不同规模的mmLIT预训练损失。当训练损失在x轴（时期）和y轴（损失值）上以对数尺度绘制时，损失下降表现出近似对数线性关系。这表明损失在每个时期呈指数级减少。
mnLIT-S显示出最低的损失，在第400个时期达到0.020404。mmLIT-XS和mmLIT-XXS在第6000和第8000个时期分别实现了0.021051和0.021514的损失。mmLIT-S的预训练稳定性显著，其训练损失的变化最小。相比之下，mmLIT-XS和mmLIT-XXS的训练损失显示出更大的波动性和更频繁、更明显的高峰，特别是在第1000个时期之后尤为明显。
On Fine-Tuning
表7比较了在300个训练或微调周期后，具有相似大小的模型的性能。mmLIT-S的准确率比Mac-VIT-C100高近1%，尽管其参数仅有Mac-VIT-C100的三分之二。mmLIT-XS的大小是Mac-VIT-C100的三分之一，其性能仅低1.5%。更有趣的是，仅有Mac-VIT-C10018%参数的mmLIT-XXS依然具有竞争力，准确率仅低3.3%。此外，mmLIT-XXS甚至略微超过了ResNet56，尽管ResNet56的参数多出18%。
表8展示了仅用Cifar100数据集进行预训练的模型的迁移学习能力，与文献[21]报告的结果进行对比。可以看出，即使参数规模为0.67M的mmLIT也具有竞争力，此时VIT-t+SSAT的规模几乎是mmLIT-XXS的9倍，而其余模型至少是mmLIT-XXS的30倍以上。此外，除了在架构中包含卷积的CVT-13之外，mLIT在Cifar100、Cifar10和Svhn数据集上，即使在小规模下也始终优于普通的VIT和Swin-T。Flowers102是一个细粒度分类数据集，因此mLIT在36 x 36的图像大小上不具备竞争力并不意外。
5 讨论
表8的结果表明，在更小的规模下，ViT可以通过简化的MoE进行更好的学习。在作者当前的 Proposal 中，在MoE层中共享。这种安排背后的主要考虑是上应用的非线性sliu。尽管如此，作者还探索了mmLIT-S上的另外两种可能的共享安排（见表9），结果表明保留会带来轻微的优势。
在作者的设计中，作者逐步减少了隐藏层的尺寸并在几个阶段增加了专家的数量。这种安排受到卷积神经网络[27]的启发，在卷积神经网络中，特征图的大小会逐层减少，而随着网络的加深，特征图的数量会增加。作者发现，如果早期层的减少不显著，且在最后几层不超过75%，MoE的这种安排下隐藏层尺寸的减少对性能的影响很小。未来的实验需要进一步确认这一点。
当作者在四个数据集上进行微调时，作者首先将用于Cifar100的微调设置应用于其他三个数据集。但是，早期阶段的过拟合成为一个一致性问题，因此作者必须修改设置以适应每个数据集。当作者尝试从头开始使用mLIT进行有监督学习时，也遇到了类似的问题。
因此，为了清晰起见，作者在表4和表5中报告了设置。值得注意的是，作者没有对配置进行彻底搜索；
作者只是调整了超参数以避免训练过程中的过拟合。此外，与[21]不同，作者在所有微调中并未使用先进的增强技术。
在mmLIT的预训练期间，作者针对不同尺寸的模型使用了不同数量的训练轮次。从图6可以看出，从损失的角度来看，应该为mmLIT-XS和mmLIT-XXS分配更多的轮次。然而， Mask 自编码器的一个缺点是模型的损失与其在后续下游任务中的性能之间的关系不明确。根据作者的经验，在后期阶段再进行几百轮的预训练，有时会导致分类性能相似或更差。这可能是在数据集较小时的特定问题。尽管如此，作者相信对于这类数据集，进行数千轮的进一步预训练应该能带来更好的模型。
在本研究中，作者发现仅用50,000张图像进行预训练的mLIT即使在最小尺寸下也能作为某种基础模型[28]。在Cifar10和Flowers102上，通过Cifar100进行预训练的模型至少提高了10%。在Cifar10上，mmLIT-XXS在不到40轮次的情况下可以达到90%的准确率。在Svhn上，改进是适度的。这部分是因为mLIT即使没有预训练也表现良好。也可能是由于Cifar100中缺少类似的图像。作者相信，在稍大一些且更多样化的数据集上进行预训练的模型能够在这些微小尺寸上的各种简单任务上具有竞争力的表现。
