华中科技 & 北航 & 清华 & 上海 AI Lab 提出 BIG.LITTLE Transformer ,用于高效视觉识别 !

科技   2024-11-23 09:03   上海  

在本文中,作者提出了大.小视觉 Transformer ,这是一种创新架构,旨在实现高效视觉识别。

该双 Transformer 系统由两个不同的模块组成:大性能模块,具有高容量和大量计算需求,以及小效率模块,设计用于速度,容量较低。

作者方法的关键创新在于其动态推理机制。在处理图像时,作者的系统确定每个 Token 的重要性,并相应地分配它们:重要 Token 由高性能大模型处理,而较不重要的 Token 由更高效的小模型处理。

这种选择性处理显著降低了计算负载,同时没有牺牲模型的整体性能,因为它确保了详细分析只保留最重要的信息。为了验证大.小视觉 Transformer 的有效性,作者在图像分类和任何东西分割任务上进行了全面的实验。

作者的结果表明,大小架构不仅保持了高精度,还实现了显著的计算节省。具体来说,作者的方法通过动态平衡性能和效率之间的权衡,实现了对大规模视觉识别任务的高效处理。

作者方法的成功强调了混合模型在优化视觉识别任务中的计算和性能的潜力,为实际应用中更实用和可扩展的先进神经网络部署铺平了道路。

1 Introduction

视觉 Transformer (ViT)自引入以来,已在计算机视觉领域产生了越来越大的影响。它在诸如图像分类(Deng等人,2009年)、图像分割和目标检测等基本任务上表现出色。此外, Transformer 架构的灵活性使ViT可以在多模态模型中充当视觉和语言信息之间的关键通道,极大地推动了它们的快速发展。另外,由于ViT的可扩展性,随着模型尺寸的增加,ViT能够有效地学习更丰富的图像表示。因此,构建大规模的ViT对于下游任务和应用来说具有很高的价值。

尽管ViT的性能令人印象深刻,但其推理速度仍然是一个显著的缺点。例如,使用ViT-Huge(具有超过600亿参数)作为核心组件的模型,如Segment Anything Model(SAM),在高性能NVIDIA A100 GPU 上可能运行在不到2 FPS(每秒帧数)。此外,具有十亿级参数的ViT模型的这一限制,在实际应用中部署基于ViT的模型将受到严重影响,因此迫切需要提高ViT模型的推理速度。

为解决此问题,近年来已开发出多种策略来提高ViT的推理速度。一些工作从模型角度出发,通过将知识蒸馏到轻量级模型(Xiong等人,2023年)或降低模型参数的精度来解决问题。相反,受到仅代表 tokens 对最终预测至关重要的发现启发,token pruning 方法出现,通过逐层减少tokens的数量来加速推理。尽管它们在图像分类任务上显示了增强的模型速度,只需要为每个图像预测一个类别标签,但直接丢弃不具有代表性的tokens可能会破坏图像tokens的空间结构并丢失上下文信息。这种不完整的信息 Stream 在执行下游感知任务(如图像分割)时可能导致模型性能下降。

因此,在保持上下文信息图像的同时,实现更高的推理速度,作者认识到并非所有的 Token 都同等重要。直观上,人类具有较大的视野范围,但每次观察世界时,作者只会关注其中一小部分。对于关注的区域,作者会更注重细节处理,同时保持对周围环境的关注。

受此观察启发,作者提出了一种名为big_LITTLE Vision Transformer (bLViT)的新系统,该系统包含在ViT架构中的少量高性能模块和大量低效能模块。在作者的设计中,每次只用性能模块更新少数重要 Token ,这确保了模型在推理过程中的性能,同时降低了计算量。对于不那么重要的区域,作者保持上下文信息,但用效率模块支付更少的计算成本,以实现高推理效率。尽管大多数图像 Token 根据其重要性从性能模块中剪裁,但效率模块确保所有 Token 按层更新,保留了图像 Token 的结构属性。一个 Token 是否由大模型处理取决于其来自预测层的 importance 分数。在整个训练过程中,作者对 Token 选择的差异性设计使预测层适当地将关键 Token 路由到性能模块,确保那些认为最重大的部分进行密集计算。

作者通过在图像分类和图像分割任务中应用作者的bLViT,使用DeiT(Touvron等,2021年)和SAM(Kirillov等,2023年)作为作者的大_LITTLE系统的基础模型,来证明其有效性。实验结果显示,计算速度和准确性之间存在权衡,突显了作者的模型在性能和效率之间的有效平衡能力。

总结一下,作者的贡献如下:

  1. 作者提出了一种大.小.LITTLE视觉Transformer(bLViT)模型,该模型在降低计算开销的同时,有效地剪枝 Token ,同时保留上下文信息,实现更好的速度-准确权衡。

  2. 作者在图像分类和图像分割任务上进行了实验,并展示了bLViT的有效性和效率。

  3. 作者进行了广泛的消融研究,以验证作者模型设计选择,并提高其性能。作者希望这些设计能为这种异构模型架构的未来发展带来帮助。

2 Related Work

视觉Transformer。视觉Transformer 在许多任务上取得了巨大成功,并展示了最先进的表现,包括图像分类(Touvron等人,2021年)、目标检测 、语义分割等。长程依赖调制使得它能够编码丰富的上下文信息,这可以通过为下游任务提供更优的图像表示来为其带来好处。因此,一系列的工作研究如何将简单的ViT适应于不同任务,以优化网络架构并提高性能,使用具有不同预训练策略的大规模数据集上的预训练模型。尽管其应用广泛且性能高,但计算负担为其推理速度和资源受限环境中的实际部署带来了挑战。对于该模型,更好的速度与准确性权衡是可取的。

计算减少。为了减少现有模型的计算量,许多工作尝试剪枝输入 Token 或合并输入 Token 。这通过识别并保留最具信息量的 Token 来实现,从而有效减少了需要处理的 Token 数量。AdaViT(Meng等人,2022)进一步尝试部分或完全删除所有 Token 的层。这种方法可以在ImageNet分类上实现良好的加速,同时性能降低非常有限。然而,很少有证据表明这种模型可以在除了图像分类之外的下游任务上工作,因为在非常早期阶段就被丢弃了许多 Token 。

利用小型模型加速。利用小型模型加速是另一种加速模型推理的方法。speculative decoding框架(Kim等人,2023年)引入了一种机制,该机制使用一个单独的大型语言模型和一个较小的模型,以提高自然语言处理中的推理速度。大-小网络(Chen等人,2018年)提出,通过使用大分支处理低分辨率输入,小分支处理高分辨率输入,来平衡图像和语音识别上的计算。混合专家也可以看作是一种通过在每次选择模型的一部分("专家")来加速推理的方法。尽管作者的方法与这些工作具有相似的精神,但作者的模型专注于开发一个单一模型,而不是两个独立的模型,并且仍然在相同的输入分辨率上工作。作者的"模型专家"具有不同的计算复杂性,这使得它更具适应性,并实现了更好的速度-准确性权衡。

也有一些研究关注模型蒸馏以及模型量化以加快计算速度。由于作者的目标是提出一个集成计算密集型和高效模块的一般模型架构,作者认为作者的模型与这些方法互补,速度可以进一步提高。

3 Big.LITTLE Vision Transformer

Overview

bLViT架构中的核心大.小模块包含两个组件:一个性能模块(P-模块)和一个效率模块(E-模块)。Token 处理流水线如图所示。在一个大.小模块中,前向函数如图1所示。作者首先考虑一组图像 Token

Token选择与路由。在双块之前,一个预测层(由线性层和softmax函数组成)估计所有图像 Token 的重要性分数,识别出图2:此模块处理一组图像 Token 作为输入。每个 Token 的重要性事先由预测层预测,允许根据重要性对 Token 进行排序。最K个 Token ,即认为最重要的 Token ,由P块处理,尽管具有更高的计算能力,但运行速度较慢。相反,整个 Token 序列通过E块处理,优先考虑效率,以较少的计算成本提供较快的处理速度。P块详细处理关键 Token ,以保持模型性能,而E块有效地更新所有 Token ,以在较低的计算成本上保持上下文信息。P块和E块的输出然后融合形成big.LITTLE模块的最终输出。

如图2所示,对于进一步处理,最重要的 Token 被确定为首要 Token 。作者使用一种基于重要性分数的顶k选择机制来选择这些 Token 。然后,这些选定的 Token 被引导到更耗计算的P块。如算法1所述,P块的注意力和FFN层只处理 Token 的一小部分,而E块则更新所有 Token 。为了使预测层能够进行反向传播,作者遵循(Raposo等人,2024)的方法,将所选 Token 的分数与P块输出相乘,公式为:

其中,score_t是当前 Token 的分数,y_t是当前 Token 的预测值,是正则化参数,用于平衡模型复杂度和泛化能力。
通过这个公式,作者可以使模型在训练过程中学习到有用的特征表示,并能够更好地泛化到新的数据。

是第 层的 Token 的重要性分数,模块可以是 P-block 中的 FFN 或注意力层, 是一个可学习的参数,初始值为 0,用于稳定训练过程。为了简洁起见,伪代码中这部分内容被省略了。

维度匹配:由于E-块和P-块具有不同的模型容量,表示的隐藏维度必然不同。为了调和这些差异,以满足效率和性能块的要求,作者在E-块的ViT块中进行了修改。具体来说,作者在FFN层的开始和结束处插入两个线性层以进行维度映射;对于注意力层,输入和输出维度直接修改以匹配主流的维度。这些操作在伪代码中的E_Attention和E_FFN中进行。

半交叉注意力在之前的 Token 削减方法中,不重要的 Token 直接被删除,导致剩余 Token 无法与注意力层中被削减的 Token 交换信息。为了解决这个问题,作者提出了一个半交叉注意力机制用于P块。具体来说,在P块的注意力层中,作者使用主要 Token 作为 Query (q),所有 Token (包括被选择的和未被选择的)作为键(k)和值(v),而不是只使用相同的 Token 作为 Query 。这样,主要 Token 仍然可以从所有 Token 中收集信息,而不仅仅是它们自己。

经过P块和E块的处理后,P块的输出与E块的输出通过全局更新的上下文进行融合。这种融合使用一个可学习的参数γ来调整 Token 对最终输出的影响,公式如下:

其中,表示当前时刻的隐藏状态,表示当前时刻的细胞状态,是一个介于0和1之间的实数,表示了 Token 对最终输出的影响。

在这里, 是一个二进制 Mask ,表示第 个 Token 是否为主 Token ()或非主 Token ()。这确保在强调最重要特征的同时,保持数据表示的整体完整性。

P-E模块的变体。在实际应用中,P模块和E模块的配置可以根据模型大小进行调整,且两者内部的尺寸都遵循ViT模块的变体。例如,作者可以将P模块和E模块的尺寸分别设置为ViT-Base模块和ViT-Tiny模块,以在保持ViT-Base性能的同时节省计算资源。在这里,作者采用1:1的堆叠比,即每一层图像 Token 都要经过一个P模块和一个E模块。在较大的模型中,如巨大的基组合模型,作者可能需要采用2:1的堆叠比或其他变体。

理论计算分析。为了降低计算需求,作者在实际中默认让性能块处理前25%最重要的 Token ,而效率块更新所有 Token ,以确保全面覆盖上下文信息。这样,作者的模型根据每个 Token 的内容自适应地分配计算资源,从而实现更好的速度与精度权衡。

作者进行了一次简单的分析,以了解作者可以在多大程度上节省计算资源:对于形状为的输入,其中是 Token 的数量,是 Token 的隐层维度,一个普通的ViT块的计算成本为(用于注意力层,用于FFN)。

性能块以的成本更新25%的 Token ,其中半交叉注意力层以的成本更新,而效率块以的隐层维度成本为(用于注意力层,用于FFN),这比将替换为在普通成本中要大,因为额外的开销来自于维度匹配)。

在处理所有 Token 时,总成本为,每层相对于理论速度的加速比为1.84倍,随着效率块变得更小,这个结果可能更高。

Training Strategy

在实际应用中,作者发现用大.小(big.LITTLE)模块无 Aware 地训练模型可能会导致次优性能,这可能是由于高剪枝比例导致的。作者在实践中发现特征蒸馏可以提高其性能。

在训练过程中,特征蒸馏被用来从预训练的vallina ViT模型中传递知识到作者的big.LITTLE ViT模型。通过将学生学习到的特征与教师学习的特征对齐,模型在应用激进剪枝时仍能保留关键信息。特征蒸馏损失的公式如下:

其中,feat_bLViT表示来自big.LITTLE模型的特征嵌入,而feat_vallinaViT表示来自预训练的教师模型的嵌入。余弦相似度函数确保了作者的模型特征表示尽可能接近教师模型。训练中使用的总损失结合了监督损失和特征蒸馏损失,权重为标量

4 Experiments

Implementation Details

在作者的实验中,作者使用了两种版本的bLViT。

在第一种变体中,作者使用ViT-Base作为P-block,ViT-Tiny作为E-block,表示为B+T。该模型包含12层,与标准的ViT-Base相同,第一层是ViT-Base层,可以看到所有 Token ,从第二层开始使用big.LITTLE模块,因此总共包含12个P-block和11个E-block。预测层在层1、4、7和10之后使用。

在第二种变体中,作者测试了一个更大的模型大小,使用ViT-Huge作为P-block,ViT-Base作为E-block,表示为H+B。这个模型遵循了标准ViT-Huge的32层架构,前9层只使用ViT-Huge,完全处理所有 Token 。从第10层开始,每两层交替使用一个big.LITTLE模块。在无E-block的层中,P-block只更新25%的 Token ,导致总共配置了32个P-block和12个E-block。在这里,预测层在层8、16和24之后使用。

对于具有窗口注意力机制的模型(如SAM,Kirillov等人,2023年),在各个窗口内进行 Token 选择,确保不同窗口内具有相同数量的 Token ,从而便于并行计算。

所有实验均在8块NVIDIA A100 GPU上进行。 初始化为。默认为。实验中应用了AdamW优化器,两个任务组的学习率均为

Baselines and Evaluation Metrics

作者将作者的方法与现有的ViT结构中的分词截短方法进行比较,例如AdaViT,ATS,A-ViT,DynamicViT,Evo-ViT,E-ViT ,高效的ViT模型,例如EfficientViT(Liu等人,2023年),MobileViT,还包括与原始ViT的比较。作者在包括图像分类和任意分割任务在内的两个任务上验证了性能。

图像分类。作者选择原始的ViT作为 Baseline 。采用Top-1准确率作为评估指标。从DeiT中采用了三个原始ViT变体。对于ATS,A-ViT,DynamicViT,Evo-ViT,E-ViT和作者方法,使用 DeiT 的预训练权重进行初始化,然后在一台8GPU的计算机上,在ImageNet-1K数据集上训练300个周期,批量大小为1024,接着进行图像分类的Top-1准确性测试。

训练细节遵循DeiT(Touvron等人,2021年)。对于AdaViT,它是由T2T-ViT(Yuan等人,2021年)初始化的,如表1中的星号所示。作者在某些方法中采用了多种设置。对于EfficientViT,作者在M5配置下的 224 和 512 分辨率下使用了相应的模型。DynamicViT使用了两个模型大小(基础和小型),而EViT使用了0.5和0.6的两个保留比例。

anything. 评估类似于SAM,其中分割是从一个前景点、一个框和多个点进行的。在这里,作者在真实 Mask 内均匀地采样随机点作为点击的点,真实框被用作 Prompt 框。作者还进行了零样本实例分割实验,遵循SAM(Kirillov等人,2023)的设置。对于 Baseline ,作者在完整的SA-1B数据集上训练了SAM的纯版本,经过2个周期。

对于Evo-ViT和E-ViT,作者将两个实验设置分为ViT-Base和ViT-Huge。在两种设置中,都使用了SAM的纯版本预训练权重进行初始化。相应地,使用了B+T和H+B的大.小配置。在训练期间,模型在占SA-1B数据集2%的10个周期上进行训练,批处理大小为8。对于测试,作者使用了LVIS数据集来评估模型的 Mask 预测性能,并在零样本实例分割中使用COCO数据集。

Image Classification

作者在ImageNet-1k分类数据集(Krizhevsky等人,2012年)上进行了实验,并报告了Top-1准确率和GFLOPs,结果如表1所示。结果表明,作者的方法取得了最佳性能。具体而言,作者的Base + Tiny bLViT将计算量降低了约50%,同时超过了ViT-B。尽管采用轻量级架构的方法相比大多数最有效的ViT方法具有显著的计算成本降低,但它们的性能受到模型容量限制。在有效ViT组中,基于ViT-Small的ATS和A-ViT的性能显著落后于作者的模型。作者的方法在基于ViT-Base的模型中取得了最佳性能和第二好的计算效率。值得注意的是,作者的模型是唯一一个超越ViT-B性能的模型,而其他类似的模型为了降低计算成本,往往会在性能上做出牺牲,如图1所示。

此外,作者可视化了在11层big.LITTLE模块中,哪些 Token 通过了P-block。如图3所示,在训练后,模型有效地选择了用于图像分类的关键区域,由高容量的P-block处理。这种能力突显了作者的bLViT架构的高效性和目标处理能力。

Segment Anything Task

在使用SA-1B数据集训练模型后,作者在两种实验上验证了它们,如表2所示。作者分别报告了三个预测 Mask 设置下的mIoU和零样本实例分割下的AP。从表中可以看出,作者的模型大大降低了计算量,体现在作者的B+T版本与ViT-B相比约减少了一半的GLOPs。此外,作者的方法在所有加速技术中表现突出,具有最高性能和效率。值得注意的是,在测试设置的三点和边界框下,作者的模型甚至超过了ViT-B和ViT-H。这种现象的潜在解释可能归因于从蒸馏损失和监督损失中获得的信号。

Ablation Study

模型设计。作者在ImageNet分类任务上进行消融研究,以验证作者的模型设计选择。除了没有进行任何 Token 剪枝的原始DeiT-Base模型外,作者还选择具有81.0 Top-1准确率和50% Token 剪枝比例的Evo-ViT作为作者的 Baseline 模型,并说明作者如何实现最终的模型设计。作者可以看到,虽然无限制地增加剪枝比例至75%并减少性能块(Early Prune)的数量(可以节省计算量,并观察到FLOPs减少)可以获得较好的性能,但性能也会严重下降。

简单地添加效率块(E-Block)可以缓解这个问题,但仍然落后于 Baseline 。然后,作者应用预测层(Predictor)和半交叉注意力(Semi-CA)来弥合这一差距。然后,作者利用预训练权重初始化,其中性能块的权重在没有任何 Token 剪枝的情况下进行预训练。作者实际上发现这可以获得更好的性能。最后,作者在训练过程中使用特征蒸馏(Feat. Dis.)来获得最佳性能,如SS 3.3中描述的那样。

蒸馏损失标量。在利用特征蒸馏损失进行模型训练时,这个损失的系数需要经验地设置,因为过大或过小的值可能会阻碍最佳性能。在表4中,可以看到2.5是一个值得采用的显著离散峰值值。

剪枝比例。在作者的模型中,当进入P-块时,一些 Token 将被丢弃,这个比例被称为剪枝比例。直观上,剪枝比例增加时,性能下降。因此,作者需要在模型性能和计算效率之间找到平衡。在表5中,作者可以大致观察到,当剪枝比例小于0.75时,随着剪枝比例的增加,性能下降趋势逐渐减弱;然而,超过这个点,性能下降变得明显加快。因此,作者实际采用的剪枝比例为0.75。

5 Conclusion

这篇论文介绍了大.小视觉 Transformer (bLViT),这是一种创新架构,旨在提高视觉识别系统的效率。

通过有策略地在高容量性能块和速度优化的高效块之间分配图像 Token ,这种架构显著降低了计算需求,同时保持了高精度。

作者的实验结果显示,bLViT不仅保留了强大的精度,还提高了计算效率,使其成为可扩展和适应性强的AI部署的实用选择。

Broader Impact

参考文献

[0]. Big.LITTLE Vision Transformer for Efficient Visual Recognition.




集智书童
书童带你领略视觉前沿之美,精选科研前沿、工业实用的知识供你我进步与学习!
 最新文章