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进行排序。
??具体而言,我们首先将图像划分为粗糙的规则网格,其中的正方形补丁数量与预期的簇数量相似。我们将网格中每个正方形补丁的中心称为空间填充锚点。空间填充曲线在锚点之间建立顺序。给定此顺序,对于属于锚点ai ∈ R2的位置p ∈ R2的tokens,我们可以定义其前一个锚点

a_{i?1}

ai?1?和下一个锚点

a_{i+1}

ai+1?。然后,我们通过对所有tokens

p计算从tokens

p到这两个锚点的距离比率r来计算。

(

)

(

)

(

)

∥

r(p) = frac{d_{i?1}(p)} {d_{i+1}(p)} = frac{∥p ? a_{i?1}∥_2}{∥p ? a_{i+1}∥_2}

r(p)=di+1?(p)di?1?(p)?=∥p?ai+1?∥2?∥p?ai?1?∥2??
??现在,在每个正方形补丁内,我们按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(QK^T + P),

A=softmax(QKT+P),
??其中

∈

Q ∈ R^{N×C/H},K ∈ R^{M×C/H},A ∈ R^{N×M},P ∈ R^{N×M}

Q∈RN×C/H,K∈RM×C/H,A∈RN×M,P∈RN×M 分别是查询、键、注意力和位置嵌入矩阵,且

(

)

P_{i,j} = w(pi - pj).

Pi,j?=w(pi?pj).
??这里，

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

pc?处token周围的邻域

(

)

N(p_c)

N(pc?)。类似于注意力块的邻域

(

)

N(p_c)

N(pc?)从其最近的R个簇获得,第i个邻居token表示为一个元组

(

)

(pi, fi)

(pi,fi),其中位置

(

)

pi = (xi, yi)

pi=(xi,yi),特征

∈

fi∈R^C

fi∈RC。在邻域内,合并使用PointConv层的修改版本执行,如下所示:

(

)

（

∑

(

)

∈

(

)

(

)

(

)

）

f_merged(p_c) = vec（sum_{(pi,fi)∈N (p_c)}overbrace{σ(l(fi))}^{S_i}·W(p_i ? p_c)f_i^?）U,

fm?erged(pc?)=vec（(pi,fi)∈N(pc?)∑?σ(l(fi))

其中

∈

′

f_{merged} ∈ R^{C′}

fmerged?∈RC′是具有

′

C^′

C′输出的合并输出,

(

)

vec(·)

vec(?)表示向量化,

(

)

σ(·)

σ(?)是sigmoid函数,

(

)

l(·)

l(?)是一个预测标量“重要性分数”的全连接层,

(

)

W(·)

W(?)是一个输出形状为

C_{mid} × 1

Cmid?×1的多层感知机(MLP),它创建输入特征

f_i

fi?在邻域中的不同加权组合。权重学习为每个邻域token的位置与合并中心位置pc之间的相对坐标的函数。最后,

∈

′

U ∈ R^{C_{mid}C×C^′}

U∈RCmid?C×C′作为一个全连接层实现。
??注意,等式(4)与PointConv

(

∑

)

f = vec(sum Wf^?)U

f=vec(∑Wf?)U类似,已被证明等价于一个连续的卷积层。我们添加了

(

)

S_i = σ(l(f_i))

Si?=σ(l(fi?))来调节这个函数,这使得模型在邻域合并期间可以“关闭”不重要的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

gi?+αSi?,这是学习到的分数

s_i

si?和网格先验gi的加权和,其中

alpha

α是超参数。虽然

s_i

si?仅基于特征向量,但网格先验帮助模型区分不同位置的相似token,从而方便模型在本地无纹理区域进行适当的、均匀步长的下采样。我们的网格先验

g_i

gi?模拟了传统网格下采样的行为——交替将1和0赋给token。然而,在token已经不规则采样的阶段,再添加规则网格先验就不再合理。因此,我们提倡一种“自适应”网格先验,它考虑采样token的局部密度。
??具体来说,自适应先验根据采样token的局部步幅选择局部网格级别。对于每个token

(

)

p_i = (x_i, y_i)

pi?=(xi?,yi?),我们将其“步幅”

t_i

ti?赋值为与其最近邻的距离,向上舍入至2的最近幂。例如,如果token与其最近邻距离为3像素,则将其步幅赋值为

t_i

ti? = 4。然后,如果

(

)

∧

(

)

(x_i mod 2t_i = 0) ∧ (y_i mod 2t_i = 0)

(xi?mod2ti?=0)∧(yi?mod2ti?=0),我们设

g_i = 1

gi?=1，也就是说,我们希望将局部步幅从ti下采样到2

t_i

ti?。因此,如果

t_i = 1

ti?=1,则

g_i = 1

gi?=1对应交替像素,如果

t_i

ti? = 2,则

g_i = 1

gi?=1对应每4个像素。此外,我们在第j阶段将

(

)

∧

(

)

(x_i mod 2^{j+1} = 0) ∧ (y_i mod 2^{j+1} = 0)

(xi?mod2j+1=0)∧(yi?mod2j+1=0)的token的网格先验设置为无穷大。我们称这些token为“保留”。我们保留这些粗网格token,以确保前向传播期间图像中远程区域之间的连通性。
??总之,自适应下采样的工作流程如下:
1)为每个token fi获得重要性分数

(

)

S_i = σ(l(fi))

Si?=σ(l(fi)),
2)为每个token计算网格先验gi,
3)选取

g_i + alpha S_i

gi?+αSi?值最高的前x% token(例如1/4或1/5),
4)在前x% token的位置周围执行Eq. (4)给出的邻域合并公式,并获得下一阶段的合并token。