[论文解读] ClusTR: Exploring Efficient Self-attention via Clustering for Vision Transformers

Yutong Xie, Jianpeng Zhang+3 more|arXiv (Cornell University)|Aug 28, 2022

Multimodal Machine Learning Applications被引用 6

一句话总结

ClusTR 提出了一种基于内容的稀疏自注意力机制，用于视觉变换器（Vision Transformers），通过基于语义相似性的聚类与聚合键（key）和值（value）令牌，降低计算复杂度，实现高效的长距离建模。该方法仅使用 22.7M 参数即在 ImageNet 上实现了 83.2% 的 Top-1 准确率，性能优于密集注意力和基于网格的稀疏注意力方法。

ABSTRACT

Although Transformers have successfully transitioned from their language modelling origins to image-based applications, their quadratic computational complexity remains a challenge, particularly for dense prediction. In this paper we propose a content-based sparse attention method, as an alternative to dense self-attention, aiming to reduce the computation complexity while retaining the ability to model long-range dependencies. Specifically, we cluster and then aggregate key and value tokens, as a content-based method of reducing the total token count. The resulting clustered-token sequence retains the semantic diversity of the original signal, but can be processed at a lower computational cost. Besides, we further extend the clustering-guided attention from single-scale to multi-scale, which is conducive to dense prediction tasks. We label the proposed Transformer architecture ClusTR, and demonstrate that it achieves state-of-the-art performance on various vision tasks but at lower computational cost and with fewer parameters. For instance, our ClusTR small model with 22.7M parameters achieves 83.2\% Top-1 accuracy on ImageNet. Source code and ImageNet models will be made publicly available.

研究动机与目标

• 为解决视觉变换器中密集自注意力机制的二次方计算复杂度问题，尤其是在密集预测任务中的应用。
• 克服基于网格的下采样方法的局限性，如丢失细粒度细节以及引入冗余的背景令牌。
• 开发一种内容感知的令牌聚合方法，在减少序列长度的同时保留语义丰富性。
• 将基于聚类的注意力机制扩展至多尺度建模，以提升在密集预测任务中的性能。
• 在减少模型大小和计算成本的前提下，实现最先进（SOTA）的性能。

提出的方法

• 基于键和值令牌的内容相似性进行聚类，通过聚合为聚类代表令牌来减少令牌数量。
• 查询令牌关注聚类后的键和值令牌，形成一种基于聚类引导的自注意力机制，从而保持长距离依赖建模能力。
• 聚类数量可灵活调节，实现对序列长度和计算成本的动态控制，优于固定网格下采样的方式。
• 通过在不同阶段应用不同的聚类数量，引入多尺度聚类，实现分层特征学习。
• 将该方法整合进一种新型变换器架构 ClusTR 中，提供小（tiny）、基础（base）和小号（small）三种变体。
• 在图像分类、语义分割、目标检测和姿态估计等基准数据集上进行训练与评估。

实验结果

研究问题

• RQ1 基于内容的键和值聚类是否能在不损失性能的前提下降低视觉变换器的计算复杂度？
• RQ2 与基于网格的或卷积下采样相比，基于聚类引导的自注意力在保留语义信息和准确率方面表现如何？
• RQ3 多尺度聚类引导的自注意力是否能提升语义分割等密集预测任务的性能？
• RQ4 基于聚类的稀疏注意力机制是否能在更少参数和浮点运算次数（FLOPs）下实现最先进（SOTA）的结果？
• RQ5 聚类数量的灵活性对不同视觉任务中模型效率与准确率的影响如何？

主要发现

• ClusTR-S 模型仅使用 22.7M 参数即在 ImageNet 上达到 83.2% 的 Top-1 准确率，优于同类模型且浮点运算次数（FLOPs）和参数量更低。
• 与基于网格的令牌聚合方法相比，该聚类引导方法将参数量减少 18%，Top-1 准确率提升 0.5 个百分点（77.2% vs. 76.7%）。
• 与 k-NN 和空间下采样注意力相比，ClusTR 在参数量减少 18%、FLOPs 减少 5% 的同时，Top-1 准确率高出 0.4%。
• 尽管浮点运算次数（FLOPs）和参数量略有增加，多尺度聚类在 ADE-20K 语义分割任务上的 mIOU 提升 1.4 个百分点（42.6 vs. 41.2），优于单尺度方法。
• 当参数量为 40.3M 时，模型在 ImageNet 上达到 84.1% 的 Top-1 准确率，展现出强大的可扩展性。
• 消融实验证实，聚类引导注意力在准确率和效率方面均优于基于网格的和 k-NN 的稀疏注意力方法。