[CVPR‘22-Oral] Not All Tokens Are Equal Human-centric Visual Analysis via Token Clustering Transform

论文：https://openaccess.thecvf.com/content/CVPR2022/papers/Zeng_Not_All_Tokens_Are_Equal_Human-Centric_Visual_Analysis_via_Token_CVPR_2022_paper.pdf

代码：https://github.com/zengwang430521/TCFormer

• 一种Transformer Block，结合了Spatial reduction layer和depth-wise convolutional layer，可以节省计算开销和显示的位置编码；
• 一种新的Transformer渐进迭代方式：通过聚类逐步减少token数量，使得后层token具有更大感受野。相较于传统CNN方法，其感受野范围形状不受限制。
• 讨论了Dynamic Token Generation问题，不同于传统token剪枝的方法，属于token融合。

摘要

• 研究背景：Vision tokens是vision transformer的重要组成部分。传统方法将图片分割为均匀块，作为vision tokens；
• 存在问题：对以人为中心的视觉任务而言，人需要更多token和更细粒度的特征表征，相对而言，则需要减少用于表征背景的token数量；
• 解决方法：提出一种适用于人体关键点检测的dynamic token generation算法，Token Clustering Transformer (TCFormer)。TCFormer逐步聚集并融合不同位置的tokens，使得TCFormer的token不仅可以聚焦重要区域，还可以调整 token 形状以适应语义，并对包含关键细节的区域采用细粒度分辨率，捕捉细节信息。
• 在whole-body pose estimation任务上（COCO-WholeBody）和3D human mesh reconstruction on 3DPM任务上（3DPW）超过部分方法。

本文贡献

• 提出Token Clustering Transformer (TCFormer)。可以生成表征不同位置、大小和形状的tokens，该步主要通过Clustering Token Merge (CTM) 模块实现，包含以下三步：1）分组，基于KNN-based density peaks clustering算法，将tokens分组；2）融合：同一组token特征的均值；3）将分组融合后的特征送入transormer模块。本文是第一个将聚类用于dynamic token generation的工作。
• 提出Multi-stage Token Aggregation (MTA) head 用于局部多步token特征，保留细节信息。

近期工作

• Dynamic Token Generation：可以分为token剪枝和token融合。本文可以归类为token融合方法。
• DynamicViT和PnP-DETR通过预测token分数，选择其中重要的token，对于不重要的token，DynamicViT丢弃不用，PnP-DETR则用稀疏特征向量表示；
• DVT 的token采样来自不同分辨率特征，采样分辨率基于分类的困难程度决定；
• PS-ViT的token基于固定大小采样，采样中心在训练过程中逐步调整。

Method

Overview Architecture

• TCFormer = 4 stages + Multi-stage Token Aggregation (MTA) head。
• stage由多个stacked transformer blocks组成，stage和stage之间通过Clustering-based Token Merge (CTM) block链接，CTM负责融合token，并产生用于下一阶段的token。
• MTA聚合所有阶段产生的token特征，并输出最终的heatmap。

Transformer Block

• spatial reduction layer，将vision tokens转换为特征图，并用strided convolution layer减低特征图分辨率，然后将处理后的特征图作为keys和values送入multi-head attention模块；
• 去除显示的位置编码，用depth-wise convolutional layer捕捉局部特征和位置信息。

Clustering-based Token Merge (CTM) block

Token Clustering：Density Peask Clustering KNN，DPC-KNN算法的变体。

1）计算局部密度ρ，基于K近邻间的欧式距离计算token的局部密度：

2）distance indicator δ：和其他有更高局部密度的token之间的最小距离，对有最高局部密度的token，它的distance indicator为它和其他token的最大距离：

每个token的分数 = ρ x δ，分数越大表明越适合作为聚类中心。本文选择分数最高的token作为聚类中心，其他token选择最近邻中心来聚类。

Feature Merging

参考Dynamic ViT，计算token的重要性分数P，基于重要性分数P计算簇间均值：

 之后，融合的token作为Q，原本的tokens作为K和V，并将重要性分数P作为计算的一部分：

 Multi-stage Token Aggregation (MTA) Head

通过之前的聚类记录，将token feature逐步上采样回原来的token数量特征。

Experiment

在COCO-WholeBody、3DPW、Human3.6M、WFLW和ImageNet-1k上进行了实验。

 

 一些token可视化