8个token能学到什么？谷歌提出《TokenLearner》，用8个token就可以达到优于数百个token的结果！

关注公众号，发现CV技术之美

本篇文章分享论文『TokenLearner: What Can 8 Learned Tokens Do for Images and Videos?』，谷歌提出《TokenLearner》，用 8 个 token 就可以达到优于数百个 token 的结果！

详细信息如下：

论文链接：https://arxiv.org/abs/2106.11297
项目链接：https://github.com/google-research/scenic/tree/main/scenic/projects/token_learner

摘要

在本文中，作者介绍了一种新的视觉表征学习方法，它依赖于少量自适应学习的tokens，适用于图像和视频理解任务。

本文的方法不是依靠手工设计的分割策略来获取视觉token，并处理大量密集采样的patch，而是学习挖掘视觉数据中的重要token。这可以高效地找到一些重要的视觉token，并在更长的视频时间范围内或图像中的空间内容上，对这些token之间的成对注意力进行建模。

实验表明，在图像和视频识别任务的几个具有挑战性的基准上，本文方法的性能都很好。重要的是，由于本文的token是自适应的，因此以显著减少的计算量实现了有竞争力的结果。在计算效率更高的同时，本文方法获得了与ImageNet上的最新技术相当的结果，并在多个视频数据集上建立了SOTA结果。

Motivation

图像和视频提供了丰富的视觉信息。图像理解是计算机视觉中一个长期存在的问题，尽管取得了惊人的进步，但为各种图像理解任务获得最佳视觉表现仍然是一个活跃的研究领域。除了解决类似的图像理解任务外，视频还需要对RGB和时间流进行有效的时空处理，以捕捉远程交互。这种理解的一个重要方面是如何快速了解输入视频流的哪些部分在空间和时间上都很重要，并将计算资源集中在这些部分上。但是什么样的基本处理机制能够成功地处理图像和视频呢？

图像理解的最新进展表明，视觉分类任务的准确性有所提高。例如，与标准卷积方法不同，Vision Transformer（VIT）将图像视为一系列patch，这里利用了类似于文本理解的Transformer架构。

视频识别的标准方法将视频视为堆叠图像（即时空体），并倾向于将2D神经结构扩展到3D（例如3D Resnet）。基于图像的视觉Transformer，一些方法提出在常规3D网格上创建3D“cubelet”视频token，这些token由Transformer进一步处理，从而生成计算量大的模型。要处理的token太多，尤其是较长的视频。

这项工作的主要问题是如何从视觉输入中自适应地学习表示，以最有效地捕捉图像的空间信息和视频的时空交互。

第一个关键观察是，能够通过学习“tokenize输入来学习表示视觉数据。这与之前的方法不同，之前的方法对图像或视频进行密集采样，例如16x16或32x32。具体来说，可以学习计算输入图像/视频中的重要区域，使token适应输入数据。作者使用空间注意机制计算每帧的多个空间权重图，并将其用于tokenization。这些map的目的是了解哪些区域是重要的。这里，每个空间权重映射与输入相乘，形成一个“token”，由后续学习模块处理。

此外，作者发现，很少有token能够胜任视觉理解任务。更具体地说，对于图像，当插入8-16个token作为中间表示（而不是保留200∼500个）时，可以显著减少ViT的计算预算。本文的TokenLearner能够将总FLOPS减少一半，同时保持甚至提高分类精度。类似地，对于视频识别，本文在三个具有挑战性的数据集上显示了比最新技术更高的性能，同时每帧仅使用8-16个中间token。

本文的方法简单、有效，并且如结果所示，优于现有技术中的卷积方法和以前的space-time Transformer方法。作者证明了本文的方法与ImageNet上以前的Transformer模型的性能相当，同时有意义地降低了计算成本。在视频理解任务中，作者在Kinetics400、Kinetics600、, Charades和AViD数据集上建立了新的SOTA性能，其性能优于以前的模型。

方法

在ViT等视觉Transformer结构中，输入图像首先通过将其拆分为小空间patch（如16x16）来tokenize，这些patch用作模型的输入。类似地，在最近的video transformer结构中，例如ViViT和TimeSformer，通过将视频切割成规则网格上的2d空间或3d时空立方体来tokenize视频。

本文的注意力模块不需要处理固定的、tokenized的输入，而是学习用于识别任务的token。通过这样做，可以获得了几个重要特性：

作者启用了自适应tokenization，以便可以根据输入动态选择token。
这也有效地减少了Transformer的token总数，考虑到视频中有许多token（例如4096），并且计算是token数量的二次方，这一点尤其有利。
最后，作者为每个后续层提供了学习依赖不同时空token的能力，这可能允许不同层捕获视频的不同方面。这些动态和自适应生成的token可用于标准Transformer结构中，如用于图像的ViT和用于视频的ViViT，或可用于专用视频结构中。

3.1 TokenLearner

上图展示了本文TokenLearner的示意图。

Compute reduction in Transformers

学习到的token被提供给视觉表征学习的后续层，例如ViT和ViViT中使用的多头自注意力（MHSA）。对于TokenLearner，这些后续层只需要处理少量的token（例如，8而不是1024），这显著减少了计算量，因为它们与token数量成二次关系。

上图（a）显示了在ViT中插入TokenLearner模块的基本架构。它可以添加到网络中的任何位置，由于token数量的巨大差异，token学习器之后Transformer层的相对计算几乎可以忽略不计。

3.2 TokenFuser

在TokenLearner生成token并随后的Transformer层（如MHSA）处理它们之后，“TokenFuser”可用于进一步融合token上的信息，以及将表示重新映射回其原始空间分辨率。这使模型能够捕获由token表示的空间（或时空）模式，并在必要时恢复原始输入张量形状。

上图展示了TokenFuser的整个过程。

实验

上图展示了不同位置插入TokenLearner和用不同数量Token的实验结果。可以发现，在网络中间插入TokenLearner（1/2）可以获得与baseline几乎相同的精度，同时将计算量减少（几乎）一半。此外，由于TokenLearner的适应性，与不使用TokenLearner的情况相比，将TokenLearner置于后网络的3/4之后可以获得更高的性能，同时比baseline更快。

上图和上表显示了使用较小ViT-S和ViT-B模型的ImageNet 微调的评估结果。

上图展示了few-shot的实验结果。

上表展示了本文方法在大模型上的实验结果。我们可以观察到，与对ViT-S和ViT-B的实验类似，TokenLearner在在中间插入时能够节省一半的计算量。

上表将TokenLearner模型与其他较大的ViT模型进行了比较。尽管使用的参数数量要少得多，但TokenLearner模型的性能与大型和巨型ViT模型相当。

作者比较了带和不带TokenFuser模块的TokenLearner模型，以确认TokenFuser的有效性，结果如上表所示。

上表比较了TokenLearner和空间池化的结果。我们可以观察到，token ‘learning’是有好处的。基于池化的token缩减确实具有与TokenLearner类似的计算，但与baseline模型相比，它失去了准确性。另一方面，尽管计算量较低，但TokenLearner的性能略优于baseline模型。