ViT自适应patch划分 ACM MM 2021

Transformer在计算机视觉方面取得了巨大的成功，而如何分割图像中的patch仍然是一个问题。现有的方法通常使用固定大小的patch embedding，这可能会破坏对象的语义。为了解决这一问题，作者提出了一种新的Deformable Patch模块（DPT），该模块以数据驱动的方式将图像自适应地分割成不同位置和大小的patch，而不是使用固定patch。该方法可以很好地保留patch中的语义。DPT是可插拔的模块，可以很容易集成到不同的Transformer中，从而实现端到端的训练。结果表明，DPT在ImageNet分类上可以达到81.9%的Top-1准确率，在MSCOCO目标检测上，结合RetinaNet可以达到43.7%的box mAP，结合Mask R-CNN可以达到44.3%。

文章来源：DPT: Deformable Patch-based Transformer for Visual Recognition (ACM MM2021)

引言

• 图1：vanilla的patch分割实例。a为原始patch embedding，这是一种对张量的固定分割，可能会破坏语义。b为DPT的patch embedding，DPT使用可学习的位置和尺度对张量进行动态分割。

transformer在NLP兴起，逐渐成为离散序列建模的主流方法，后来被应用于视觉研究。transformer通常将输入图像分成固定大小的patch序列（patch大小16×16），利用多头注意力对patch序列建模上下文关系。目前，ViT依然存在改进挑战，目前的工作都基于固定的patch，这种hard patch会带来两个问题：

• 图像中局部结构崩溃，如图1a所示，hard patch难以捕捉到完整的与对象相关的局部结构，因为对象在不同的图像中具有不同的比例。
• 图像之间的语义不一致，不同图像中的同一对象可能具有不同的几何变换（缩放，旋转），固定的patch分割会捕获不同图像中同一个对象的不一致信息。

因此，固定的patch分割会破坏语义信息，导致性能下降。

为了解决上述问题，作者提出了一个名为DePatch的新模块，它以可变形的方式分割图像为patch。通过这种方式，我们可以在一个patch中更好地保留语义，减少固定patch分割造成的语义破坏。为了实现这一点，DPT学习了feature map空间中每个patch的偏移和比例。偏移量和比例是基于输入feature map学习的，并为每个patch生成，如图1b所示。

所提出的模块是轻量级的，并且引入了非常少量的参数和计算。更重要的是，它可以作为一个即插即用模块（plug-and-play）工作，可以很容易地集成到其他Transformer架构中。带有DePatch模块的Transformer被命名为基于可变形patch的Transformer DPT。在这项工作中，作者将DePatch模块集成到Pyramid Vision Transformer（PVT）中，以验证其有效性，因为PVT在像素级预测任务（如目标检测和语义分割）中实现了最先进的性能。通过可变形patch调整，DPT基于局部上下文结构为每个patch生成完整、鲁棒和有判别性的特征。因此，它不仅可以在分类任务上实现高性能，而且在高度依赖于局部特征的任务（例如目标检测，语义分割）上也优于其他方法。与PVT Tiny相比，DPT在ImageNet分类上实现了2.3%的改进，在用于MSCOCO对象检测的RetinaNet和Mask R-CNN框架下，box mAP分别提高了2.8%和3.5%。

方法

前言：Vision Transformer

ViT由三个部分组成，一个patch embedding模块，多头注意力模块，前馈多层感知机MLP。网络从patch embedding模块开始，该模块将输入张量转换为token序列，然后输入多头自注意力模块和MLP以获得最终表示。下面主要阐述patch embedding模块，并回顾多头自注意力模块。

patch embedding模块将图像分割成具有固定大小和位置的patch，并用线性层embedding每个patch。假设输入张量为$A\in R^{H\times W\times C}$，为了化简，假设$H=W$。过去的工作将$A$分割为具有$N$个patch的序列，patch的大小为$s\times s$，其中$s=[H/\sqrt{N}]$。这个序列被表示为${\left\{z^{(i)}\right\}}_{1\le i\le N}$。

为了更好解释patch分割过程，作者重新定义patch embedding模块。每个patch $z^{(i)}$可以被看作输入图像中一个矩形区域的表示。我们记其中心坐标为$(x_{ct}^{(i)},y_{ct}^{(i)})$。由于patch的大小固定，左上角坐标和右下角坐标可以直接得到：$$(x_{ct}^{(i)}-s/2,y_{ct}^{(i)}-s/2),(x_{ct}^{(i)}+s/2,y_{ct}^{(i)}+s/2)$$在当前patch $i$ 中，有$s\times s$个像素，它们的坐标被表示为${\vec{p}}^{(i,j)}=(p_x^{(i,j)},p_y^{(i,j)})$。所有的坐标${\vec{p}}^{(i,j)}$都是整数。每个像素对应的特征被记为${\left\{{\tilde{a}}^{(i,j)}\right\}}_{1\le j\le s\times s}$。这些特征被flatten并被线性处理得到patch $i$ 的新表示：$$z^{(i)}=W_{patch}\cdot concat\left\{{\tilde{a}}^{(i,1)},...,{\tilde{a}}^{(i,s\times s)}\right\}+b_{patch}$$多头自注意力模块聚合整个输入序列上的相关信息，为每个token提供全局视图。该模块学习每个注意力头的三组代表性特征：query $Q_h\in R^{N\times d}$，key $K_h\in R^{N\times d}$，value $V_h\in R^{N\times d}$。$Q_h$和$K_h$相乘得到attention map $Att{n}_h$，其反应了不同patch之间的相似性。attention map再作为权重改变$V_h$。计算不同头部的独立结果，以获得不同的特征，然后将所有头部的结果连接并转换为新表示$Z^{\prime }$。$$Att{n}_h=Softmax(Q_h{K}_h^T/\sqrt{d})$$ Z ′ = c o n c a t { A t t n 1 V 1 , . . . , A t t n H V H } W p r o j + b p r o j Z'=concat\left\{Attn_{1}V_{1},...,Attn_{H}V_{H}\right\}W_{proj}+b_{proj} Z′=concat{Attn1​V1​,...,AttnH​VH​}Wproj​+bproj​

DePatch模块

对于固定patch分割的方式，patch的位置$(x_{ct}^{(i)},y_{ct}^{(i)})$和大小$s$是固定的，因此矩形区域对于每个patch是不变的。每个patch的特征直接用其内部像素表示。为了更好地定位重要区域并适应对象的形变，作者开发了DePatch。

• 图2：DePatch模块说明。利用局部特征预测偏移量和尺度，通过双线性插值得到新的embedding。

首先，将每个patch的位置和size转化为基于输入内容的预测参数。对于位置，预测原始patch中心的偏移量$({\delta }_x,{\delta }_y)$。对于size，预测$s_h$和$s_w$替换$s$。这样，得到了一个新的矩形区域，并记其左上角和右下角坐标为$(x_1,y_1)$和$(x_2,y_2)$。为了清楚，先省略上标$(i)$，因此注意，$({\delta }_x,{\delta }_y,s_w,s_h)$对于每个patch都是不同的。$$x_1=x_{ct}+{\delta }_x-\frac{s_w}{2},y_1=y_{ct}+{\delta }_y-\frac{s_h}{2}$$$$x_2=x_{ct}+{\delta }_x+\frac{s_w}{2},y_2=y_{ct}+{\delta }_y+\frac{s_h}{2}$$如图2所示，作者添加了新的分支来预测这些参数。基于输入特征图，为每个patch预测$({\delta }_x,{\delta }_y,s_w,s_h)$，然后再embedding。

偏移量和尺度的预测如下：$${\delta }_x,{\delta }_y=Tanh(W_{offset}\cdot f_p(A))$$$$s_w,s_h=ReLU(Tanh(W_{scale}\cdot f_p(A)+b_{scale}))$$其中，$f_p$可以是任何特征提取器，作者使用的是单个线性层。在训练初始阶段，$W_{offset}$和$W_{scale}$和$b_{scale}$被初始化为0。

确定矩形区域后，提取每个patch的特征。主要问题是区域大小不同，预测的坐标通常是float数。因此用抽样-插值的方法（sampling-and-interpolation）来解决这个问题。给定矩形坐标$(x_1,y_1)$和$(x_2,y_2)$，在区域内均匀抽样$k\times k$个points，$k$是一个超参数。每个抽样的位置为${\vec{p}}^{(j)}=(p_x^{(j)},p_y^{(j)})$，$1\le j\le k\times k$。这些抽样点对应的特征为${\left\{{\tilde{a}}^{(j)}\right\}}_{1\le j\le k\times k}$被输入线性层得到patch embedding：$$z^{(i)}=W\cdot concat\left\{{\tilde{a}}^{(1)},...,{\tilde{a}}^{(k\times k)}\right\}+b$$注意，抽样点的index是float数，为了得到该index的特征，需要bilinear interpolation，假设要获取任意一个float index $(p_x,p_y)$ 的特征，则有：$$A(p_x,p_y)=\sum _{q_x,q_y}G(p_x,p_y;q_x,q_y)\cdot A(q_x,q_y)$$$$G(p_x,p_y;q_x,q_y)=max(0,1-|p_x-q_x|)\cdot max(0,1-|p_y-q_y|)$$其中，$G$是双线性插值kernel，它只在靠近$(p_x,p_y)$的四个整型index上取特征。这可以被快速计算。

整体架构

DePatch是一个自适应模块，可以改变patch的位置和尺寸。由于DePatch的可插拔性，作者将其集成到各种ViT中。在文中，作者将其集成到PVT中，PVT分为4个stage，它利用spatial-reduction注意力来降低处理高分辨率特征图的成本。DePatch+PVT成为DPT，通过将stage 2，3，4中的patch embedding换成DPT而构建。整体架构见图3。

• 图3：左边是PVT，右边是DPT。

实验

实验包括图像分类和目标检测。

• 表1：ImageNet上的分类结果。

• 表2：MSCOCO目标检测结果（DePatch+RetinaNeT）。

• 图4：对于DPT-Small的stage 4，可视化学习的patch以及偏移量。可以看出，DPT能够根据内容自适应调整每个patch的位置和比例。