通过学习令牌化提高视觉 Transformer 的效率和准确率

会员服务 ·

通过学习令牌化提高视觉 Transformer 的效率和准确率

2022 年 1 月 19 日 TensorFlow

发布人：Google 机器人团队研究员 Michael Ryoo 和 Google Research 研究员 Anurag Arnab

Transformer模型在计算机视觉任务（包括目标检测和视频分类）领域不断取得领先成果。

Transformer 模型
https://ai.googleblog.com/2017/08/transformer-novel-neural-network.html
目标检测
https://arxiv.org/pdf/2103.14030v2.pdf
视频分类
https://arxiv.org/pdf/2103.15691.pdf

不同于逐像素处理图像的标准卷积方法，视觉 Transformer (ViT) 将图像视为一系列的块 token（即由多个像素组成的图像的较小部分或“块”）。这意味着在每一层，ViT 模型都会使用多头自注意力机制，根据每对 token 之间的关系重新组合和处理块 token。这样，ViT 模型就能够对整个图像构建全局表征。

卷积方法
http://deeplearning.csail.mit.edu/cvpr2017_tutorial_kaiminghe.pdf
多头自注意力机制
https://arxiv.org/abs/1706.03762

在输入级，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：8 个学习的 Token 可以为图像和视频做什么？
https://arxiv.org/pdf/2106.11297.pdf

NeurIPS 2021
https://nips.cc/Conferences/2021

TokenLearner

我们使用一种简单的空间注意力方法来实现 TokenLearner。为了生成每个学习的 token，我们计算了一个空间注意力图，突出显示重要区域（使用卷积层或 MLP）。然后将这个空间注意力图应用到输入中，来为每个区域进行不同的加权（忽略不必要的区域），并将结果空间池化，来生成最终学习的 token。这个过程会并行重复多次，从而从原始输入中得出少量（10 个左右）token。这也可视为基于权重值，对像素进行软选择，随后进行全局平均池化。请注意，用于计算注意力图的函数由不同的可学习参数集控制，并以端到端的方式进行训练。这让注意力函数可以在捕获输入中不同空间信息时进行优化。下图说明了该过程。

TokenLearner 模块学习为每个输出 token 生成空间注意力图，并使用它来抽象输入以进行令牌化。在实践中，模型将学习多个空间注意力函数，将其应用到输入，并行生成不同的 token 向量

因此，TokenLearner 能让模型处理与特定识别任务相关的少量 token，而不是处理固定的、均匀地令牌化的输入。也就是说，(1) 我们启用了自适应令牌化，这样可以根据输入动态地选择 token，(2) 这一做法有效地减少了 token 的总数，显著减少了网络执行的计算量。这些动态自适应生成的 token 可用于标准的 Transformer 架构，例如针对图像的 ViT 和针对视频的 ViViT。

ViT
https://arxiv.org/pdf/2010.11929.pdf

TokenLearner 放在哪里

在构建 TokenLearner 模块后，我们必须确定将其放置在哪个位置。首先，我们尝试将其放置在标准 ViT 架构内的不同位置，输入 224x224 大小的图像。TokenLearner 生成的 token 数量是 8 和 16，远远少于标准 ViT 使用的 196 或 576 个 token。下图显示了在 ViT B/16 的不同相对位置插入 TokenLearner 时，模型的 ImageNet 小样本分类准确率和 FLOPS，其中 ViT B/16 是有 12 个注意力层的基础模型，在 16x16 块 token 上运行。

ImageNet
https://image-net.org/

上图：进行 JFT 300M 预训练后，TokenLearner 处于 ViT B/16 内的不同相对位置时，对应的ImageNet 5 样本转移准确率。位置 0 意味着 TokenLearner 放置在任何 Transformer 层之前。Base 指原始 ViT B/16。

下图：在每个 TokenLearner 相对位置，以数十亿浮点算子 (GFLOPS) 衡量的计算量

JFT 300M
https://arxiv.org/pdf/1707.02968v2.pdf

我们发现，在网络的前四分之一处之后（含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 模型性能相当。

ViT G/14 模型
https://arxiv.org/pdf/2106.04560v1.pdf

左图：大规模 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 在视频表征学习任务中特别有效。我们这项研究的预印本以及代码已公开。