图像未必值16x16词：可变序列长度的动态视觉Transformer来了

©作者 | 王语霖

单位 | 清华大学自动化系

研究方向 | 机器学习、计算机视觉

本文主要介绍刚刚被 NeurIPS 2021 会议录用的一篇关于动态 Transformer 的最新工作：Not All Images are Worth 16x16 Words: Dynamic Vision Transformers with Adaptive Sequence Length，全部代码和预训练模型已经在 Github 上开源。

论文标题：

Not All Images are Worth 16x16 Words: Dynamic Vision Transformers with Adaptive Sequence Length

论文来源：

NeurIPS 2021

论文链接：

https://arxiv.org/abs/2105.15075

代码链接：

https://github.com/blackfeather-wang/Dynamic-Vision-Transformer

以 ViT 为代表的视觉 Transformer 通常将所有输入图像表征为固定数目的 tokens（例如 16x16 和 14x14）。这项工作发现采用定长的token序列表征数据集中所有的图像是一种低效且次优的做法，并提出一种可针对每个样本自适应地使用最合适的 token 数目进行表征的动态 ViT 模型。该方法在 ImageNet 上将 T2T-ViT 的平均推理速度（GPU 实测）加快了 1.4-1.7 倍。其主要思想在于利用级联的 ViT 模型自动区分“简单”与“困难”样本，实现自适应的样本推理。为了减少级联模型中的冗余计算，文章还提出了特征重用与关系重用的模型设计思路。

▲ 图1 Dynamic Vision Transformer (DVT) 示例

研究动机及简介

近半年来，以 Google 的工作《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》为代表的 Vision Transformer (ViT) 系列视觉模型受到了学界的广泛关注，这些模型通常将图像数据划分为固定数目的 patch，并将每个 patch 对应的像素值采用线性映射等方式嵌入为一维的 token，作为 Transformer 模型的输入，示意图如下所示。

▲ 图2 Vision Transformer (ViT) 对输入的表征方式

假设模型结构固定，即每个 token 的维度大小固定，将输入表征为更多的 token 可以实现对图片更为细粒度的建模，往往可以有效提升模型的测试准确率；然而与此相对，由于 Transformer 的计算开销随 token 数目成二次方增长，增多 token 将导致大量增长的计算开销。为了在精度和效率之间取得一个合适的平衡，现有的 ViT 模型一般将 token 数设置为 14x14 或 16x16。

而论文则提出，一个更合适的方法应当是，根据每个输入的具体特征，对每张图片设置对其最合适的 token 数目。具体而言，不同图片在内容、远近、物体大小、背景、光照等诸多方面均存在较大的差异，将其切分为相同数目和大小的 patch 没有考虑适应这些要素的变化，因而极有可能是次优的。以下图为例，左侧的苹果图片构图简单且物体尺寸较大，右侧图片则包含复杂的人群、建筑、草坪等内容且物体尺寸均较小。显然，前者只需要少量 token 就可以有效表征其内容，后者则需要更多的 token 描述不同构图要素之间的复杂关系。

▲ 图3 根据具体输入确定token数目

这一问题对于网络的推理效率是非常关键的。在下表中，文章使用比原文推荐值（14x14）更少的 token 数目训练了一个 T2T-ViT-12 模型，并报告了对应的测试精度和计算开销。从结果中可以看到，若将 token 数目设置为 4x4，准确率仅仅下降了 15.9%（76.7% v.s. 60.8%），但计算开销下降了 8.5 倍（1.78G v.s. 0.21G）。这一结果表明，正确识别占数据大多数的较“简单”的样本只需 4x4 或更少的 token，相当多的计算浪费在了使用存在大量冗余的 14x14 token 表征他们。

▲ 表1 T2T-ViT-12使用更少token时的测试精度和计算开销

 

方法详述

受到上述现象的启发，论文提出了一个动态视觉 Transformer 框架（Dynamic Vision Transformer, DVT），意在针对每个样本选择一个合适数目的 token 来进行表征。

首先介绍 DVT 的推理过程，如图 2 所示。论文使用从小到大的 token 数目训练了一组 Transformer 模型，他们具有相同的基本结构，但是参数相互独立，以分别适应逐渐增多的 token 数目。对于任意测试样本，首先将其粗略表征为最小数目的 token，输入第一个 Transformer，以低计算成本迅速得到初步预测。而后判断预测结果是否可信，若可信，则直接输出当前结果并终止推理；若不可信，则将输入图片表征为更多的 token，激活下一个 Transformer 进行更细粒度、但计算开销更大的的推理，得到结果之后再次判断是否可信，以此类推。

论文采用将预测的置信度（confidence）与一个固定阈值进行比较的方式作为准出的判断准则，关于这一方法的细节和其合理性的验证，由于空间有限，请参见 paper。

值得注意的是，处于下游的 Transformer 可以复用先前模型已产生的深度特征和注意力图（attention maps），以尽可能避免进行冗余和重复的计算，关于这一点，将在下文详述。

▲ 图4 Dynamic Vision Transformer (DVT)

在训练时，论文简单地训练网络在所有出口都取得正确的预测结果，训练目标如下式所示。其中 和 分别代表数据和标签， 代表第 个出口的softmax预测概率， 代表交叉熵损失。

显然，DVT 框架中所有的 Transformer 都具有相同的训练目标。因此当一个处于下游位置的 Transformer 被激活时，一个显然更为高效的做法是，应当训练其在先前 Transformer 已提取的特征的基础上进行进一步提升，而非完全从 0 开始重新提取特征。出于这一点考虑，论文提出了一个特征复用机制，如下图所示。将上游 Transformer 最后一层输出的特征取出，经 MLP 变换和上采样（Upsample）后，作为上下文嵌入（Context Embedding）以 Concat 的方式整合入下游模型每一层的 MLP 模块中。

▲ 图5 特征复用（Feature Reuse）机制

除特征外，Transformer 中还有一个关键的组分，即自注意力机制中的注意力图（attention maps）。视觉 Transformer 通过这些注意力图建模 token 之间的关系，从而实现对每个具体的 token 进行来自全局的信息整合。在论文所提出的 DVT 框架中，位于下游的模型同样可以借助上游 Transformer 已经得到的注意力图来进行更准确的全局注意力关系建模，论文称之为关系复用，示意图如下所示。

将上游模型的全部注意力图以 logits 的形式进行整合，经 MLP 变换和上采样后，加入到下层每个注意力图的 logits 中。这样，下游模型每一层的 attention 模块都可灵活复用上游模型不同深度的全部 attention 信息，且这一复用信息的“强度”可以通过改变 MLP 的参数自动地调整。

▲ 图6 关系复用（Relationship Reuse）机制

值得注意的是，对特征图进行上采样需要对其行或列进行重组后分别完成，以确保其几何关系的对应性，下图给出了一个例子。

▲ 图7 对注意力图进行上采样

由于空间有限，关于 DVT 的更多细节以及对两种复用机制有效性的验证、结构设计的探究和更进一步的讨论，请参见 paper。

实验结果

DVT 的一个显著优势在于，大多数的视觉 Transformer 均可作为其的 backbone 以获得更高的计算效率，在实验中，论文基于 T2T-ViT 和 DeiT 测试了所提出的的方法。

DVT（T2T-ViT-12/14）在 ImageNet 图像识别任务上的计算效率如下，可见该方法对 backbone 的提速比在 1.6-3.6x，对大模型效果尤为明显。

▲ 图8 DVT（T2T-ViT）在ImageNet上的计算效率

DVT（T2T-ViT-14）在 ImageNet 上的实际推理速度如下，基于 Nvidia 2080Ti GPU，batch size=128，在每个出口移除终止推理的样本。对于小模型，DVT 可以在不减慢速度的前提下显著提升效果；对于大模型，DVT 可以在不降低表现的前提下显著提速。

▲ 表2 DVT (T2T-ViT) 在ImageNet上的实际速度

DVT（T2T-ViT-12/14）在 CIFAR 图像识别任务上的计算效率如下。

▲ 表3 DVT (T2T-ViT) 在CIFAR上的计算效率

DVT（DeiT）在 ImageNet 图像识别任务上的计算效率如下，该方法效果同样明显。

▲ 图9 DVT (DeiT) 在ImageNet上的计算效率

那么，什么样的样本适合更少的 token，什么样的样本适合更多的 token 呢？下图给出了可视化的结果，“easy”和“hard”分别代表需要少和多的 token 数目。可见，后者往往包含复杂的场景、较小的物体尺寸、以及一些非常规的姿态和角度。

▲ 图10 可视化结果

结语

这篇工作的价值至少体现在两个方面：

1. 针对视觉 Transformer 提出了一个自然、通用、有效的自适应推理框架，理论和实际效果都比较显著。

2. 提出了一个颇具启发性的思路，即目前大多数视觉 Transformer 采用的、对全部图片以固定方式划分 patch 的表征方式，是不够灵活和次优的，一个更合理的策略是，应当根据输入数据动态调整表征方式。这或为开发高效、可解释性强、迁移性好的通用视觉 Transformer 提供了新方向。

如有任何问题，欢迎留言或者给我发邮件，附上我的主页链接：
http://www.rainforest-wang.cool/

特别鸣谢

感谢 TCCI 天桥脑科学研究院对于 PaperWeekly 的支持。TCCI 关注大脑探知、大脑功能和大脑健康。

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：hr@paperweekly.site 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·