发布人：Google 机器人团队研究员 Michael Ryoo 和 Google Research 研究员 Anurag Arnab
Transformer 模型 在计算机视觉任务（包括 目标检测 和 视频分类）领域不断取得领先成果。
不同于逐像素处理图像的标准 卷积方法，视觉 Transformer (ViT) 将图像视为一系列的块 token（即由多个像素组成的图像的较小部分或“块”）。这意味着在每一层，ViT 模型都会使用 多头自注意力机制，根据每对 token 之间的关系重新组合和处理块 token。这样，ViT 模型就能够对整个图像构建全局表征。
在输入级，token 是通过将图像均匀地分割为多个部分来形成的，例如，将 512x512 像素的图像分割成多个 16x16 像素的块。在中间级，上一层的输出变成下一层的 token。对于视频，视频“管道”会变成 token，例如 16x16x2 的视频片段（2 帧的 16x16 图像）。视觉 token 的质量和数量决定了视觉 Transformer 的整体质量。
许多视觉 Transformer 架构面临的主要挑战是，它们通常需要过多的 token 才能获取合理的结果。例如，即使使用 16x16 的块令牌化，一个 512x512 的图像也会产生 1024 个 token。对于具有多帧的视频，这会导致每一层都需要处理数以万计的 token。考虑到 Transformer 计算量随着 token 的数量增加而呈指数级增长，这通常会使 Transformer 难以处理更大的图像和更长的视频。这就引出了一个问题：真的有必要在每一层处理那么多的 token 吗？
在 “TokenLearner：8 个学习的 Token 可以为图像和视频做什么？ (TokenLearner: What Can 8 Learned Tokens Do for Images and Videos?)”（其早期版本已在 NeurIPS 2021 上发表）一文中，我们展示了通过 自适应 方式生成较少数量的 token，而不是全部依赖于均匀分割形成的 token，可以提升视觉 Transformer 的运行速度和性能。TokenLearner 是一个可学习的模块，它会获取类似图像的张量（即输入）并生成一小组 token。此模块可以放置在目标模型内的不同位置，可显著减少所有后续层中要处理的 token 数量。实验表明，TokenLearner 可以在不损害分类性能的情况下，节省至少一半的内存和计算量，且因为其适应输入的能力，甚至提高了准确率。
TokenLearner
我们使用一种简单的空间注意力方法来实现 TokenLearner。为了生成每个学习的 token，我们计算了一个空间注意力图，突出显示重要区域（使用卷积层或 MLP）。然后将这个空间注意力图应用到输入中，来为每个区域进行不同的加权（忽略不必要的区域），并将结果空间池化，来生成最终学习的 token。这个过程会并行重复多次，从而从原始输入中得出少量（10 个左右）token。这也可视为基于权重值，对像素进行软选择，随后进行全局平均池化。请注意，用于计算注意力图的函数由不同的可学习参数集控制，并以端到端的方式进行训练。这让注意力函数可以在捕获输入中不同空间信息时进行优化。下图说明了该过程。
TokenLearner 模块学习为每个输出 token 生成空间注意力图，并使用它来抽象输入以进行令牌化。在实践中，模型将学习多个空间注意力函数，将其应用到输入，并行生成不同的 token 向量
因此，TokenLearner 能让模型处理与特定识别任务相关的少量 token，而不是处理固定的、均匀地令牌化的输入。也就是说，(1) 我们启用了自适应令牌化，这样可以根据输入动态地选择 token，(2) 这一做法有效地减少了 token 的总数，显著减少了网络执行的计算量。这些动态自适应生成的 token 可用于标准的 Transformer 架构，例如针对图像的 ViT 和针对视频的 ViViT。
TokenLearner 放在哪里
在构建 TokenLearner 模块后，我们必须确定将其放置在哪个位置。首先，我们尝试将其放置在标准 ViT 架构内的不同位置，输入 224x224 大小的图像。TokenLearner 生成的 token 数量是 8 和 16，远远少于标准 ViT 使用的 196 或 576 个 token。下图显示了在 ViT B/16 的不同相对位置插入 TokenLearner 时，模型的 ImageNet 小样本分类准确率和 FLOPS，其中 ViT B/16 是有 12 个注意力层的基础模型，在 16x16 块 token 上运行。
上图： 进行 JFT 300M 预训练后，TokenLearner 处于 ViT B/16 内的不同相对位置时，对应的ImageNet 5 样本转移准确率。位置 0 意味着 TokenLearner 放置在任何 Transformer 层之前。Base 指原始 ViT B/16。
下图： 在每个 TokenLearner 相对位置，以数十亿浮点算子 (GFLOPS) 衡量的计算量
我们发现，在网络的前四分之一处之后（含1/4 处）插入 TokenLearner，可以获得与基线几乎相当的准确率，同时将计算量减少到基线的三分之一以下。此外，将 TokenLearner 放在稍后的层（网络的 3/4 处后），与不使用 TokenLearner 相比，也能获得更好的性能，同时由于其适应性，运行速度更快。由于 TokenLearner 前后 token 数量差别巨大（例如，之前 196 个，之后 8 个），TokenLearner 模块之后的 Transformer 相对计算量几乎可以忽略不计。
与 ViT 的比较
我们将带有 TokenLearner 的标准 ViT 模型和普通 ViT 模型进行对比，同时在 ImageNet 小样本转移上采用相同的设置。我们将 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）块作为初始 token。下图显示，尽管使用更少的参数和计算量，TokenLearner 的性能与具有 48 层的巨型 G/14 模型性能相当。
左图： 大规模 TokenLearner 模型与 ViT G/14 在 ImageNet 数据集上的分类准确率比较。右图： 参数数量和 FLOPS 的比较
高性能视频模型
视频理解是计算机视觉中的关键挑战之一，因此我们在多个视频分类数据集上评估了 TokenLearner。我们将 TokenLearner 添加到视频视觉 Transformer (ViViT) 中，ViViT 可视为 ViT 的时空版本。TokenLearner 每个时步学习 8（或 16）个 token。
通过与 ViViT 结合，TokenLearner 在多个流行的视频基准测试（包括 Kinetics-400、Kinetics-600、Charades 和 AViD）中获得了最高水平 (state-of-the-art，SOTA) 的性能，在 Kinetics-400 和 Kinetics-600 上优于先前的 Transformer 模型，在 Charades 和 AViD 上优于先前的 CNN 模型。
使用 TokenLearner 的模型在流行的视频基准测试上优于先前最先进的模型（2021 年 11 月采集）。左： 流行的视频分类任务。右： 与 ViViT 模型的比较
随着时间的推移，TokenLearner 中空间注意力图的可视化。当人在场景中移动时，TokenLearner 会关注不同的空间位置，以进行令牌化
结论
虽然视觉 Transformer 是计算机视觉领域的强大模型，但大量的 token 和庞大的计算量一直是将其应用于更大图像和更长视频的瓶颈。在此项目中，我们展示了在整个层集中保留如此大量的 token 并进行完全处理是没有必要的。并且，我们证明了通过学习基于输入图像自适应地提取 token 的模块，可以在节省计算量的同时获得更好的性能。通过在多个公开数据集上取得的结果，证明我们提出的 TokenLearner 在视频表征学习任务中特别有效。我们这项研究的预印本以及 代码 已公开。
致谢
感谢以下合著者：AJ Piergiovanni、Mostafa Dehghani 和 Anelia Angelova。我们还要感谢 Google 机器人团队成员的积极讨论。
原文：Improving Vision Transformer Efficiency and Accuracy by Learning to Tokenize
中文：TensorFlow 公众号