AAformer: Auto-Aligned Transformer for Person Re-Identification 解读

局部特征的方法让Re-ID任务的精度得到了飞跃，那么 使用Transformer还能不能使用局部方法呢？当然可以，本文作者提出了一个Transformer Re-ID局部框架，据说有了此框架，CNN中的局部方法就可以迁移到Transformer，岂不美哉？

文章标题：AAformer: Auto-Aligned Transformer for Person Re-Identification

论文地址：https://arxiv.org/abs/2104.00921

代码地址：即将开源

1 概述

  最近陆续出现基于Transformer的ReID工作，众所周知，在基于CNN的ReID工作里，part对齐/局部特征是很有效的方法，但是这些方法不能直接用在Transformer里(主要还是因为原来的池化不适用于Transformer)，可是对ReID来说，学习局部特征表示或者对齐部分特征是很有必要的。基于此动机，本文提出了一个用于ReID的自动对齐Transformer框架，作者说此框架很灵活，可以由其他人自定义，对此说法我的理解是：有了此框架，之前CNN中的一些对齐方法就可以转移到Transformer上了，代码即将开源。

2 贡献

• 提出首个基于对齐的Transformer ReID 框架，引入part tokens来学习part特征，并将part对齐集成到自注意力中；
• 提出了在线自动对齐的Transformer (AAformer)，自适应地为相同的part token分配相同部件的patch embedding，AAformer同时学习部件对齐和部件特征表示，而不是采用固定的分块方案。

3 详细内容

3.1 主体框架(ViT)

  主体框架就是Vision Transformer (ViT)，作者遵循了ViT的设置，这里再复述一次便于记忆：

  给定输入图像$\mathbf{x}\in {\mathbb{R}}^{H\times W\times C}$，将其reshape为展平的2D patches：$\mathbf{x}\in {\mathbb{R}}^{N\times \left(I^2\cdot C\right)}$，其中$(I,I)$是每个patch的size，$N=(H\cdot W)/I^2$是patch序列的长度(其实就是patch数量)，而后经过可学习的线性投影将patches映射为$D$维向量作为patch embedding，此后加入一个class token (CLS)嵌入以提取全局特征表示，最后，结果向量序列$\mathbf{Z}\in {\mathbb{R}}^{L\times D}$被馈入Transformer Encoder中，其中$L=1+N$，此外，1D的可学习的位置嵌入和上述向量序列加和(对应元素)以获得位置信息。下面展示部分伪代码：

# 只能当成伪代码看
img_size=[256,256]   #假设图片[8,3,256,256]
patch_size=[16,16]
embed_dim=768
class patch_embed(nn.Module):
    #16x16
    num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])   
    self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, 
                          stride=patch_size) 
    #proj: [8, 768, 16, 16]
    ...
    x = self.proj(x).flatten(2).transpose(1, 2)    #[8,256,768]
    
#cls_token [1, 1, 768]
cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))  
#位置嵌入 [1, 257, 768]
pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) 

## forward  x
## 输入图像x
B = x.shape[0]   #8
x = patch_embed(x)  #[8,256,768]
cls_tokens = cls_token.expand(B, -1, -1)  #[8,1,768]
x = torch.cat((cls_tokens, x), dim=1)     #[8,257,768]
x = x + self.pos_embed                    #[8,257,768]

  对于下图所示的Transformer的Encoder block，这里就不再详述，就是layer norm—>多头自注意力->layer norm—>MLP，网上有铺天盖地的教程，其网络层的堆叠如下代码所示：

(0): Block(
        (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
        (attn): Attention(
          (qkv): Linear(in_features=768, out_features=2304, bias=True)
          (attn_drop): Dropout(p=0.0, inplace=False)
          (proj): Linear(in_features=768, out_features=768, bias=True)
          (proj_drop): Dropout(p=0.0, inplace=False)
        )
        (drop_path): Identity()
        (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
        (mlp): Mlp(
          (fc1): Linear(in_features=768, out_features=3072, bias=True)
          (act): GELU()
          (fc2): Linear(in_features=3072, out_features=768, bias=True)
          (drop): Dropout(p=0.0, inplace=False)
        )
      )

3.2 Auto-Aligned Transformer

  上一节提到了Transformer的Encoder结构，本文最核心的创新就是把其中的MSA(多头自注意力)换成了Multi-head Auto-Alignment ，在序列中作者添加了$P$个可学习的part tokens，结合ViT之后，整个框架便如下图所示，其实很巧妙，Transformer中class token是用来表示全局信息的，作者添加几个part tokens来表示局部信息，而且不是每一个part token都对应一个patch embedding，part token只与patch embedding的子集有联系，这样序列$\mathbf{Z}$的长度就不再是$1+N$而是$1+L+N$。

  具体地，定义${\mathbf{q}}_p,p\in \{1,\dots ,P\}$作为part token $p$的$query$向量，${\Phi }_p$为分配给$p$的patch embeddings子集，和原来的自注意力一样($q,k,v$)，如公式(1)所示，作者说这不像CNN那样只在backbone顶部进行一次对齐，这个Transformer框架会在整个模型中对齐，从而增强对齐效果，且灵活性强，利用这个框架，可以容易地将CNN中的一些局部方法用到Transformer中，例如PCB这种水平条带方法，也就是说本文方法可以作为基于part的Transformer ReID框架。
$$\text{ Attention }\left({\mathbf{q}}_p,{\mathbf{K}}_{{\Phi }_p},{\mathbf{V}}_{{\Phi }_p}\right)=\mathrm{Softmax}\left(\frac{{\mathbf{q}}_p{\mathbf{K}}_{{\Phi }_p}^T}{\sqrt{D}}\right){\mathbf{V}}_{{\Phi }_p}\text{\hspace{1em}(1)}$$
  Multi-head Auto-Alignment：作者发现基于条带的patch分配不能够很好地对齐行人特征，且利用预训练模型提供的语义信息不能识别非人类部件如：背包，为此，作者提出了一个在线的Auto-Aligned Transformer以自适应地将同一part的patch embedding分配给相应的part token。作者的基本思想是将part tokens视为part prototypes，然后在线自适应将patch embedding分组到patch token上(通俗说就是聚类思想实现patch分配)，如下图所示。

  结合自注意力，part tokens的$query$向量可以视为part prototypes $\mathbf{C}\in {\mathbb{R}}^{P\times D}$，作者说需要关注的是如何将图像patches的$key$向量$\mathbf{S}\in {\mathbb{R}}^{N\times D}$映射到part prototypes，记此映射为$\mathbf{Y}\in {\mathbb{R}}^{P\times N}$，位置$(p,n)$处的值表示patch embedding $n$属于part token $p$的概率，为了最大化图像patches的$key$向量和part prototypes的相似度，$\mathbf{Y}$由公式(2)优化：
$$\underset{\mathbf{Y}\in \mathcal{Y}}{\max }\mathrm{Tr}\left({\mathbf{Y}}^{\top }{\mathbf{C}\mathbf{S}}^{\top }\right)+\epsilon H(\mathbf{Y})\text{\hspace{1em}(2)}$$
其中$H(\mathbf{Y})=-{\sum }_{ij}{\mathbf{Y}}_{ij}\log {\mathbf{Y}}_{ij}$是熵函数，$ϵ$是控制映射平滑度的参数，实验中设为0.05。为了进一步防止所有patches分配给同一part token，作者遵循[1,2]的设置，使用约束矩阵$\mathbf{Y}$来强制等分，[1]旨在将所有的图像分类到class prototypes，作者采用此方法，通过将transportation polytope(运输多面体)限制在图像的patches中来实现patch分配，如公式(3)所示，
$$\mathcal{Y}=\left\{\mathbf{Y}\in {\mathbb{R}}_{+}^{P\times N}\mid \mathbf{Y}{1}_N=\frac{1}{P}{1}_P,{\mathbf{Y}}^{\top }{1}_P=\frac{1}{N}{1}_N\right\}\text{\hspace{1em}(3)}$$
其中$l_P$表示$P$维1向量，这些约束强制一张图中的每个part prototypes被选择至少$N/P$次，软分配${\mathbf{Y}}^{\ast }$是问题2(公式(2))在集合$\mathbb{Q}$上的解，采用标准化指数矩阵的形式，如公式(4)所示，
$${\mathbf{Y}}^{\ast }=\mathrm{Diag}(\mathbf{u})\exp \left(\frac{{\mathbf{C}\mathbf{S}}^{\top }}{\epsilon }\right)\mathrm{Diag}(\mathbf{v})\text{\hspace{1em}(4)}$$
其中$\mathbf{u}$和$\mathbf{v}$分别是${\mathbb{R}}^P$和${\mathbb{R}}^N$的重正则化向量，重正则化向量通过少量矩阵相乘计算，用的是Sinkhorn-Knopp算法[3]的变体，该算法在GPU上运行十分高效，最后作者用上述的自动对齐方法构建了多头自动对齐MAA，替换掉了Transformer Encoder中的MSA (Multi-head Auto-Alignment)。

3.3 目标函数

  如上一节框架图所示，训练阶段，对于Re-ID head的输入Transformer层的输出，CLS $Z_0^M$表示全局特征，$Z_{1:P}^M$表示局部特征，损失函数就用最常见的三元组和交叉熵，然后加和。测试阶段，CLS和局部特征连接在一起表示图片。

3.4 实验

  用的模型是ViT base，输入图像大小为$384\times 384$，patch长宽均是16，采用了随机裁剪、水平翻转和随机擦除三个数据增强方案，作者使用了ImageNet上预训练的ViT模型，由于MAA没给ViT添加额外的参数，所以预训练模型可以直接拿来用，作者发现仅仅24个epochs模型就可以收敛，学习率前两轮从$3.5\times 10^{-5}$到$3.5\times 10^{-4}$线性增加(warm up)，此后在第8轮和第14轮学习率衰减为原来的10分之一，batch大小为128，使用Adam优化器。

4 心得

  本文提出的方法思想很简单，但是MAA部分有些难以理解，因为牵扯到机器学习里的最优传输知识，总的来说，本文探索了Transformer+局部方法实现ReID的可行性，等到开源后便可以借鉴此框架。

参考文献

1. Asano Y M, Rupprecht C, Vedaldi A. Self-labelling via simultaneous clustering and representation learning[J]. arXiv preprint arXiv:1911.05371, 2019.
2. Caron M, Misra I, Mairal J, et al. Unsupervised learning of visual features by contrasting cluster assignments[J]. arXiv preprint arXiv:2006.09882, 2020.
3. Cuturi M. Sinkhorn distances: Lightspeed computation of optimal transport[J]. Advances in neural information processing systems, 2013, 26: 2292-2300.