AdaViT: Adaptive Vision Transformers for Efficient Image Recognition 解读

一、论文基本信息论文题目AdaViT: Adaptive Vision Transformers for Efficient Image Recognition作者Lingchen Meng、Hengduo Li、Bor-Chun Chen、Shiyi Lan、Zuxuan Wu、Yu-Gang Jiang、Ser-Nam Lim发表会议CVPR 2022官方代码MengLcool/AdaViT。论文 PDF 中也明确给出该代码地址GitHub 页面说明这是 CVPR 2022 论文的官方实现。这篇论文容易和A-ViT: Adaptive Tokens for Efficient Vision Transformer混淆。你前面看的 A-ViT 主要是token halting让每个 patch token 自己决定算到第几层而这篇AdaViT更全面它同时动态选择哪些 patch tokens 保留。哪些 self-attention heads 激活。哪些 Transformer blocks / sublayers 执行。所以 AdaViT 不是单纯 token pruning而是一个patch head block 三粒度联合动态推理框架。二、论文要解决什么问题Vision Transformer 的计算量来自三个主要维度patch token 数量。self-attention head 数量。Transformer block 层数。传统 ViT 对所有输入图像都使用相同计算路径同样数量的 tokens、同样数量的 heads、同样深度的 blocks。但图像难度差异很大。论文举例说明简单、目标居中的图像可能只需要少量 patch 和 head 就能分类复杂背景、遮挡、多物体场景则需要更多上下文和更深计算。论文图 1 也用 “White Stork” 和 “Barbershop” 展示了简单样本与复杂样本的不同计算需求。所以这篇论文的核心问题是ViT 是否可以像动态网络一样根据每张输入图像的难度自适应决定使用多少 patch、多少 head、多少 blockAdaViT 的答案是可以。它在每个 Transformer block 前加入轻量decision network为当前输入动态产生使用策略从而减少冗余计算。三、核心思想AdaViT 的核心思想是给每张图像生成一套 instance-specific usage policy让简单图像走轻量路径复杂图像走更完整路径。具体来说在第 l 个 Transformer block 前AdaViT 会产生三类二值决策Patch selection当前 block 中哪些 patch tokens 继续参与计算。Head selection当前 multi-head self-attention 中哪些 attention heads 被激活。Block selection当前 block 的 MSA 和 FFN 是否执行。这就形成了一种输入自适应的计算图。对于简单图像模型可能只保留少量 patch、少量 head并跳过部分 block对于复杂图像模型会保留更多 token、head 和 block。论文摘要明确说AdaViT 学习的是每个输入样本对应的 patch、head、block 使用策略并通过轻量 decision network 在线生成决策。四、它“剪”的是什么AdaViT 不是传统参数剪枝。它不直接剪掉模型权重也不是训练后得到一个固定小模型。它做的是动态结构选择动态 token pruning每张图像、每一层保留不同 patch tokens。动态 head pruning每张图像、每一层激活不同 attention heads。动态 block skipping每张图像、每一层可以跳过 MSA 或 FFN。所以它更准确的定位是dynamic inference / adaptive computation for ViTs。如果从剪枝角度分类可以叫输入自适应的多粒度结构化动态剪枝。但要注意它的剪枝结果不是一个固定结构而是每张图像一套计算路径。五、Decision Network如何产生动态策略AdaViT 在每个 Transformer block 前插入一个轻量决策网络。论文说第 l 个 block 的 decision network 由三个线性层组成分别输出 patch selection、head selection 和 block selection 的策略。每个输出经过 sigmoid 后表示对应结构被保留或激活的概率。这三个输出分别是patch policy长度为 N对应 N 个 patch tokens。head policy长度为 H对应 H 个 attention heads。block policy控制当前 block 是否执行。更细一点论文把 block policy 扩展成两个维度分别控制MSA 是否执行。FFN 是否执行。这样比整块跳过更灵活。因为有些情况下 attention 可能有用但 FFN 可以跳过或者反过来模型可以按需保留不同子层。论文公式 8 就把原始残差结构改成了带二值 mask 的 MSA / FFN 控制形式。六、Patch Selection如何减少 token在每个 block 输入处AdaViT 会决定哪些 patch tokens 保留。如果某个 patch 的 mask 为 0它就不进入当前 block 的计算。class token 始终保留因为最终分类依赖 class token。论文在 patch selection 部分明确说明class token always kept。这和 DynamicViT、ATS、EViT 属于同一类 token reduction 思路减少后续 attention 和 FFN 要处理的 token 数。但 AdaViT 的特殊之处是它不是只做 token selection而是和 head selection、block selection 一起联合学习。也就是说模型可以同时回答三个问题这个 patch 要不要算这个 head 要不要开这个 block 要不要走七、Head Selection部分关闭与完全关闭AdaViT 对 attention head 提出了两种关闭方式。第一种Partial deactivation。不真正删除 head而是把 attention softmax 输出替换成预定义矩阵例如 identity matrix。这样可以省掉 attention map 的计算但仍保留 value 相关输出。精度更稳但省计算有限。第二种Full deactivation。直接移除整个 attention head并相应减少 MSA 输出维度。这样节省更多计算但更容易造成精度损失。论文也指出在相同关闭比例下full deactivation 更省计算但由于动态改变 embedding 结构精度风险更高。实验表明在关闭 50% heads 时partial deactivation 可以达到81.7% Top-1 / 6.9 GFLOPsfull deactivation 是80.3% Top-1 / 5.1 GFLOPs也就是说partial 更准full 更省。这个设计很有意思因为它说明 head pruning 并不是只有“删掉 head”一种方式。AdaViT 把 head selection 做成了可按预算选择的两种工程实现。八、Block Selection跳过 MSA 和 FFNAdaViT 还可以跳过 Transformer block 的子层。由于 Transformer 有残差连接如果某个 MSA 或 FFN 被认为冗余可以直接跳过这一部分让输入通过 residual path 传到下一层。论文没有只做整块 block skipping而是分别控制MSA 子层。FFN 子层。这比传统 SkipNet / LayerDrop 式整层跳过更细。因为 ViT 中 MSA 和 FFN 的功能不同有些图像可能仍需要 attention 做全局聚合但不一定需要完整 FFN也可能某些层整体都较冗余。九、训练难点二值决策不可导AdaViT 的策略都是二值决策保留或丢弃激活或关闭执行或跳过。这种二值选择不能直接反向传播。论文没有采用强化学习因为策略维度很高policy gradient 方差大、收敛慢。作者选择使用Gumbel-Softmax对离散采样过程做可微松弛让 decision network 和 backbone 可以端到端联合训练。训练目标由两部分组成分类交叉熵损失。保证 ImageNet 分类精度。usage loss。约束 patch、head、block 的平均使用率接近目标预算。这里有三个预算超参数γp目标 patch 保留比例。γh目标 head 激活比例。γb目标 block / sublayer 使用比例。论文明确说γp、γh、γb 控制希望保留的 patch、head、block 百分比因此 AdaViT 可以通过调节这些超参数适配不同计算预算。十、为什么 AdaViT 能自适应不同图像因为 decision network 的输入来自当前 block 的 token 表示而这个表示已经包含当前图像的信息。如果图像简单比如主体清晰、背景干净decision network 会倾向于保留少量 patch、少量 head、跳过更多 block。如果图像复杂比如背景杂乱、多物体、遮挡严重decision network 会保留更多计算。论文分析中也发现AdaViT 会给 “shoe shop”“barber shop”“toyshop” 这类复杂场景分配更多计算而给 “parachute”“kite”等目标更简单、物体更中心的类别分配更少计算。这正是 AdaViT 的核心价值不是所有图像都走同样算力路径。十一、实验设置论文主要在ImageNet-1K上实验。Backbone 使用T2T-ViT-19因为它在 ImageNet 上性能较强且计算量适中。实验设置中T2T-ViT backbone 有19 个 Transformer blocks。每个 MSA 有 7 个 heads。输入 token 数 N196。Decision network 从第 2 个 block 开始插入。训练时使用 T2T-ViT 官方预训练权重初始化 backbone然后联合训练 150 epochs。论文报告 Top-1 accuracy 和 GFLOPs。十二、主要实验结果论文最核心结果在表 1。原始T2T-ViT-19Top-181.9%FLOPs8.5GAdaViTTop-181.1%FLOPs3.9G也就是说AdaViT 用约46% 的计算量保持了很高精度精度只下降约0.8%。论文也明确说相比原始 T2T-ViTAdaViT 达到超过2× efficiency improvement只带来约 0.8% accuracy drop。和其他模型相比AdaViT 在 3.9G FLOPs 下达到 81.1%比很多静态 ViT / CNN 的效率—精度折中更好。例如 DeiT-S 是 79.9% / 4.6GPVT-Small 是 79.8% / 3.8GSwin-T 是 81.3% / 4.5G。所以 AdaViT 的卖点是比原始 T2T-ViT 省很多计算。比相近 FLOPs 的静态模型保持更强精度。十三、随机策略对比学出来的策略真的有用吗为了证明不是“随便少算一点也能行”论文做了随机策略对比。在与 AdaViT 类似的计算预算下Random33.0% Top-1Random71.5% Top-1AdaViT81.1% Top-1Random 表示随机策略后再 fine-tune仍然比 AdaViT 低很多。论文进一步指出AdaViT 在 3.9 GFLOPs 下比 Random 高48.1%比 Random 高9.6%。这说明 AdaViT 学到的 usage policy 确实有意义不是简单随机删 patch、删 head、跳 layer。十四、三类策略各自是否有效论文把三类策略分别随机替换观察精度变化。结果是随机 patch policy49.2% Top-1随机 head policy57.4% Top-1随机 block policy64.7% Top-1完整 AdaViT81.1% Top-1这说明三类策略都很重要。随机替换任何一类都会严重破坏结果。论文还指出在类似计算预算下AdaViT 比随机 patch/head/block 分别高31.9% / 23.7% / 16.4%。从这个消融可以看出patch selection 最敏感。head selection 也很关键。block selection 相对鲁棒一些但仍不能随机。十五、计算分配有什么规律论文对 learned usage policies 做了可视化和统计得到几个很有价值的结论。第一patch 数量随着层数加深逐渐减少。论文发现平均保留 patch 数从浅层到深层逐步下降。这符合 ViT 的信息聚合过程经过多层 attention 后patch 之间已经交换了信息后面只需要少量判别性 patch 就能代表整张图像。第二head 和 block 在最后几层保留更多。与 patch 不同head selection 和 block selection 在最后几层倾向于保留更多计算。论文推测这是因为最后几层更直接负责最终分类预测因此不能过度跳过。第三简单类别用更少计算复杂类别用更多计算。例如 parachute、kite 等简单目标类别计算较少barbershop、shoe shop、toyshop 等复杂场景计算较多。这个分析证明 AdaViT 学到的是一种符合直觉的动态计算分配。十六、和 DynamicViT 的区别DynamicViT 主要做token sparsification用 prediction module 预测 token 重要性并逐层删除冗余 tokens。AdaViT 更广DynamicViT主要动态删 patch tokens。AdaViT同时动态选择 patch tokens、attention heads、Transformer blocks。DynamicViT 的计算控制更集中在序列长度AdaViT 则同时控制序列长度、注意力分支数量和网络深度。所以 AdaViT 更像多粒度动态 ViT而 DynamicViT 是token-level 动态 ViT。十七、和 A-ViT 的区别A-ViT 是 token halting。它让每个 token 学习自己在哪一层停止计算不加额外 halting network核心是 ACT 式 token early stopping。AdaViT 不做 halting depth而是在每个 block 处用 decision network 产生三类 usage policy。可以这样区分A-ViT每个 token 决定“我算到第几层”。AdaViT每个 block 决定“哪些 patch、哪些 head、哪些子层要用”。A-ViT 更轻主要控制 token 深度AdaViT 更全面但也更复杂需要额外 decision network 和 Gumbel-Softmax 训练。十八、和 ATS 的区别ATS 是 parameter-free 的 adaptive token sampling。它利用 class token attention 和 value norm 给 token 打分再通过 inverse transform sampling 选择 tokens。AdaViT 不是 parameter-free。它需要训练一个轻量 decision network并且同时控制 patch、head、block。所以ATS无参数主要做 token sampling。AdaViT有决策网络做 patch/head/block 联合策略。ATS 更适合低侵入插入已有模型AdaViT 更适合端到端训练一个多粒度动态推理模型。十九、和 ToMe / TokenLearner 的区别ToMe是 token merging。它把相似 tokens 合并通常可以 training-free 地加速已有 ViT。TokenLearner是 learned tokenization。它学习生成 8 或 16 个新 tokens后续层只处理 learned tokens。AdaViT不合并 token也不生成新 token而是决定已有 token、head、block 是否使用。所以三者的问题不同ToMe哪些 tokens 相似可以合并TokenLearner应该学习生成哪些少量新 tokensAdaViT当前图像需要哪些 patch、head、block二十、它是不是剪枝严格说AdaViT 不是传统静态剪枝。传统剪枝得到一个固定子网络所有输入都使用同一个结构。AdaViT 得到的是动态使用策略不同输入使用不同结构。因此它更接近动态推理。条件计算。输入自适应结构化剪枝。如果写论文综述可以把它放在ViT dynamic inference / adaptive computation。而不是只放在 token pruning 里。因为它不只减少 tokens还动态控制 heads 和 blocks。二十一、方法优点第一多粒度动态计算。它同时控制 patch、head、block比单纯 token pruning 更全面。第二输入自适应。简单图像少算复杂图像多算符合视觉样本复杂度差异。第三端到端训练。通过 Gumbel-Softmaxdecision network 和 backbone 可以联合优化。第四预算可控。γp、γh、γb 可以控制 patch、head、block 的目标使用率从而适配不同 FLOPs 预算。第五实验结果较强。在 T2T-ViT-19 上FLOPs 从 8.5G 降到 3.9GTop-1 从 81.9% 到 81.1%实现超过 2× 的效率提升。二十二、方法局限第一训练复杂。需要额外 decision network、Gumbel-Softmax、usage loss 和预算超参数。相比 ToMe、ATS 这类低侵入方法更重。第二动态结构部署不一定简单。不同输入保留不同 patch、head、block实际 GPU batch 推理时可能需要动态 shape、mask、padding 或分组调度否则理论 FLOPs 节省不一定完全转化为真实吞吐。第三主要验证 ImageNet 分类。分类任务只需要最终 class token因此适合动态删 patch 和跳 block。检测、分割等 dense prediction 任务需要更谨慎因为空间细节不能随意丢。第四仍有精度损失。论文也在 limitation 中承认相比 upperbound baseline 仍存在小幅 accuracy drop。第五backbone 主要是 T2T-ViT。虽然思想可以迁移到其他 ViT但论文核心实验主要围绕 T2T-ViT-19而不是现代更大规模 ViT、Swin、MAE、DINOv2 等 backbone。二十三、整体评价AdaViT 的核心价值在于它把 ViT 的动态推理从单一 token 维度扩展到了 patch、head、block 三个维度。前面很多方法都是问哪些 token 可以删而 AdaViT 同时问哪些 token 要算哪些 attention heads 要开哪些 Transformer 子层要执行这使它成为一篇比较典型的multi-granularity adaptive computation论文。它的思路很直观也很符合 ViT 的冗余结构图像中有冗余 patchattention 中有冗余 head深层网络中也有冗余 block。不同样本的冗余程度不同因此应该动态分配计算量。不过从部署角度看AdaViT 也比 ToMe、ATS 这类方法更复杂。它的理论 FLOPs 降低很明显但真正落地时需要处理动态 token 数、动态 head 数和动态 block 跳过带来的运行时调度问题。二十四、一句话总结《AdaViT: Adaptive Vision Transformers for Efficient Image Recognition》提出一种多粒度动态 ViT 推理框架在每个 Transformer block 前加入轻量 decision network为每张输入图像动态决定保留哪些 patch tokens、激活哪些 self-attention heads、执行哪些 MSA/FFN 子层并用 Gumbel-Softmax 和 usage loss 实现端到端训练与预算控制。它不是单纯 token pruning而是 patch、head、block 联合自适应计算在 ImageNet 上以 T2T-ViT-19 为 backbone将 FLOPs 从 8.5G 降到 3.9GTop-1 仅从 81.9% 降到 81.1%证明 ViT 中存在可按输入动态利用的大量结构冗余。