论文阅读|TCFormer

Not All Tokens Are Equal: Human-centric Visual Analysis via Token Clustering Transformer

Abstract

大多数的vision transformer将图片分割成规则和固定的网格并且将每一个cell都作为token来产生visual tokens，然而，在以人为中心的视觉任务中，并非所有区域都同等重要，例如，人体需要许多tokens来精细表示，而图像背景可以通过几个tokens来建模。为了解决这个问题，提出了Token Clustering Transformer (TCFormer),通过渐进式聚类progressive clustering来合并tokens，其中tokens可以从具有灵活形状和大小的不同位置合并。TCFormer中的token不仅可以聚焦重要区域，还可以调整token形状以适应语义概念，并对包含关键细节的区域采用精细的分辨率，有利于捕获详细信息。 大量实验表明，TCFormer 在不同的具有挑战性的以人为中心的任务和数据集上始终优于同类网络，包括 COCO-WholeBody 上的全身姿态估计和 3DPW 上的 3D 人体网格重建。

Introduction

由于传统的transformer使用一系列tokens作为输入，大多数现有的vision transformer都遵循这种范式，将输入图像划分为规则且固定的网格，其中每个单元（图像块）被视为一个token，如图1a所示。基于网格的token生成方法简单，在图像识别、目标检测和分割等计算机视觉任务[Swin,PVT]中取得了巨大的成功。

*图1：Comparisons between vision tokens generated by (a) standard grids and (b) TCFormer. 不同token的token区域，或者visual token所代表的图像区域，用不同的颜色进行可视化。 从左到右，不同的图像代表不同的阶段。 Grid-based tokens对所有区域一视同仁，如 (a) 所示。 而 (b) 中的tokens动态地处理图像区域。 tokens在人体上分布更密集。 对于背景区域，大区域由单个标记（蓝色）表示，而对于包含重要细节的区域，例如面部区域，则使用具有精细空间大小的标记（红色）。

然而，基于固定网格的visual tokens对于以人为中心的视觉分析来说不是最优的。 在以人为中心的视觉分析中，人体的图像区域比图像背景更重要，这促使我们用动态形状和大小的视觉标记来表示不同的图像区域。但是基于网格的视觉token的token区域是具有固定位置、形状和大小的矩形区域。统一视觉token分配无法将更多tokens分配到重要区域。

为了解决这个问题，我们提出了一种新颖的vision transformer，Token Clustering Transformer（TCFormer），它通过渐进式令牌聚类progressive token clustering生成tokens。 TCFormer在每个阶段动态生成tokens。 如图 1 (b) 所示，它能够生成具有各种位置、大小和形状的tokens。 首先，与基于网格的tokens不同，聚类后的tokens不限于规则形状，可以集中在重要区域，例如人体。 其次，TCFormer动态生成具有适当大小的tokens来表示不同的区域。 对于人脸等充满重要细节的区域，分配更精细的token。 相比之下，单个标记（例如，图 1 (b) 中的蓝色标记）用于表示背景的大面积区域。

在 TCFormer 中，特征图中的每个像素在第一阶段都被初始化为一个visual token，其token区域是该像素所覆盖的区域。 我们逐步合并具有相似语义的tokens，并在不同阶段获得不同数量的tokens。 为此，我们精心设计了一个Clustering Token Merge (CTM) block。 首先，给定前一阶段的tokens，CTM 通过对token features应用基于 k 近邻的密度峰值聚类算法 [11] 对它们进行分组。 其次，分配给同一集群的tokens通过对token features进行平均来合并为单个token。 最后，tokens被输入到一个transformer block中以进行特征聚合。 合并tokens的token区域是输入tokens区域的并集。

多阶段特征的聚合被证明有利于以人为中心的分析 [HRNet,HRFormer]。 大多数先前的工作 [Swin,PVT,HRFormer] 将visual tokens转换为特征图并以特征图的形式聚合features。 然而，当我们将动态visual tokens转换为特征图时，多个tokens可能位于同一个像素网格中，导致细节丢失。 为了解决这个问题，我们提出了一个 Multi-stage Token Aggregation (MTA) head，它能够以有效的方式保留所有阶段的图像细节。 具体来说，MTA head 从最后一个阶段的tokens开始，然后逐步对tokens进行上采样并聚合前一阶段的token特征，直到所有阶段的特征都被聚合。 聚合后的tokens与特征图中的像素一一对应，并被重新reshape为特征图以供后续处理。

Contributions

• 我们提出了 Token Clustering Transformer (TCFormer)，通过渐进式群集和合并tokens为每个图像生成不同位置、大小和形状的visual tokens。据我们所知，这是第一次将clustering集群用于动态token生成。
• 我们提出了一种 Multi-stage Token Aggregation(MTA) head来聚合多个阶段的token features，有效地保留了各个阶段的详细信息
• 大量的实验表明，TCFormer在不同具有挑战性的以人类为中心的任务和数据集上都表现出了一致的性能，包括在Coco-WholeBody上的全身姿势估计和在3DPW上的3D人体网格重建。

Related works

2.1Transformers in Human-Centric Vision Tasks

之前的工作基本分为两类，细化图像特征和聚合局部特征。第一种方法利用transformer来提取更好的图像特征。例如，Transpose[69]使用额外的transformer encoder layer来优化CNN提取的图像特征。HRFormer[74]将transformer block应用于HRNet[50]结构。我们提出的TCFormer也属于第一种。

第二种方法是利用transformer来聚合不同部分的特征。例如，TFPose和TokenPose使用堆叠的transformer来生成基于关键点的特征，根据这些特征分别通过回归坐标和热图来预测关键点位置。Mesh Graphormer设计了一种带有transformer和图形卷积层的head。head从图像特征图中聚合关节和顶点特征，并预测人体关节和网格顶点的3D位置。

这两种类型的大多数已有工作生成的vision tokens都带有固定的网格标记区域，这对于以人为中心的任务来说是次优的。相比之下，TCFormer的令牌区域是自动学习的，而不是使用手动设计。学习的token区域根据语义在位置、形状和大小上都是灵活的。

2.2 Dynamic Token Generation

DynamicViT[48]和PnP-DETR[59]通过预测token级分数来获取重要的tokens。对于不重要的tokens，DynamicViT简单地丢弃它们，而PNP-DETR用稀疏的特征向量来表示它们。DVT[61]根据分类难度的不同，用不同分辨率的网格构建visual tokens。PS-VIT[75]构建具有固定网格大小的tokens，并在处理过程中逐步调整网格中心。

上面提到的方法都是基于网格的token生成的变体。它们专门修改网格的数量、分辨率和中心。相比之下，我们提出的TCFormer的token区域不受网格结构的限制，在位置、形状和大小三个方面都更加灵活。首先，我们基于语义相似度而不是空间贴近度将图像区域分配给特定的visual tokens。其次，token的图像区域不限于矩形。具有相同语义的区域，即使它们是不相邻的，也可以由单个tokens来表示。第三，同一阶段的visual tokens可能有不同的大小。例如，在图1(B)中，用于背景的token呈现大区域，而用于人脸的token仅呈现小区域。这样的属性有助于在TCFormer的所有阶段中保留重要的细节。

2.3Clustering for Feature Aggregation

基于聚类的特征聚集方法在点云[47]和图表示[71]的研究领域得到了广泛的探索。PointNet++[47]通过最远的点采样对点云进行下采样，然后聚合k个最近邻点的要素。DIFFPOOL[71]预测输入图的软聚类分配矩阵，该矩阵用于分层粗化该图。

这些方法是专门针对点云和图形数据设计的，不能直接应用于基于图像的vision transformer。相反，对于基于图像的以人为中心的视觉任务，我们在不同阶段对tokens进行聚类。据我们所知，这是首次将聚类用于visual tokens的生成。

3.Method

3.1 Overview Architecture

如图2所示，我们的TCFormer遵循流行的多阶段体系结构。TCFormer由4个分级阶段和一个 Multi-stage Token Aggregation (MTA) head组成。每个级包含几个堆叠的transformer blocks。在两个相邻阶段之间，插入Clustering-based Token Merge(CTM)块以合并tokens并为下一阶段生成tokens。MTA head聚合所有阶段的token features，并输出最终热图。

*图2：在给定一幅输入图像的情况下，利用单层卷积提取特征图，并将特征图中的像素作为初始的visual tokens。这些初始的visual tokens以多阶段的方式处理并且每个阶段由多个transformer blocks组成，在每个阶段的最后，vision tokens与CMT块合并以生成下一阶段的tokens。最终，所有的visual tokens被送入MTA head来预测热图。在CMT block中，输入的tokens首先会被聚类，然后通过加权平均进行合并。合并的tokens和原始的tokens然后被传入到一个transformer block中，以便更好地进行特征聚合。

3.2Transformer block

图3展示了TCFormer中的transformer block。类似PVT，首先采用一个空间缩减层来减少计算复杂度，这个空间缩减层首先将vision tokens转换为特征图，然后通过带有步长的卷积层来减小特征图的分辨率。经过处理的特征图中的像素比visual tokens少得多，作为keys和values送入MHSA module。MHSA module聚合了tokens之间的特征。受LovalViT,CvT,HRFormer的启发，我们利用一个深度卷积层来捕获局部特征和位置信息，并且移除了显式的位置编码。

*图3：TCFormer中的transformer block的结构。在MHSA module之前采用一个空间缩减模块来减小计算量。一个深度卷积层在注意力模块之后来捕捉局部信息

3.3Clustering-based Token Merge (CTM) Block

如图2所示，我们的CTM块有两个过程：token聚类和特征合并。

我们首先利用token clustering来根据token features将visual tokens分成具有一定数量的clusters，然后利用特征融合来将同一cluster中的tokens合并成一个token。

Token Clustering

在token clustering过程，我们使用了一种基于K-nearest neighbor的密度峰值聚类算法（DPC-KNN）

给定一系列的tokens X，我们根据它的k-nearest neighbours来计算每个token的局部密度：

 KNN(xi)代表token i的k-nearest neighbors，xi和xj是他们对应的token features。

然后，对于每一个token，我们计算距离指标作为它与具有较高局部密度的任何其他token之间的最小距离。 对于具有最高局部密度的tokrn，其指标设置为它与任何其他tokens之间的最大距离。

是距离指标，是局部密度。

我们结合局部密度和距离指标来得到每个token的得分。高的得分意味着更高的几率成为聚类的中心。我们通过挑选具有最高得分的tokens来定义聚类的中心，然后根据特征距离将其他tokens分配给它们最近的聚类中心。

Feature Merging

对于tokens的特征融合，一个直观的方法是直接对聚类中的token特征平均。然而，即使对于具有相似语义的tokens，它们也不是同等重要的。

受DynamicViT启发，我们引入了一个importance score P来显式地表示每个token的重要性，它是根据token特征来估计的。在token importance的指导下来对token特征进行平均：

 Ci代表第i个集群的集合，xj和pj是初始的token features和相对应的importance score，yi是合并tokens的特征。合并token的token区域是原始token区域的并集。

如图2所示，融合的tokens被传入一个transformer block作为Q，原始的tokens被作为K和V。为了保证重要的tokens能够对输出有更大的贡献，importance score P被加入到注意力权重当中：

 dk是queries的通道数。为了清楚起见，我们在这里省略了多头设置和空间缩放层。引入token importance score使我们的CTM模块能够在合并visual tokens时专注于关键图像特征。

我们的CMT block采用了一种高效的GPU实现。集群和特征融合部分只占TCFromer前向时间的9.4%。

3.4Multi-stage Token Aggregation (MTA) Head

前人的工作[HRNet,HRFormer]证明了特征聚合在多个阶段对以人类为中心的视觉任务的好处。为了聚合特征，我们提出了一种基于transformer的Multi-stage Token Aggregation(MTA)head，该头部能够在各个阶段保持细节。

图4a展示了token上采样过程，在token merging过程，每个token被分配到一个聚类并且每个聚类用一个融合的token来表示。我们记录原始token和融合token之间的关系，使用这个记录信息将融合token特征复制到相应的上采样tokens。

图4b显示，在token上采样之后，MTA heads将上一个阶段的token特征加入到上采样的visual tokens中。然后visual tokens被一个transformer block处理。这样的处理被逐步进行直到所有的visual tokens都被融合。最后的tokens，它的token区域是高分辨率特征图的一个单独的像素，可以被很容易地reshape为特征图来进行进一步处理。

*图4a：在token上采样过程，我们使用记录的token关系来将融合的token复制到相应的上采样tokens。图4b:MTA head从最后一个阶段开始，并且通过堆叠的上采样过程和transformer block来逐步融合特征。最后的tokens与特征图是一对一的关系并且reshape成特征图用于后续的处理。

Why not CNN?

大多数前人的工作[Swin,PVT,HRFormer]首先将visual tokens转换为特征图，然后用卷积层聚合多阶段特征。然而，如图5所示，具有重要细节的区域(如人脸)的visual tokens是小尺寸的。将visual tokens转换为低分辨率特征图涉及到对这些tokens进行特征平均，这将导致细节的丢失。通过将所有阶段的visual tokens转换为高分辨率的特征图，可以避免这种损失，但这将导致不可接受的复杂性和存储成本。我们的MTA Head以token的形式聚合特征，并在最后将最终的tokens转换为高分辨率的特征图，在相对较低的复杂度下保留了图像各个阶段的细节。

3.5 Implementation Details

Transform between Vision Tokens and Feature Maps

对于非规则tokens，卷积过程不是平凡的。我们需要在此过程之前将tokens转换为feature maps，并在此过程之后执行逆转换。我们将特征图中的每个像素都看作是一个矩形网格区域的表示。在将vision tokens转换为feature maps时，首先找到每个像素对应的visual tokens，即visual tokens的token区域与像素的网格区域。然后根据重叠区域对token features进行平均，得到每个像素的图像特征。当将特征图转换为visual tokens时，我们直接将token features设置为token区域的平均图像特征。

Cluster Numbers

在CTM block中，cluster的数目可以是任何小于输入token数目的值。在这项工作中，我们将cluster个数设置为输入token数的四分之一，以保持每个阶段的token数与之前的工作 [Swin,PVT] 相同。

4.Experiments

4.1. 2D Whole-body Pose Estimation

全身姿态估计的目标是定位包括面部、手和脚在内的整个人体的细粒度关键点，这要求能够捕获详细信息。

表1显示了将 TCFormer 与 COCOWholeBody V1.0 数据集 [23] 上的最新模型进行比较的实验结果。 结果表明，TCFormer 的全身姿态估计精度（57.2% AP 和 67.8% AR）在很大程度上高于SOTA的自顶向下方法，例如 HRNet。

输入图像中手的尺寸相对较小，这使得手关键点的估计非常困难，并且严重依赖模型捕捉细节的能力。 如表 1 所示，大多数模型的手部性能远低于其他部分。 我们的 TCFormer 在手部关键点估计方面取得了很大的进步，即比 HRNetw32 [50] 高 6.2% AP，比 SBL-Res152 [66] 高 5.3% AP，这证明了 TCFormer 在捕获小尺寸关键图像细节方面的出色能力。

4.2 3D Human Mesh Reconstruction

以往的三维人体网格重建工作可分为基于模型的方法和无模型方法。我们结合TCFormer和HMR Head构建了一种基于模型的方法。

我们将TCFormer与最先进的基于SMPL的方法进行了比较，并将结果显示在表2中。TCFormer在结构和模型复杂性方面优于大多数以前的工作[26，27，29，77]。与使用额外密集监督的DSR方法相比，TCFormer甚至获得了与之相当的结果[13]。 Human Mesh Reconstruction的结果进一步验证了TCFormer在捕捉重要图像特征方面的有效性。它不仅适用于密集预测任务，而且适用于基于全局特征的回归任务。

4.3 2D Face Keypoint Localization

 如表 3 所示，与完整测试集和所有子集上的其他最先进的方法相比，TCFormer 实现了卓越的性能（4.28% NME）。 TCFormer 甚至比具有额外信息的方法具有更低的错误，例如使用强数据增强的 PDB [14] 和使用额外边界信息的 LAB [64]。 人脸对齐的性能验证了 TCFormer 超越人体估计的多功能性。

4.4 Image Classification

尽管我们模型的目标不是图像分类，但在 ImageNet-1K 上的实验结果表明，与最先进的架构相比，TCFormer 实现了具有竞争力的性能（82.4% Top-1 Acc.），这表明我们的 CTM 块 在提取一般图像特征方面也很有效。

4.5 Ablation Study

我们对全身姿态估计的任务进行了消融实验，如表 1 所示。

Effect of CTM. 

为了验证 CTM 的效果（第 3.3 节），我们通过用跨步卷积层替换 CTM 来构建一个baseline transformer network，并发现性能显着下降（-3.7% AP 和-3.9% AR）。 依赖详细信息的部分（例如脚和手）的性能下降比其他部分更显著。 脚和手的估计 AP 分别下降了 13.6% 和 5.6%，而人体的 AP 仅下降了 2.4%。 它证明了性能下降主要是由图像细节的丢失引起的，并验证了我们的 CTM 块在捕获小尺寸图像细节方面的有效性。

Effect of MTA Head. 

为了验证 MTA Head 的效果（第 3.4 节），我们将其替换为反卷积头 [66]，并注意到 -1.9% AP 和 -1.6% AR 的性能下降。 尤其是脚（-4.0% AP）和手（-3.6% AP）的性能下降比身体（-1.2% AP）更明显。 结果表明，MTA Head 有利于保留身体部位的图像细节。

4.6 Qualitative Evaluation

 图 6 显示了人体全身姿态估计、3D 网格重建和面部对齐的一些定性结果。 图 7 显示了上述任务中visual tokens分布的一些示例。 如图 7 所示，visual tokens集中在前景区域，并且只有很少的tokens来标记背景区域，即使背景很复杂。 具有精细空间尺寸的visual tokens用于具有重要细节的区域，例如全身估计和3D网格重建中的面部和手部区域。 对于人脸对齐，TCFormer 将精细标记分配给人脸边缘区域。 动态visual tokens使 TCFormer 能够更有效地捕捉图像细节并获得更好的性能。

5.Ayalysis

在这一部分中，我们将解释为什么 TCFormer 专注于具有重要细节的区域。 以人体全身姿态估计任务中的人脸区域为例，为了区分人脸密集的关键点，TCFormer倾向于针对不同的人脸区域学习不同的特征。 由于代表不同面部部分的tokens的特征不同，CTM块倾向于将它们分组到不同的簇中。 所以它们在接下来的合并过程中不会相互合并，而是保持在一个很好的空间大小，这反过来又有助于特征学习。

我们通过训练两个具有不同任务的模型来验证解释。 第一个旨在仅估计手部关键点，而另一个旨在仅估计面部关键点。 我们在图 8 (b) 和 (c) 中可视化这两个模型的tokens分布，并发现tokens分布是特定于任务的，这与我们的解释相符。

6. Conclusions and Limitations

在本文中，我们提出了 Token Clustering Transformer (TCFormer)，这是一种新颖的基于 Transformer 的架构，用于以人为中心的视觉任务。我们提出了一个 Clustering-based Token Merge（CTM）块，以使transformer具有在重要区域保留更精细细节的能力，并且较少关注无用的背景信息。实验表明，我们提出的方法显着改进了其baseline，并在几个以人为中心的视觉任务上实现了竞争性能，即全身姿态估计、人体网格恢复和人脸对齐。 TCFormer 的主要限制是 KNN-DPC 算法的计算复杂度与token数成二次方，这限制了 TCFormer 在大输入分辨率下的速度。这个问题可以通过将令牌分成多个部分并执行部分token聚类来缓解。我们设想所提出的方法是通用的，可以应用于广泛的视觉任务，例如对象检测和语义分割。未来的工作将集中在探索 CTM 在这些视觉任务上的有效性。