Transformer 模型在计算机视觉任务（包括物体检测和视频分类）中始终取得最佳效果。与逐像素处理图像的标准卷积方法不同， Vision Transformers (ViT) 将图像视为一系列补丁标记（即由多个像素组成的图像的较小部分或“补丁”）。这意味着在每一层，ViT 模型都会使用多头自注意力机制，根据每对标记之间的关系重新组合和处理补丁标记。通过这样做，ViT 模型能够构建整个图像的全局表示。
在输入层，标记是通过将图像均匀分割成多个片段而形成的，例如，将 512 x 512 像素的图像分割成 16 x 16 像素的块。在中间层，上一层的输出成为下一层的标记。对于视频，视频“小管”（例如 16x16x2 视频片段（2 帧上的 16x16 图像））将成为标记。视觉标记的质量和数量决定了 Vision Transformer 的整体质量。
许多 Vision Transformer 架构面临的主要挑战是，它们通常需要太多标记才能获得合理的结果。例如，即使使用 16x16 块标记化，单个 512x512 图像也对应 1024 个标记。对于具有多帧的视频，这会导致每一层都需要处理数万个标记。考虑到 Transformer 计算量会随着标记数量的二次方增加，这通常会使 Transformer 难以处理更大的图像和更长的视频。这引出了一个问题：真的有必要在每一层处理那么多标记吗？
在“ TokenLearner：8 个学习到的 token 能为图像和视频做什么？ ”（其早期版本已在NeurIPS 2021上展示）中，我们展示了自适应地生成较少数量的 token，而不是始终依赖于均匀分割形成的 token，可以让 Vision Transformers 运行得更快、性能更好。TokenLearner 是一个可学习的模块，它接受类似图像的张量（即输入）并生成一小组 token。此模块可以放置在相关模型内的不同位置，从而显著减少所有后续层中要处理的 token 数量。实验表明，使用 TokenLearner 可以在不损害分类性能的情况下节省一半或更多的内存和计算量，而且由于它能够适应输入，它甚至可以提高准确率。
TokenLearner
我们使用一种简单的空间注意力方法实现了 TokenLearner。为了生成每个学习到的 token，我们计算一个突出显示重要区域的空间注意力图（使用卷积层或MLP）。然后将这种空间注意力图应用于输入，以对每个区域赋予不同的权重（并丢弃不必要的区域），并将结果进行空间池化以生成最终的学习到的 token。此过程并行重复多次，从而从原始输入中产生几个（~10）token。这也可以看作是基于权重值对像素进行软选择，然后进行全局平均池化。请注意，计算注意力图的函数由不同的可学习参数集控制，并以端到端的方式进行训练。这允许优化注意力函数以捕获输入中的不同空间信息。下图说明了该过程。
TokenLearner 模块学习为每个输出 token 生成一个空间注意力图，并使用它来抽象要进行 token 化的输入。实际上，会学习多个空间注意力函数，将其应用于输入，并并行生成不同的 token 向量。
因此，TokenLearner 不再处理固定的、统一标记化的输入，而是使模型能够处理与特定识别任务相关的较少数量的标记。也就是说，(1) 我们启用自适应标记化，以便可以根据输入动态选择标记；(2) 这有效地减少了标记的总数，大大减少了网络执行的计算。这些动态和自适应生成的标记可用于标准转换器架构，例如用于图像的ViT和用于视频的 ViViT 。
TokenLearner 的放置位置
构建 TokenLearner 模块后，我们必须确定将其放置在何处。我们首先尝试将其放置在具有 224x224 图像的标准 ViT 架构中的不同位置。TokenLearner 生成的标记数量为 8 和 16，远少于标准 ViT 使用的 196 或 576 个标记。下图显示了在 ViT B/16 中的各个相对位置插入 TokenLearner 的模型的ImageNet小样本分类准确率和FLOPS，ViT B/16 是具有 12 个注意层的基础模型，在 16x16 块标记上运行。
顶部：使用JFT 300M预训练的ImageNet 5 次迁移准确率，相对于 ViT B/16 内的相对 TokenLearner 位置。位置 0 表示 TokenLearner 位于任何 Transformer 层之前。底部是原始的 ViT B/16。底部：计算量，以每个相对 TokenLearner 位置的数十亿次浮点运算 (GFLOPS) 来衡量。
我们发现，在网络的初始四分之一之后（1/4 处）插入 TokenLearner 可实现与基线几乎相同的准确率，同时将计算量减少到基线的三分之一以下。此外，与不使用 TokenLearner 相比，将 TokenLearner 放置在较后的层（3/4 处之后）可实现更好的性能，同时由于其自适应性，执行速度更快。由于 TokenLearner 之前和之后的 token 数量差异很大（例如，之前为 196 个，之后为 8 个），TokenLearner 模块之后的 transformer 的相对计算量几乎可以忽略不计。
与 ViT 比较
我们在 ImageNet 小样本迁移上按照相同设置比较了带有 TokenLearner 的标准 ViT 模型和不带有 TokenLearner 的标准 ViT 模型。TokenLearner 被放置在每个 ViT 模型中间的不同位置，例如 1/2 和 3/4。下图显示了带有和不带有 TokenLearner 的模型的性能/计算权衡。
带有和不带有 TokenLearner 的 ViT 模型在 ImageNet 分类上的表现。这些模型使用 JFT 300M 进行预训练。模型越靠近每个图的左上角越好，这意味着它运行得更快，性能更好。观察 TokenLearner 模型在准确性和计算方面的表现如何优于 ViT。
我们还将 TokenLearner 插入到更大的 ViT 模型中，并将其与巨型ViT G/14 模型进行比较。在这里，我们将 TokenLearner 应用于 ViT L/10 和 L/8，它们是具有 24 个注意层的 ViT 模型，以 10x10（或 8x8）块作为初始标记。下图显示，尽管使用更少的参数和更少的计算，TokenLearner 的性能与具有 48 层的巨型 G/14 模型相当。
左图：大规模 TokenLearner 模型与 ViT G/14 在 ImageNet 数据集上的分类准确率对比。右图：参数数量与 FLOPS 对比。
高性能视频模型
视频理解是计算机视觉领域的关键挑战之一，因此我们在多个视频分类数据集上评估了 TokenLearner。这是通过将 TokenLearner 添加到Video Vision Transformers (ViViT) 中实现的，后者可以被视为 ViT 的时空版本。TokenLearner 每个时间步学习 8 个（或 16 个）标记。
与 ViViT 结合使用时，TokenLearner 在多个流行视频基准测试（包括 Kinetics-400、Kinetics-600、Charades 和 AViD）上获得了最先进 (SOTA) 的性能，超越了 Kinetics-400 和 Kinetics-600 上之前的 Transformer 模型以及 Charades 和 AViD 上之前的 CNN 模型。
使用 TokenLearner 的模型在流行视频基准测试中的表现优于最先进的模型（拍摄于 2021 年 11 月）。左图：流行的视频分类任务。右图：与 ViViT 模型的比较。
TokenLearner 中空间注意力图随时间变化的可视化。当人物在场景中移动时，TokenLearner 会关注不同的空间位置以进行标记。
结论
虽然 Vision Transformers 是强大的计算机视觉模型，但大量的 token 及其相关的计算量一直是其应用于更大图像和更长视频的瓶颈。在这个项目中，我们说明了保留如此大量的 token 并在整个层上完全处理它们是不必要的。此外，我们证明，通过学习基于输入图像自适应提取 token 的模块，可以获得更好的性能，同时节省计算量。我们提出的 TokenLearner 在视频表示学习任务中特别有效，我们用多个公共数据集证实了这一点。我们的工作的预印本以及代码都是公开的。
致谢
我们感谢我们的合著者：AJ Piergiovanni、Mostafa Dehghani 和 Anelia Angelova。我们还感谢 Google 机器人团队成员的积极讨论。
评论