#Mixture-of-Experts, MoE

ViT的轻量化还能做到什么程度?MoE方法立大功!使用专家混合(Mixture-of-Experts, MoE)来简化而非增强视觉 Transformer 。

在本文中,作者探讨了一种策略,该策略使用专家混合(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 编码器层的整体结构。

51c视觉~合集8_视觉

2.1.1 Grouped Query Attention (GQA)

分组 Query 注意力将 Query 头分为个组,每个组共享一个键头和一个值头。例如,GQA-1,只有一个组,等同于多 Query 注意力。

51c视觉~合集8_视觉_02

2.1.2 专家混合(MoE)

MoE的损失函数。为了鼓励所有专家具有同等的重要性,引入了两种损失函数:负载均衡损失和重要性损失。为了计算负载均衡损失,作者首先通过以下概率确定

在mLIT中,每个专家中的和及其对应偏置在MoE层中是共享的。

2.1.3 Depth-wise Scaling

作者的视觉 Transformer 在 Transformer 编码器层之间并没有固定的隐藏尺寸。相反,隐藏尺寸是从第一层到最后一层线性减少的。

此外,mLIT中的专家数量在不同阶段会增加,具体来说,在每3、4或5层之后,从3增加到5,这取决于模型的大小。更多详细信息请参见图4。

51c视觉~合集8_视觉_03

mmLIT

在mmLIT中,作者将在mLIT上应用 Mask 自动编码器。作者严格遵循文献[7]中概述的流程,该流程基于原始MAE论文[8]。Mask 比例设定为0.75。然而,与文献[7]不同的是,作者没有在解码器的输入中合并可分离的学习位置嵌入(见图5)。与文献[7]类似,作者计算自动编码器的损失为 Mask 块上的损失与未 Mask 块上的额外折减损失的之和。总损失由以下公式给出:

51c视觉~合集8_视觉_04

3 Experimental Setup

作者研究了mLIT和mmLIT在三种不同规模:S、XS和XXS上的性能表现。表1展示了每个模型的编码器细节,表2展示了解码器细节。对于mmLIT,作者在Cifar100 [14] 训练数据集上进行了自监督预训练,并在Cifar100、Cifar10[14]、Flowers102 [15]和Svhn [16]上进行微调。此外,作者还对上述四个数据集进行了有监督学习(不进行任何预训练)。

51c视觉~合集8_视觉_05

51c视觉~合集8_视觉_06

51c视觉~合集8_视觉_07

图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] 来调整学习率,并遵循线性学习率缩放规则,基础学习率为:

51c视觉~合集8_视觉_08

微调

对于每个预训练模型,作者进行了两组微调实验。首先,作者对在 Cifar100 上预训练的模型进行了针对 Cifar100 分类任务的微调,共300个周期。其次,作者通过在 Cifar100、Cifar10、Flowers102 和 SVHN 上对每个模型进行100个周期的微调,来评估模型的迁移学习能力。表4展示了第一组微调的配置,而表5则展示了第二组微调相对于表4的配置差异。

51c视觉~合集8_视觉_09

51c视觉~合集8_视觉_10

此外,作者还评估了不同尺寸的 mLIT 模型(仅使用监督学习从头开始训练)的性能。表6展示了在四个数据集上对 mLIT 进行监督学习时的配置差异。作者遵循线性学习率缩放规则(方程16)进行所有微调和监督学习。所有预训练和微调都使用混合精度(torch.float16)进行。

51c视觉~合集8_视觉_11

4 Result

本研究的结果部分旨在详尽地展示实验数据和模拟结果的详细信息。以下各节将分别阐述所采用的方法在不同测试集上的表现,并对所得结果进行分析讨论。

On Pre-Training

图6展示了不同规模的mmLIT预训练损失。当训练损失在x轴(时期)和y轴(损失值)上以对数尺度绘制时,损失下降表现出近似对数线性关系。这表明损失在每个时期呈指数级减少。

51c视觉~合集8_视觉_12

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%。

51c视觉~合集8_视觉_13

表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的图像大小上不具备竞争力并不意外。

51c视觉~合集8_视觉_14

5 讨论

表8的结果表明,在更小的规模下,ViT可以通过简化的MoE进行更好的学习。在作者当前的 Proposal 中,在MoE层中共享。这种安排背后的主要考虑是上应用的非线性sliu。尽管如此,作者还探索了mmLIT-S上的另外两种可能的共享安排(见表9),结果表明保留会带来轻微的优势。

51c视觉~合集8_视觉_15

在作者的设计中,作者逐步减少了隐藏层的尺寸并在几个阶段增加了专家的数量。这种安排受到卷积神经网络[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中缺少类似的图像。作者相信,在稍大一些且更多样化的数据集上进行预训练的模型能够在这些微小尺寸上的各种简单任务上具有竞争力的表现。

在文献中,通常会将卷积层[29, 23]集成进来,或者采用卷积“教师”模型[24]来赋予归纳偏置。似乎使用简化的MoE(多专家系统)可以帮助缓解归纳偏置的不足。然而,正如[7, 21]所示,通过 Mask 自动编码器的设置,可以克服归纳偏置的问题。

6 Conclusion

在本文中,作者展示了精简的MoE架构在创建极为轻量级的视觉 Transformer 方面的潜力。

通过在MoE层间共享参数并采用深度可分离缩放,作者实现了模型复杂度与性能之间的平衡。

此外,作者的研究结果表明,在稍大且更多样化的数据集上进行预训练可以增强模型在各种任务中的通用性和有效性。

精简的MoE方法在与 Mask 自动编码器设置结合使用时,似乎在缓解归纳偏置不足的问题上具有很大前景。



#TAPTR

让跟踪最细粒度的视觉元素像检测一样简单

本文介绍了两篇关于视觉元素跟踪的系列工作,均收录于ECCV 2024。本文提出了一种新的网络架构TAPTR,将跟踪任务视为检测任务,通过将每个跟踪点视为一个visual prompt,并使用point-query和cost-volume等技术进行跟踪。

本文介绍一下我们最近的两篇系列工作:

TAPTR: Tracking Any Point with Transformers as Detection(收录ECCV 2024)

TAPTRv2: Attention-based Position Update Improves Tracking Any Point

主页: https://taptr.github.io

代码:github.com/IDEA-Research/TAPTR

TAPTR: https://arxiv.org/abs/2403.13042

TAPTRv2: https://arxiv.org/abs/2407.16291

TAPTR(收录ECCV 2024)

51c视觉~合集8_视觉_16

图1. TAPTR的网络架构图。

如图1所示,我们首先提出了TAPTR这个简单高效的网络架构。与前人的工作不同,我们不再将每个被追踪点的各种特征 (包括语义特征、光流、光流的映射、可见性、cost-volume等) 混为一谈,然后一股脑地送入到网络中进行处理,使得对点追踪的建模黑盒化。在TAPTR中,我们首次提出将每个要被追踪的点都清晰地看作是point-level的visual prompt,而TAP任务则可以被看作是从视频的每一帧中将这个visual-prompt所指代的点给检测出来。这使得我们对追踪点的建模十分清晰,同时基于前人对检测任务的详细研究也使得整个架构简单且高效。TAPTR可以被大致地分为三个阶段,接下来我会按次序进行讲解。

第一个阶段为准备阶段。首先是对video中每一个frame的特征图进行提取。这里我们使用的是一个ResNet以及一个Transformer Encoder。在这之后,我们需要为每一个被追踪点获取一个能够准确描述该点的特征。为了方便描述,不失一般性的,我们在这里假设我们只需要追踪一个点。在TAPTR中,我们会在被追踪点开始被追踪的位置l_e进行双线性插值从而采样出描述该点的特征f_e。为了将被追踪点从每一帧中检测出来,我们需要在每一帧 都构建属于该被追踪点且负责将该点检测出来的point-query。我们将这些point-query初始化为 (l_e, f_e).

由于tracking-any-point任务与optical flow estimation任务的相似性,我们将在optical flow以及前人工作中被验证十分有效的cost-volume (有些工作也称其为correlation map) 也放入到我们的framework中。为了构建每一个点的cost-volume,我们将每一个被追踪点的初始特征f_e与每一帧的特征图进行点乘。(由此可见,随着被追踪点的数量以及输入图像的增大,这里的cost将会急速增高,这也是我们v2版本致力于解决的问题之一)

第二个阶段是通过我们的point-decoder将被追踪点在各个帧中检测出来。在point-decoder中,主要包含四个部分。由于cost-volume提供了最基础的被追踪点与图像的相似度关系,为被追踪点的检测提供了最基本的感知,我们首先会通过cost-volume aggregation模块将point query和cost-volume进行交互。接着,我们也会将point query与视频图像的深度特征图之间通过cross-attention进行交互从而补充更加详细的几何结构信息。接着,在同一帧内的所有point query之间会互相交互,从而扩大各个point query的感受野,提供更加丰富的上下文信息。此外,我们会将属于同一个被追踪点的在各个frame下的point query在temporal-attention中互相进行交互,从而补充时序信息。每一个point query最终都会经过content updater以及position updater来更新它的语义特征、位置以及可见性。

第三阶段:由于一个视频通常都很长,为了节省计算资源,同时也是为了适配不同长度的视频,在TAPTR中我们使用滑动窗口的策略。在每一次forward的时候,我们仅仅需要将所有点在某一个window内的视频帧中检测出来即可。检测结果将会通过window post-processing模块更新到完整的trajectory中。当滑动窗口滑动过整个视频之后,我们的追踪任务也就完成了。

51c视觉~合集8_视觉_17

表1. TAPTR与其他方法的性能对比。

我们在TAP-Vid benchmark上对TAPTR的性能进行了测试,结果如表1所示。可以看到,TAPTR在各个数据集中都达到了最优的性能,以及最快的速度。

51c视觉~合集8_视觉_18

表2. TAPTR的消融实验。

由于TAPTR是TAP任务的baseline工作,如表2所示,我们提供了丰富的ablation实验,验证了TAPTR中各个模块的有效性来作为后续工作的参考,从而方便后续的研究工作。

TAPTRv2

正如我们上文对TAPTR的讲述中所提到的那样。受optical flow的影响,现有的方法,包括TAPTR,都十分依赖cost-volume。这样的依赖不仅仅会导致额外的计算量(尤其是在被追踪点数以及视频的分辨率较大时),还会污染我们的point-query的语义特征。

但是,cost-volume却对TAPTR的性能有着很大的影响。因此,在TAPTRv2中,我们主要探究了:1. cost-volume为什么work?2. 能否有更加轻量化地方式来得到cost-volume或者它的等价替代者?3. cost-volume或者它的代替者要怎么被使用才能不污染我们的point-query?

51c视觉~合集8_视觉_19

表3. cost-volume对domain gap的缓解作用的实验证明。

如表3所示,我们发现cost-volume真正起作用的原因可能是在于它对domain gap的缓解作用。

51c视觉~合集8_视觉_20

图2. 在APU帮助下的TAPTRv2的decoder层。

我们认为,cost-volume与注意力机制中的注意力有着本质的相似性。因此我们可以用attention value来作为cost-volume的替代者。同时,为了防止其污染point-query,我们不再简单地将其融入到point-query的语义特征中,而是将其直接转化成一个对point-query的位置的一次更新。为此,如图2所示,我们提出了Attention-based Position Update (APU) 。

51c视觉~合集8_视觉_21

图3. TAPTRv2的架构图。

51c视觉~合集8_视觉_22

图4. TAPTRv2、TAPTR以及前人的工作的方法的对比。

如图3以及图4所示。在APU的帮助下,我们不仅仅能够在不污染point-query的情况下保留原本的domain gap的缓解能力,同时,我们也无需再额外计算cost-volume,也无需在decoder中加入额外的cost-volume aggregation操作了。整个模型十分的简洁高效。

51c视觉~合集8_视觉_23

表4. TAPTRv2与其他方法在性能上的对比。

此外,如表4所示,在保持了对domain gap的缓解能力并解决了point-query被污染的问题之后,TAPTRv2相比于TAPTR在各项指标上有了进一步的提升。