CVPR 2023 autofocusformer核心部分翻译

CVPR 2023 autofocusformer 核心部分翻译

前文摘要引言和部分方法见：乄洛尘的文章
代码链接见：https://github.com/apple/ml-autofocusformer

3. 方法

在backbone中，从patch embedding module（2层步长为2的3x3 conv）开始，然后聚类算法（3.1节），然后几个局部注意力的transformerblocks（3.2节），然后新的自适应下采样模块（3.3节）。ps：分割头本文不说。

3.1 聚类和邻域

  一个局部邻域可以方便地通过在2D网格上切出一个正方形窗口来定义。与此形成对比的是,对于一个无序的点集,确定3D点云中的邻域的传统方法依赖于诸如k近邻(kNN)等算法,它计算点之间的成对距离。朴素的kNN算法具有二次时间复杂度。有趣的是,许多加速kNN的算法的第一步都是在给定的点上进行k-means聚类,以减少邻域搜索的空间。受此启发,我们也使用簇来定义局部邻域,即我们将tokens分为若干簇,并将邻域定义为包含几个相邻的簇。不幸的是,传统的聚类算法如k-means或局部敏感哈希并不直接适合我们的目的。首先,它们需要多次迭代或哈希轮数,因此相比深度网的前向传播而言速度较慢。更重要的是,它们导致的簇中点的数量不一。不均匀大小的簇导致需要用零填充的张量,这将浪费内存和时间。注意我们仅对2D位置进行聚类而不是更高的维度,我们提出一种平衡聚类方法来解决上述问题。

3.1.1 平衡聚类

  我们的方法不是采用传统的迭代方法,而是将图像上的2D位置排列成1D数组,然后将数组划分成等大小的组。为此,我们考虑了空间填充曲线,它将2D位置排列在1D线上,同时尽量保留2D距离测度。因此,在线上接近的点在2D空间中也比较接近。然而,由于从2D到1D的转换,完全保留2D度量是不可能的,如果我们直接利用空间填充曲线将tokens划分为簇,就会出现伪影。为了部分减轻这一问题,我们采用了2阶段过程进行聚类。其思想是仅在粗糙级别上利用空间填充曲线来获得稀疏、正则采样的2D图像位置的顺序。然后,根据与这些位置的2D距离对tokens进行排序。
  具体而言,我们首先将图像划分为粗糙的规则网格,其中的正方形patches数量与预期的簇数量相似。我们将网格中每个正方形patches的中心称为空间填充锚点。空间填充曲线在锚点之间建立顺序。根据这个顺序,对于属于锚点$a_i\in R^2$的位置$p\in R^2$的tokens,我们可以定义其前一个锚点$a_{i-1}$和下一个锚点$a_{i+1}$。然后,我们计算所有token $p$计算从$p$到这两个锚点的距离比率r。
$$r(p)=\frac{d_{i-1}(p)}{d_{i+1}(p)}=\frac{\Vert p-a_{i-1}{\Vert }_2}{\Vert p-a_{i+1}{\Vert }_2}，（1）$$
  现在,在每个正方形补丁内,我们按r的升序对tokens进行排序,使得接近前一个锚点的tokens排在前面。此过程因此建立了所有tokens的顺序。最后,为了获得平衡的聚类,我们简单地将数组划分成等大小的组。图3说明了整个算法。

  因为我们可以在O(1)时间内通过简单量化它们的坐标来找到每个tokens对应的空间填充锚点,所以聚类的总体时间复杂度不超过对局部补丁中的所有tokens位置进行一次排序,与网络的时间复杂度相比可以忽略不计,因为没有涉及到特征通道。注意,聚类算法只需要在每个阶段开始时执行一次,簇信息可以用于该阶段的所有注意力块以及末端的下采样模块。与先前的平衡聚类工作不同,该算法不是迭代的,并产生完全平衡的簇。它也保证每个tokens只属于一个簇,不同于RoutingTransformer,其中一些tokens可能不属于任何簇。但是,请注意我们的平衡聚类仅适用于低维点。在这项研究的早期阶段,我们探索了使用嵌入特征而不是仅2D位置进行聚类,但性能差异可以忽略。因此,我们决定仅对位置进行聚类,这使我们能够利用所提出的算法生成完美平衡的簇。

3.1.2 聚类中的邻域

  为了促进整个图像中的信息流动,重要的是注意力不应仅限于同一簇内的位置。例如,在Swin Transformers中,连续层之间的窗口移动允许像素在不同层中关注不同的邻居。然而,在我们的情况下,对每一层重新进行聚类会增加不必要的计算量。因此,我们选择使用较小的簇,并允许每个tokens关注来自附近R个簇的tokens。为实现这一点,我们使用的邻域大小是簇大小的几倍。这是有利的,因为邻域是重叠的,保证了簇之间的信息交换。

3.2 Transformer Attention Blocks

  在任何阶段,设tokens数量为N,每个tokens的邻居数量为M,头数为H,特征通道数为C。具有相对位置嵌入的单个头的局部注意力A通过每个tokens关注其邻域中的所有其他tokens来计算,即
$$A=softmax(Q{K}^T+P),（2）$$
  其中$Q\in R^{N\times C/H},K\in R^{M\times C/H},A\in R^{N\times M},P\in R^{N\times M}$ 分别是查询、键、注意力和位置嵌入矩阵,且
$$P_{i,j}=w(pi-pj).（3）$$
  这里，$pi$,$pj$是两个相邻tokens的(x,y)坐标,$w$(·)是一个返回该头的标量位置嵌入的函数。对于固定形状邻域的模型,相对位置嵌入可以存储在矩阵中,并在需要时读取。但是对于我们的模型,邻居的位置事先是未知的,因此我们需要通过可学习函数将坐标差投影到位置嵌入中。我们用一个全连接层实现$w$。

3.3 Adaptive downsampling

  在每个阶段的末尾,除了最后一个阶段,我们采用一个下采样模块。下采样模块包含两部分,一个用于选择最重要tokens的可学习评分组件,以及周围选定tokens的邻域合并步骤。

3.3.1 Learnable Importance Score

  我们学习预测每个token i的标量分数si,以指示其相对于任务损失的重要性。学习该分数的主要困难在于下采样限制了选定tokens的数量。如果任务损失仅反向传播到选定的tokens,则未选定的tokens可能不会接收到任何梯度。为了向所有tokens反向传播梯度,我们提出学习每个token的重要性分数与邻域合并过程结合。
  具体而言,假设我们正在合并一个特定位置$p_c$处token周围的邻域$N(p_c)$。类似于注意力块的邻域$N(p_c)$从其最近的R个簇获得,第i个邻居token表示为一个元组$(pi,fi)$,其中位置$pi=(xi,yi)$,特征$fi\in R^C$。在邻域内,合并使用PointConv层的修改版本执行,如下所示:
$$f_{merged}(p_c)=vec（\sum _{(pi,fi)\in N(p_c)}\stackrel{S_i}{\overbrace{\sigma (l(fi))}}\cdot W(p_i-p_c)f_i^{\top }）U,$$
其中$f_{merged}\in R^{C\prime }$是具有$C^{\prime }$输出的合并输出,$vec(\cdot )$表示向量化,$\sigma (\cdot )$是sigmoid函数, $l(\cdot )$是一个预测标量“重要性分数”的全连接层,$W(\cdot )$是一个输出形状为$C_{mid}\times 1$的多层感知机(MLP),它创建输入特征$f_i$在邻域中的不同加权组合。权重学习为每个邻域token的位置与合并中心位置pc之间的相对坐标的函数。最后,$U\in R^{C_{mid}C\times C^{\prime }}$作为一个全连接层实现。
  注意,等式(4)与PointConv $f=vec(\sum W{f}^{\top })U$类似,已被证明等价于一个连续的卷积层。我们添加了$S_i=\sigma (l(f_i))$来调节这个函数,这使得模型在邻域合并期间可以“关闭”不重要的token,使它们的特征不被用于下一阶段。因此,分数si可以看作是一个指示器,表示模型认为一个token的重要程度。这个分数接下来可以用于选择合并中心。
  这种形式化让我们可以直接从任务损失中学习分数si,而不必诉诸于先前工作中使用的昂贵技术,如策略梯度。一些以前的工作使用Gumbel Softmax/Sigmoid来获得一个可微的二值掩码,以保持在下采样后向“删除”的token的梯度流;由于我们采用合并策略而不是删除,我们不需要这样一个硬二值掩码。我们在图4中提供了邻域合并过程的可视化。

3.3.2 Grid priors and Selection of Tokens to Sample

  最终用于选择作为合并中心的token的分数是$g_i+\alpha S_i$,这是学习到的分数$s_i$和网格先验gi的加权和,其中$\alpha$是超参数。虽然$s_i$仅基于特征向量,但网格先验帮助模型区分不同位置的相似token,从而方便模型在本地无纹理区域进行适当的、均匀步长的下采样。我们的网格先验$g_i$模拟了传统网格下采样的行为——交替将1和0赋给token。然而,在token已经不规则采样的阶段,再添加规则网格先验就不再合理。因此,我们提倡一种“自适应”网格先验,它考虑采样token的局部密度。
  具体来说,自适应先验根据采样token的局部步幅选择局部网格级别。对于每个token $p_i=(x_i,y_i)$,我们将其“步幅” $t_i$赋值为与其最近邻的距离,向上舍入至2的最近幂。例如,如果token与其最近邻距离为3像素,则将其步幅赋值为$t_i$ = 4。然后,如果$(x_{i}mod2t_i=0)\wedge (y_{i}mod2t_i=0)$,我们设$g_i=1$，也就是说,我们希望将局部步幅从ti下采样到2$t_i$。因此,如果$t_i=1$,则$g_i=1$对应交替像素,如果$t_i$ = 2,则$g_i=1$对应每4个像素。此外,我们在第j阶段将$(x_{i}mod{2}^{j+1}=0)\wedge (y_{i}mod{2}^{j+1}=0)$的token的网格先验设置为无穷大。我们称这些token为“保留”。我们保留这些粗网格token,以确保前向传播期间图像中远程区域之间的连通性。
  总之,自适应下采样的工作流程如下:
1)为每个token fi获得重要性分数$S_i=\sigma (l(fi))$,
2)为每个token计算网格先验gi,
3)选取$g_i+\alpha S_i$值最高的前x% token(例如1/4或1/5),
4)在前x% token的位置周围执行Eq. (4)给出的邻域合并公式,并获得下一阶段的合并token。