CVPR2022: Self-Supervised Learning of Object Parts for Semantic Segmentation

Abstract

自监督学习的进步带来了强大的图像表示学习方法。然而到目前为止，它主要集中在图像级学习上。反过来，无监督图像分割等任务并没有从这一趋势中受益，因为它们需要空间多样化的表示。在本文中，我们认为对象部分的自我监督学习是解决这个问题的方法。

尽管自监督 ViT 擅长提取对象，但它们并没有学习到空间标记嵌入空间，即区分不同的对象类别。为此, 我们提出了一个密集聚类 pretext task 来学习语义丰富的空间标记，缩小监督 ViT 和自我监督 ViT 之间的差距。

我们探索使用 Vision Transformer (VIT) 和我们的新损失，使其完成自我监督的分割。虽然视觉转换器显示出巨大的潜在统一架构，并且可以很好地将数据扩展到数十亿，但它们主要被证明适用于以监督方式的图像级任务或密集任务。我们的工作旨在通过自监督学习密集的 VIT 模型来缩小这一差距。

Method

我们的目标是学习一个嵌入空间，该空间对包含对象相同部分的图像块进行分组。原因是我们假设对象的部分表示比对象表示更通用, 因为其可以以多种方式重新组合到不同的对象。例如，轮子表示可以组合成汽车表示，也可以组合公共汽车表示。

我们没有使用需要大量 GPU 预算的从头开始预训练模型，而是使用我们的损失来微调预训练神经网络 DINO。

Untitled

Fine-tuning loss for spatial tokens

Untitled

为了训练 vit 网络提取空间标记，我们首先将图像随机裁剪 $V$ 次, 每次得到 $v_g$ global view和 $v_l$ local view。给定一个图像 $x ∈\mathbb{R}^{3×H×W}$, 那么剪裁后的图像表示为 $x_{t_j} , j ∈ \{1,., V \}$
使用一个 MLP + L2 Norm 作为 bottleneck 进行提取特征, 得到 $g \left( f \left( x _ { t _ { j } } \right) \right) = Z _ { t _ { j } } \in \mathbb { R } ^ { D \times N }$

<aside> 💡 我们使用的视觉编码器是 ViT ，它将图像块 $x_i$ 映射到 $N$ 个空间标记的向量 $f (x) = [f (x_1), . . . f(x_N)]$。

</aside>
接下来初始化每个图像块特征 $Z$ 的 $K$ 个聚类中心, $\left[ c _ { 1 } , \ldots , c _ { K } \right] = C \in \mathbb { R } ^ { D \times K }$, 然后使用 OTA 最优传输匹配, 计算每个图像块特征 $Z_{t_j}$ 属于每个类别的概率 $Q_{t_j}$, 由于全局作物捕获了图像的大部分，我们仅使用 $v_g$ 它们来计算 $Q_{tj}$ ，因为空间标记可以关注全局场景信息，从而提高整体预测目标质量。
为了将 $v_g$ 与 $v_l$ 的区域对齐，$α(·)$ 使用 ROI Align 将 $v_g$ 与 $v_l$ 对齐.
然后, 损失为:

$$ L \left( x _ { t _ { 1 } } , \ldots , x _ { t _ { V } } \right) = \sum _ { j = 0 } ^ { v _ { g } } \sum _ { i = 0 } ^ { V } \mathbb { 1 } _ { k \neq j } l \left( x _ { t _ { i } } , x _ { t _ { j } } \right) \\ l \left( x _ { t _ { i } } , x _ { t _ { j } } \right) = H \left[ \left( s _ { \tau } \left( \alpha _ { B _ { j , i } } \left( g \left( \Phi \left( x _ { t _ { i } } \right) \right) ^ { T } C \right) , \alpha _ { B _ { i j } } \left( Q _ { t _ { j } } \right) \right] \right. \right. $$

其中 $H$ 是交叉熵，$s_τ$ 是按温度 $τ$ 缩放的 softmax。我们使用 $L$ 联合优化类别向量 $C$ 以及神经网络 $f$ 和 $g$。 $C$ 在每个梯度步骤之后进一步 $L2$ 归一化，使得 $Z^T C$ 为图像块特征与聚类中心之间的余弦相似度。

<aside> 💡 类似于 Self-Supervised Visual Representation Learning with Semantic Grouping. 但这里的 Q 是由 OTA 匹配的来的, 而另一篇文章的 Q 计算方式与 student 网络一致.

</aside>