代码: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可视化