豪取四个SOTA，谷歌魔改Transformer登NeurIPS 2021

谷歌改造Vision Transformer的新作被NeurIPS 2021收录了。在这篇文章里，谷歌提出了TokenLearner方法，Vision Transformer用上它最多可以降低8倍计算量，而分类性能反而更强！

目前，Transformer模型在计算机视觉任务（包括目标检测和视频分类等任务）中获得了最先进的结果。

不同于逐像素处理图像的标准卷积方法，Vision Transformer（ViT）将图像视为一系列patch token（即由多个像素组成的较小部分图像）。

这也就意味着在每一层神经网络中，ViT模型使用多头自注意力（multi-head self-attention），基于每对token之间的关系来处理patch token。

这样，ViT模型就能够构建整个图像的全局表示。

在输入端，将图像均匀地分割成多个部分来形成token，例如，将512×512像素的图像分割成16×16像素的patch token。在中间层，上一层的输出成为下一层的token。

这里插一句。如果处理的是视频，则视频「管道」如16x16x2视频片段（2帧16x16图像）就成为了token。视觉token的质量和数量决定了Vision Transformer的整体性能。

许多Vision Transformer结构面临的主要挑战是，它们通常需要太多的token才能获得合理的结果。

例如，即使使用16x16patch token化，单个512x512图像也对应于1024个token。对于具有多个帧的视频，每层可能都需要处理数万个token。

考虑到Transformer的计算量随着token数量的增加而二次方增加，这通常会使Transformer难以处理更大的图像和更长的视频。

这就引出了一个问题：真的有必要在每一层处理那么多token吗？

谷歌在「TokenLearner：What Can 8 Learned Tokens Do for Images and Videos?」中提到了「自适应」这个概念。这篇文章将在NeurIPS 2021上进行展示。

论文地址：https://arxiv.org/pdf/2106.11297.pdf

项目地址：

https://github.com/google-research/scenic/tree/main/scenic/projects/token_learner

实验表明，TokenLearner可以自适应地生成更少数量的token，而不是总是依赖于由图像均匀分配形成的token，这样一来，可以使Vision Transformer运行得更快，性能更好。

TokenLearner是一个可学习的模块，它会获取图像张量（即输入）并生成一小组token。该模块可以放置在Vision Transformer模型中的不同位置，显著减少了所有后续层中要处理的token数量。

实验表明，使用TokenLearner可以节省一半或更多的内存和计算量，而分类性能却并不会下降，并且由于其适应输入的能力，它甚至可以提高准确率。

TokenLearner是啥？

TokenLearner其实是一种简单的空间注意力方法。

为了让每个TokenLearner学习到有用的信息，先得计算一个突出的重要区域的空间注意力图（使用卷积层或MLP）。

接着，这样的空间注意力图会被用来对输入的每个区域进行加权（目的是丢弃不必要的区域），并且结果经过空间池化后，就可以生成最终的学习好了的token。

应用于单个图像的TokenLearner模块的直观图示

TokenLearner学习在张量像素的子集上进行空间处理，并生成一组适应输入的token向量。

这种操作被并行重复多次，就可以从原始的输入中生成n个（10个左右）token。

换句话说，TokenLearner也可以被视为基于权重值来执行像素的选择，随后进行全局平均。

值得一提的是，计算注意力图的函数由不同的可学习参数控制，并以端到端的方式进行训练。这样也就使得注意力函数可以在捕捉不同输入中的空间信息时进行优化。

在实践中，模型将学习多个空间注意力函数，并将其应用于输入，并平行地产生不同的token向量。

TokenLearner模块学习为每个输出标记生成一个空间注意力图，并使用它来抽象化输入的token

因此，TokenLearner使模型能够处理与特定识别任务相关的少量token，而不是处理固定的、统一的token化输入。

也就是说，TokenLearner启用了自适应token，以便可以根据输入动态选择token，这一做法有效地减少了token的总数，大大减少了Transformer网络的计算。

而这些动态自适应生成的token也可用于标准的Transformer架构，如图像领域的ViT和视频领域的ViViT（Video Vision Transformer）。

TokenLearner放在哪？

构建TokenLearner模块后，下一步就必须要确定将其放置在哪个位置。

首先，研究人员尝试将它放置在标准ViT架构中的不同位置，输入图像使用224x224的大小。

TokenLearner生成的token数量为8个和16个，远远少于标准ViT使用的196个或576个token。

下图显示了在ViT B/16中的不同相对位置插入TokenLearner的模型的ImageNet 5-shot分类精度和FLOPs，其中ViT B/16是一个基础模型，有12个注意力层。其运行时使用16x16大小的patch token。

在JFT 300M的预训练下，ImageNet的5-shot精度与ViT B/16中TokenLearner的相对位置有关

位置0意味着TokenLearner被置于任何Transformer层之前。其中，baseline是标准的ViT B/16的ImageNet 5-shot分类精度和FLOPs。

计算量以数十亿次浮点运算（GFLOPS）衡量

我们发现，在网络的最初四分之一处（1/4处）插入TokenLearner，实现了与基线几乎相同的准确性，同时将计算量减少到基线的三分之一以下。

此外，将TokenLearner放在后面一层（网络的3/4之后），与不使用TokenLearner相比，取得了更好的性能，同时由于其适应性，性能更快。

由于TokenLearner前后的token数量相差很大（例如，前196个，后8个），TokenLearner模块后的相对计算量几乎可以忽略不计。

TokenLearner VS ViT

将带有TokenLearner的ViT模型和普通的ViT模型进行对比，同时在ImageNet的few-shot上采用相同的设置。

TokenLearner会被放置在每个ViT模型中间的不同位置，如网络的1/2和3/4处。其中，模型通过JFT 300M进行预训练。

从图上观察可以得知，TokenLearner模型在准确率和计算量方面的表现都比ViT要好。

不同版本的ViT模型在ImageNet分类上的表现

在更大的ViT模型中插入TokenLearner，如具有24个注意力层，并以10x10（或8x8）个patch作为初始token的L/10和L/8。

之后，将这两个模型与48层的ViT G/14模型进行比较。

可以看到，在表现和G/14模型相当的情况下，TokenLearner只需要非常少的参数和计算量。

左：大规模TokenLearner模型与ViT G/14在ImageNet数据集上的分类精度对比；右：参数量和FLOPS的对比

高性能视频模型

视频理解是计算机视觉的关键挑战之一，TokenLearner在多个视频分类数据集基准上取得了SOTA的性能。

其中，在Kinetics-400和Kinetics-600上的性能超过了以前的Transformer模型，在Charades和AViD上也超过了之前的CNN模型。

通过与视频视觉Transformer（Video Vision Transformer，ViViT）结合，TokenLearner会在每个时间段学习8（或16）个token。

左：视频分类任务；右图：不同模型的对比

随着时间的推移，当人物在场景中移动时，TokenLearner会注意到不同的空间位置变化从而进行token化。

TokenLearner的空间注意力图的可视化

结论

虽然Vision Transformer是计算机视觉领域的一个强大模型，但大量的token及庞大的计算量一直是将ViT应用于更大图像和更长视频的瓶颈。

本文中作者表明，保留如此大量的token并在整个层集上完全处理它们是没有必要的。

此外，作者还证明了通过学习一个基于输入图像自适应提取token的模块，可以在节省计算的同时获得更好的性能。

最后，多个公共数据集上的验证也表明了TokenLearner在视频表征学习任务中的表现十分优异。