CVPR 2022 Oral | TCFormer：基于动态token聚类的transformer

近年来，Transformer网络被广泛应用到了计算机视觉领域，并获得了许多惊艳的效果。Transformer网络在处理图片时，往往需要将图片转换为一组视觉令牌（vision token）。如下图（a）所示，现有的方法大多使用基于网格的划分方式，将图片转化为视觉令牌。即：将图片划分为标准的网格区域，且每一个网格用一个令牌表示。

然而，基于网格的划分方式忽视了图片的语义信息。基于网格的划分方式仅仅根据空间关系产生视觉令牌，单个令牌可能需要代表具有不同语义信息的多个物体，不利于令牌特征的提取。我们认为应当根据图片语义信息划分令牌，将具有相近语义信息的图片区域用同一个令牌表示。

另外，基于网格的划分方式也忽视了不同区域之间的差异。比如，在以人体为中心的视觉任务中，人体区域往往比背景区域更加重要，应当使用更多的令牌来表示。并且，一些区域（如人脸、人手区域）比其他区域（如躯干区域）包含更多的细节信息，应当采用更精细的令牌表示。

基于上述观察，我们提出了一种基于特征聚类的视觉令牌生成方式，能够动态地调整视觉令牌的形状、大小和位置，从而更好地提取图像特征。

下图展示了TCFormer的整体框架。TCFormer采用了多阶段（stage）的网络框架，共由4个不同的阶段组成，每个阶段又由多个串联的transformer模块组成。将图片输入TCFormer后，我们首先通过卷积层提取图像特征图，并将特征图中的每一个像素当作一个视觉令牌，从而得到初始的视觉令牌。在相邻的两个阶段之间，我们通过令牌融合的方式减少令牌数目，并增加令牌特征的通道数。最后，我们通过一个多阶段令牌特征聚合模块（MTA）融合各个阶段的令牌特征。

Transformer模块

下图展示了 TCFormer中使用的transformer模块。在标准的transformer模块的基础上，我们仿照PVT[1] 加入了空间缩减模块。该模块通过卷积层，减少了作为key、value的令牌的数目，从而减少了注意力模块的计算复杂度。另外，我们受到 CVT[2]等工作的启发，在transformer模块中加入了逐通道的卷积层，从而更好地抽取局部图像特征。

基于特征聚类的令牌融合模块 （Clustering-based Token Merge, CTM）

基于特征聚类的令牌融合模块（CTM）是TCFormer的核心模块，我们在下图展示了它的结构。输入前一阶段的视觉令牌后，CTM首先根据令牌特征进行聚类。随后，CTM根据聚类结果，将同一类的视觉令牌融合成一个新的视觉令牌，从而减少视觉令牌的数目。融合后的视觉令牌的特征是融合前令牌特征的加权平均，视觉令牌的对应区域则是融合前令牌对应区域的并集。最后，我们将融合后的视觉令牌作为query，将融合前的视觉令牌作为key和value，输入到一个transformer模块中，进一步加强从融合前令牌到融合后令牌的信息传递。

我们首先在COCO-WholeBody[4]这一数据集上测试了TCFormer在人体全身关键点检测任务上的效果。如下表所示，TCFormer取得了最先进的效果，并在如人手、人脸关键点的细节重建精度上，取得了明显的提升。这一结果证明了TCFormer在图像细节信息提取上的优势。

在视觉transformer领域，注意力机制得到了广泛的研究，而视觉令牌本身却并没有受到足够的关注。本文提出了一种transformer网络——TCFormer。TCFormer采用了一种基于特征聚类的动态令牌生成方式，和一种基于transformer的特征聚合模块。TCFormer在基于图像的人体全身关键点估计、人脸关键点估计和人体三维网格重建任务上都获得了最先进的效果，并在人体细节的重建精度上取得了明显的提升。我们希望这项工作能够启发后续的相关研究。