[论文解读] TokenLearner: What Can 8 Learned Tokens Do for Images and Videos?

Michael S. Ryoo, AJ Piergiovanni+3 more|arXiv (Cornell University)|Jun 21, 2021

Human Pose and Action Recognition参考文献 47被引用 12

一句话总结

TokenLearner 引入了一个可学习模块，可从图像和视频中自适应地选择 8–16 个重要视觉标记，取代 Vision Transformers 中固定的 patch 基 token 化方式。通过使用空间注意力机制聚焦显著区域，并减少每层处理的标记数量，该方法在 ImageNet 和多个视频基准测试中实现了最先进性能，同时将计算量和显存使用量减少一半以上。

ABSTRACT

In this paper, we introduce a novel visual representation learning which relies on a handful of adaptively learned tokens, and which is applicable to both image and video understanding tasks. Instead of relying on hand-designed splitting strategies to obtain visual tokens and processing a large number of densely sampled patches for attention, our approach learns to mine important tokens in visual data. This results in efficiently and effectively finding a few important visual tokens and enables modeling of pairwise attention between such tokens, over a longer temporal horizon for videos, or the spatial content in images. Our experiments demonstrate strong performance on several challenging benchmarks for both image and video recognition tasks. Importantly, due to our tokens being adaptive, we accomplish competitive results at significantly reduced compute amount. We obtain comparable results to the state-of-the-arts on ImageNet while being computationally more efficient. We also confirm the effectiveness of the approach on multiple video datasets, including Kinetics-400, Kinetics-600, Charades, and AViD. The code is available at: https://github.com/google-research/scenic/tree/main/scenic/projects/token_learner

研究动机与目标

• 为解决 Vision Transformers 处理数千个固定 patch 或 tubelets 所带来的高计算成本问题。
• 通过学习一组少量代表性标记而非均匀采样 patch，提升图像和视频识别的效率与准确率。
• 在计算开销最小化的情况下，实现对长距离空间和时间依赖关系的有效建模。
• 证明仅使用 8–16 个学习得到的标记即可匹配甚至超越使用数百个固定标记的模型性能。

提出的方法

• TokenLearner 使用可学习的空间注意力机制识别输入特征中的重要区域，取代固定的 patch 基 token 化方式。
• 通过卷积层或 MLP 计算注意力图，突出显示特征图中的显著空间区域。
• 将注意力图应用于输入特征图，以自适应地加权区域，随后通过空间平均池化生成一组少量学习得到的标记。
• 生成的标记被送入标准的 Transformer 编码器，通过向量化自注意力机制进行处理，同时保留通道信息。
• 该模块可插入 Vision Transformer 的任意层，支持灵活集成，并在全网络范围内实现计算节省。
• 该方法通过在空间和时空特征图上应用相同机制，同时支持图像和视频输入。

实验结果

研究问题

• RQ1 少量自适应学习的标记是否能在不损失性能的前提下替代 Vision Transformers 中密集的 patch token 化？
• RQ2 自适应标记选择是否能在图像和视频识别任务中提升准确率并降低计算成本？
• RQ3 在 FLOPS 和准确率方面，TokenLearner 与固定 patch 方法及其他最先进模型相比表现如何？
• RQ4 为最大化效率与性能，学习标记的最佳数量（如 8 与 16）和放置位置是什么？

主要发现

• TokenLearner 在实现与最先进模型相当的 ImageNet top-1 准确率的同时，将 FLOPS 减少超过 50%。
• 在 Kinetics-400、Kinetics-600、Charades 和 AViD 数据集上，基于 TokenLearner 的模型仅使用每帧 8–16 个学习标记，即创下新的最先进结果。
• 在网络早期插入 TokenLearner（例如在第 2 或第 3 层）可带来更高的准确率和更大的效率提升。
• 消融实验表明，向量化 Transformer 公式化与学习标记的结合优于标准多头自注意力（MHSA），在 Charades 数据集上达到 59.6% 的准确率。
• TokenLearner、向量化注意力与 TokenFuser 模块的组合实现了最佳性能（Charades 数据集上为 59.6%），证明了完整流水线的有效性。