MoR-ViT打破ViT参数冗余魔咒 | token级动态递归对决DynamicViT,70%裁参2.5倍提速

机器学习算法人工智能与算法

picture.image

精简阅读版本

本文主要解决了什么问题

    1. 标准视觉Transformer(ViT)架构存在参数冗余严重和计算成本高的问题,限制了其实际部署。
    1. 现有高效ViT方法(如静态模型压缩或token级Sparse化)对所有token应用固定的计算深度,无法根据token的语义复杂度或重要性进行差异化处理。
    1. 现有方法无法实现灵活且与输入相关的计算资源分配,导致在琐碎区域进行不必要的计算。

本文的核心创新是什么

    1. 提出了MoR-ViT,一种新颖的视觉Transformer框架,首次将受Mixture-of-Recursions(MoR)范式启发的token级动态递归机制融入其中。
    1. 设计了轻量级路由器,使每个token能够自适应地确定其处理深度,实现基于token语义的计算自适应分配。
    1. 实现了两种路由方案:专家选择和Token选择,前者在每次递归中采用分层top-k Token选择,后者在开始时为每个Token分配固定的递归步骤。
    1. 引入了参数共享机制,仅缓存被选择的token的键值(KV)对,并参与后续的自注意力计算,提高了计算效率。

结果相较于以前的方法有哪些提升

    1. 在ImageNet-1K和迁移基准上的实验表明,MoR-ViT以高达70%的参数减少和2.5倍的推理加速实现了最先进的精度。
    1. 在同等条件下超越了DynamicViT和TinyViT等领先的效率ViT Baseline。
    1. 在多个下游分类数据集(CIFAR-10、CIFAR-100和Flowers-102)上表现出强大的泛化能力和迁移能力。
    1. 在不依赖额外领域特定资源(如大规模蒸馏或大量外部预训练数据)的情况下实现了这些改进。

局限性总结

    1. 当路由器退化时(几乎所有token都被分配了最小的递归深度),动态机制有效地简化为统一的静态计算调度,潜在效率提升将丧失,类似于固定深度模型。
    1. 禁用token级的动态递归会导致性能下降(0.7个百分点),表明该机制对整体性能至关重要。
    1. 关闭参数共享会使模型参数量大幅增加(达到8600万参数),推理速度降低约50%,尽管精度提升有限(仅比基础ViT提升0.9个百分点)。
深入阅读版本

导读

视觉Transformer(ViT)在图像识别领域取得了显著成功,但标准ViT架构因参数冗余严重和计算成本高而限制了其实际部署。尽管近期高效ViT的研究主要集中于静态模型压缩或token级Sparse化,但它们仍受限于所有token的固定计算深度。在这项工作中,作者提出了MoR-ViT,一种新颖的视觉Transformer框架,首次将受Mixture-of-Recursions(MoR)范式启发的token级动态递归机制融入其中。该方法使每个token能够自适应地确定其处理深度,从而实现灵活且与输入相关的计算资源分配。在ImageNet-1K和迁移基准上的大量实验表明,MoR-ViT不仅以高达70%的参数减少和2.5倍的推理加速实现了最先进的精度,还在同等条件下超越了DynamicViT和TinyViT等领先的效率ViT Baseline 。这些结果确立了动态递归作为高效视觉Transformer的有效策略,并为现实场景中可扩展和可部署的深度学习模型开辟了新途径。

1 引言

视觉Transformer(ViT)已成为视觉识别领域的一种强大范式,在某些情况下甚至超越了传统的卷积神经网络(CNN)如ResNet [1]。开创性的视觉Transformer(ViT)[2]将图像建模为一系列块 Token ,并应用标准的Transformer编码器,在多个基准测试中取得了优异的性能。然而,普通的ViT具有显著的参数冗余和高计算成本,限制了其效率和实际应用。

为解决这些局限性,近年来该领域见证了高效视觉Transformer架构研究的激增。DynamicViT [3] 引入了token Level 的动态Sparse化,通过在每一层自适应地丢弃无信息token来减少计算量。Adaptive Token Sampling (ATS) [4]、EViT [5] 以及相关的动态token剪枝技术进一步深化了这一方向,实现了输入自适应计算,并在降低FLOPs的同时保持高精度。EfficientFormer [6] 和 PoolFormer (MetaFormer) [7] 通过优化网络结构和算子来提升延迟和参数效率,证明基于Transformer的模型可以实现移动端 Level 的速度和紧凑性。MobileViT [8] 结合卷积和Transformer模块,提供了轻量级、移动友好的 Backbone 网络。TinyViT [9] 的近期研究利用大规模蒸馏和优化的架构,提升了小ViT的精度,在紧凑设置中可与甚至超越更大的模型相媲美或表现更优。其他创新如TokenLearner [10] 提出了动态token聚合模块,该模块通过少量潜在token学习总结全局信息,而SparseFormer [11] 则探索了使用有限可学习 Query 的极Sparse视觉表示。

尽管取得了显著进展,几乎所有现有的高效视觉Transformer方法都受到一个关键约束的限制:它们对所有token应用固定的计算深度和相同的处理方式,而不管其语义复杂度或重要性如何。因此,即使经过高度优化的视觉Transformer也可能未能以token自适应的递归方式分配计算,从而无法充分利用潜在的节省。

为解决这一差距,作者提出了MoR-ViT,一种全新的视觉Transformer架构,首次集成了受Mixture-of-Recursions(MoR)范式启发的token级动态递归机制。在MoR-ViT中,每个token通过一个轻量级路由器进行网络传输,其递归深度由输入依赖的灵活机制决定,从而实现基于token语义的计算自适应分配。作者在ImageNet-1K和迁移学习基准测试上的实验表明,MoR-ViT在精度上达到了与当前最先进的高效ViTs(包括DynamicViT、TinyViT等)相当或更优的水平,同时实现了高达70%的参数减少和2.5倍的推理加速。这些结果突显了动态递归作为一种可扩展、资源高效视觉建模的新而有效策略,为未来的研究和部署开辟了广阔的前景。

值得注意的是,MoR-ViT在不依赖额外领域特定资源(如大规模蒸馏或大量外部预训练数据)的情况下实现了这些改进,并且在多个情况下,在关键评估指标上甚至超越了 Baseline 性能。

2 相关工作

2.1 视觉Transformer的全面综述

近期,多项综合综述已对视觉 Transformer (ViTs)及其扩展的快速演进进行了系统梳理[12, 13]。这些研究系统地分类了视觉 Transformer 架构,分析了自注意力机制变体、混合结构,并总结了在分类、检测和生成建模中的应用。同时,它们还探讨了效率、可扩展性和部署方面的开放性挑战,强调了在资源受限和自适应计算环境下开展进一步研究的必要性。

2.2 TokenSparse化和动态Token剪枝方法

一项重要的研究方向集中于通过自适应选择或剪枝输入token来减少ViTs中的冗余。DynamicViT [3] 引入了一种动态tokenSparse化机制,允许模型在每一层丢弃无信息的token。Adaptive Token Sampling (ATS) [4] 和 EViT [5] 通过学习根据内容选择显著patch来进一步优化token剪枝。TokenLearner [10] 和 SparseFormer [11] 探索了token聚合和Sparse潜在token表示,实现了高效的全局信息建模。这些方法在显著节省计算资源的同时保持了高精度,但仍然对所有token应用统一的处理深度。

2.3 高效架构设计与模型压缩

另一个主要方向是高效Transformer架构的设计和参数压缩。EfficientFormer [6] 和 PoolFormer(MetaFormer)[7] 提出了算子级和结构级的优化,大幅降低了推理延迟和模型尺寸。MobileViT [8] 将卷积先验与Transformer模块相结合,为移动端友好的视觉 Backbone 。TinyViT [9] 采用大规模蒸馏和紧凑设计,以最少的参数实现最先进的精度。尽管这些工作证明了ViT可以变得轻量化和快速,但它们仍然对所有token使用固定的计算调度。

2.4 现有方法的局限性及本工作的动机

尽管取得了显著进展,几乎所有现有的高效视觉Transformer(ViT)方法都存在一个关键局限性:它们对所有token应用固定的计算深度和相同的处理方式,而不管它们的个体语义重要性或复杂性如何。这种统一性限制了资源分配的灵活性,并可能导致在琐碎区域进行不必要的计算。作者的工作通过引入MoR-ViT来弥补这一差距,MoR-ViT能够实现token级的动态递归和在Transformer内部的适应性深度分配,从而实现更资源高效的视觉建模。

3 方法

3.1 视觉Transformer(ViT)概述

视觉Transformer(ViT)[2]的架构和主要计算流程如图1所示。ViT直接将Transformer框架应用于图像识别。

picture.image

给定输入图像

,将其分割为

个不重叠的块,每个块的大小为

。每个块

被展平为一个向量,并通过可学习的线性投影 E E R(P23)xD 映射到一个

维的token嵌入。

Patch 嵌入序列表示为

。在序列前添加一个类别 Token

,并添加可学习的位置嵌入

由此产生的序列

被输入到一个堆叠了

个标准Transformer编码器块的序列中。

3.2 混合递归机制 (MoR)

混合递归(MoR)框架基于标准Transformer流程,通过引入动态路由机制,自适应地为每个token分配递归计算。输入嵌入、patch划分和位置编码与vanilla ViT和Transformer架构相同。

给定输入token

(例如,具有位置编码的图像块嵌入),MoR采用轻量级路由器来确定分配给每个token的递归深度。在每个递归步骤

,路由器计算token

的路径选择得分。

其中

是在递归步骤

时 Token

的隐藏状态

是路由参数,而

是一个激活函数(例如 sigmoid 或 tanh)。

计算步骤

中所有路由分数的

百分位数阈值

,选择满足

的token继续进入下一递归步骤,而其他token则退出:

这里

是共享的 Transformer 模块,

是其参数。用于 token 选择的 Mask 矩阵在每个递归步骤中都会更新。为了实现效率,仅缓存被选择的 token 的键值(KV)对,并参与后续的自注意力计算,遵循 [14] 中的做法。

MoR支持专家选择和 Token 选择两种路由方案,前者在每次递归中采用分层top-k Token 选择,后者在开始时为每个 Token 分配固定的递归步骤。

目标函数由主要任务损失(例如分类中的交叉熵)和可能的辅助损失组成,以鼓励稳定的路由和平衡的计算。

符号摘要:

: Token

在递归

时的路由得分

..

-百分位数阈值用于在步骤

进行路由

: 递归

时的隐藏状态

: 递归块的共享参数

: Transformer块函数

3.3 MoR-ViT 架构

所提出的MoR-ViT架构的整体工作流程,该架构集成了ViT风格的图像块嵌入与混合递归(MoR)机制,以实现图像识别的逐 Token 自适应计算,如图3所示。

picture.image

给定输入图像

,图像首先被划分为

个大小为

的非重叠块,其中

。每个块被展平并映射到一个

维的token嵌入:

一个可学习的类别 Token

被添加在前面,并添加位置嵌入

序列

作为 MoR 模块的输入 Token 。

在MoR模块中,每个token都经过一个共享变换。在每一步递归步骤

中,一个轻量级路由器为每个token计算一个门控分数:

其中

表示在步骤

时 Token

的隐藏状态,

表示Sigmoid激活函数。

为了动态选择用于进一步递归的 Token ,通过在自适应的

-百分位数处对门控分数进行阈值处理来构建一个二进制 Mask :

Extra open brace or missing close brace

具有

的 Token 在步骤

时退出递归,并保留其输出,而具有

的 Token 在下一步中递归处理。这种动态过程将更深层次的递归分配给更模糊的 Token ,并允许更简单的 Token 提前退出。

为了提高实现效率,仅缓存所选token的键值对以供后续自注意力计算使用,遵循[14]。

在所有递归步骤结束后,所有 Token 的输出被聚合,最终的[CLS] Token 通过一个线性分类器(例如MLP)生成类别logits。训练目标是定义如下:

其中

是标准任务损失(例如交叉熵),

是路由正则化项,

是权衡系数。对于图像分类任务,如CIFAR-10/100和Oxford Flowers-102,该协议遵循ViT/DeiT [2,15]。

动态路由的可解释性。为了更深入地理解作者的动态路由机制如何在 Token Level 分配计算深度,作者可视化和对比了两种路由场景:(1)动态路由有效的典型情况,以及(2)所有 Token 被分配均匀浅层递归深度的退化情况。

相比之下,图5描绘了一种退化场景,其中几乎所有token都被分配了最小的递归深度,动态机制有效地简化为统一的静态计算调度。在这种情况下,空间图和直方图都显示出缺乏适应性:几乎所有token在第一次递归步骤后退出,无论其语义内容如何。

picture.image

这种比较突出了MOR-VIT的一个关键优势:当路由器成功学会区分token复杂度时,MoR-ViT能够实现高效的资源分配和强大的泛化能力;而当路由器退化时,潜在效率提升将丧失,类似于固定深度模型。这些可视化结果为机制的有效性提供了有力的直观证据,同时也为未来研究中识别失败案例或训练病理提供了诊断工具。

4 实验

作者系统地评估了MoRViT在表示能力和效率方面的表现,并与一系列最新的视觉Transformer架构和高效 Baseline 进行了比较。除了ResNet [1]、视觉Transformer (ViT) [2]和高效图像Transformer (DeiT) [15]等经典模型外,作者的评估还包括了与领先的ViT高效变体的全面横向比较,具体包括DynamicViT [3]、Adaptive Token Sampling (ATS) [4]、EfficientFormer [6]、PoolFormer (MetaFormer) [7]、MobileViT [8]、TinyViT [9]、TokenLearner [10]和SparseFormer [11]。这些模型代表了在tokenSparse化、动态路由、架构压缩和移动友好型Transformer设计方面的最先进方法。

为确保公平且严谨的比较,所有Transformer模型均在标准监督场景下进行训练,仅使用ImageNet数据集[16],且不使用额外的大规模数据或自监督预训练。对于下游迁移学习评估,作者遵循既定协议[2,15],并在公开的CIFAR[17]和Oxford Flowers[18]数据集上进行微调和测试。通过在这套多样化的模型和基准上进行性能评估,作者旨在突出MoRViT相对于传统及最新高效ViT方法的优缺点和权衡。

4.1 实验设置

MOR-VIT基于标准的ViT Backbone 网络(包含patch embedding和transformer encoder),通过无缝集成Mixture-of-Recursions(MoR)机制来增强模型的适应能力。这种模块化设计既保留了ViT的优点,又通过MoR模块引入了显著的新颖性。所有模型均在ImageNet数据集[16]上使用监督学习从头开始训练,未使用任何外部大规模数据或自监督预训练。

为了进行公平比较,作者重点关注基础级变体(ViT-B/16和MoR-B/16),并使用数据高效图像Transformer(DeiT-B/16)以及其他高效Transformer Baseline (TinyViT、EfficientFormer、MetaFormer、TokenLearner)作为外部比较基准。请注意,MOR-VIT架构并未继承任何DeiT特定的设计或修改;DeiT仅用于基准参考。所有模型均使用Adam优化器[19]进行200个epoch的训练,并采用原始ViT工作[2]中描述的标准数据增强和正则化协议。在下游评估中,模型遵循常见的迁移学习程序,在CIFAR-10/100[17]和Oxford Flowers102[18]上进行微调和测试。

在MoR-ViT的训练过程中未观察到任何异常波动或收敛问题。训练曲线与ViT-B/16相当,标准的提前停止标准有效。整体训练时间和资源消耗与同等规模的 Baseline Transformer模型相似。

所有对比模型的关键架构配置汇总于表1,从而实现全面的横向比较。

picture.image

对于所有Transformer模型,均使用基础级(B/16)或同等配置。TinyViT/MetaFormer/EfficientFormer的数值引用自其官方出版物。

4.2 实验结果

作者进行了广泛的实验,并系统地比较了所提出的MoR-ViT与十二种具有代表性的高效视觉Transformer方法(总共13种方法)。这些代表性方法可以分为两大类:第一类基于tokenSparse化和动态剪枝(例如DynamicViT、Adaptive Token Sampling、TokenLearner、SparseFormer、Token Clustering Transformer等),而第二类则专注于高效架构设计和模型压缩(例如EfficientFormer、PoolFormer、MobileViT、TinyViT等)。

与所有现有方法不同,MoR-ViT首次引入了token Level 的动态递归机制,直接针对先前方法中所有token使用固定深度和相同计算的根本性局限性。该机制允许为每个token自适应计算深度,从而有效解决了大多数高效ViT模型无法根据token重要性分配资源的问题。

值得注意的是,所有MoR-ViT结果均未使用任何蒸馏、额外的自监督预训练或外部大规模数据。仅通过标准监督训练即实现了极具竞争力的结果。

所有对比方法及其缩写如下:

MoR-ViT(MOR-VIT)ResNet-50(ResNet)Vision Transformer(ViT)数据高效图像Transformer(DeiT)

动态视觉 Transformer

自适应 Token 采样 (ATS)

高效Former(高效Former)

池化Transformer(元Transformer)

移动视觉 Transformer

微型视觉Transformer

TokenLearner

SparseFormer

Token Clustering Transformer (TCT)

主要表和补充表的格式如下(已移除表头,仅保留方法名称和空白占位符):

创新与发现:本研究首次为视觉Transformer引入了一种基于token的动态路由(动态递归)机制,使每个token能够根据特征复杂性和信息内容自适应地确定其计算深度。与其他所有方法相比,该机制能够实现更有效的计算资源分配,显著提升参数效率、推理速度和能耗,并在主要评估指标上达到了与最新高效ViT Baseline 高度竞争甚至超越的性能。MoR-ViT为未来高效视觉Transformer的设计和部署提供了新的技术范式和理论基础。

4.3 工程优势

近年来,视觉Transformer的进展主要集中在两个关键工程挑战上:传统架构的模型尺寸过大和计算负担过重,以及在实际部署约束下对高吞吐量推理的需求。作者的MoR-ViT通过集成紧凑的参数化、低复杂度的计算方案和高效的动态递归机制,直接针对这两个问题。

4.3.1 模型大小(参数)

如图6所示,MoR-ViT在模型参数数量上相较于传统Transformer Baseline 实现了显著减少,同时保持了与近期轻量级设计的竞争力。这种紧凑性不仅降低了内存消耗和存储需求,还简化了在实际部署中的硬件适配和模型更新。

picture.image

4.3.2 计算复杂度(FLOPs)

如图7所示,MoR-ViT在FLOPs方面始终低于大多数主流Transformer,使其能够在计算资源有限的设备上高效运行。该设计在大规模部署时进一步降低了功耗和运营成本。

picture.image

4.3.3 推理速度

MoR-ViT的动态递归设计在推理吞吐量上带来了显著优势,如图8所示。通过自适应分配计算资源,MOR-VIT实现了高处理速度而不牺牲准确性,使其非常适合需要快速且可靠预测的场景。

picture.image

4.3.4 泛化与可迁移性分析

为评估MoR-ViT在不同视觉领域的泛化能力和迁移能力,作者在多个下游分类数据集上进行了实验,包括CIFAR-10、CIFAR-100和Flowers-102。图9展示了Top-1准确率的结果。

picture.image

如图所示,MoR-ViT在所有评估数据集上均取得了极具竞争力的性能,其表现优于或与DeiT-B/16和ViT-B/16等其他强大的基于transformer的 Baseline 模型相当。值得注意的是,MoR-ViT在粗粒度和细粒度分类任务上均保持高度准确率,这表明其具有较强的泛化能力和适应多种下游场景的能力。

4.3.5 消融研究:动态递归和参数共享的影响

为了量化MoR-ViT中每个核心组件的贡献,作者逐一禁用它们,并在相同的训练设置下对ImageNet-1K进行评估。

主要发现禁用token Level 的动态递归导致了最大降幅(0.7 pp),证实了自适应深度是MoR-ViT的核心。关闭参数共享将模型膨胀至8600万参数,使推理速度降低约50%,但仅比基础ViT提升0.9 pp,突显了权重绑定的效率-精度平衡点。

5 结论

在这项工作中,作者提出了MoR-ViT,这是首个引入token级动态递归的视觉Transformer架构,从根本上解决了现有高效ViT模型应用固定计算深度和所有token进行统一处理的局限性。MOR-VIT使每个token能够自适应地确定其计算深度,从而实现更灵活和资源高效的计算分配。

在ImageNet-1K和各种下游基准测试上的大量实验表明,MoR-ViT在高达70%的参数减少和2.5倍的推理加速下实现了最先进的准确率,在可比设置下优于领先的高效ViT Baseline 模型。值得注意的是,这些提升是在不依赖外部大规模预训练或蒸馏的情况下实现的。消融研究证实了动态递归机制和参数共享对整体效率和性能的关键贡献。

展望未来,所提出的动态递归范式为视觉Transformer中的自适应计算未来研究开辟了充满希望的途径,包括将其扩展到更大规模的模型、更多样化的视觉任务以及实时边缘部署。未来的工作可以探索自动化路由策略、与神经架构搜索的集成以及递归深度分配的理论分析。作者希望MoR-ViT成为迈向真正可扩展和可部署的视觉识别系统的基石。

参考

[1]. MOR-VIT: Efficient Vision Transformer with Mixture-of-Recursions.

0
0
0
0
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论