超越所有SOTA达11%！媲美全监督方法 | UC伯克利开源UnSAM

文章链接：https://arxiv.org/pdf/2406.20081

github链接：https://github.com/frank-xwang/UnSAM

SAM 代表了计算机视觉领域，特别是图像分割领域的重大进步。对于需要详细分析和理解复杂视觉场景(如自动驾驶、医学成像和环境监控)的应用特别有价值。今天和大家一起学习的是无监督SAM（UnSAM），用于能够即时启动并自动进行整体图像分割，而无需人工标注数据的方法。UnSAM利用分而治之的策略来“发现”视觉场景的层次结构。首先，利用自顶向下的聚类方法将未标注的图像分割成实例级和语义级的segments。对于每个segment内的所有像素，采用自底向上的聚类方法来迭代合并它们，从而形成更大的组，建立层次结构。这些无监督的多粒度masks随后用于监督模型训练。

在七个流行数据集上评估结果显示，UnSAM在语义平均召回率（AR）方面与有监督的SAM模型取得了竞争力的结果，并且在无监督分割领域超过了先前的最新技术，提升了11%。此外，还展示了有监督的SAM模型也能从本文的自监督标签中受益。通过将无监督伪masks集成到SA-1B的真实masks中，并仅使用SA-1B的1%进行训练，轻度半监督的UnSAM经常能够分割出有监督SAM忽视的实体，其在SA-1B数据集上的AR提升超过了6.7%，AP提升了3.9%。

整体亮点如下：

UnSAM的能力在七个主要的整体实体和部分分割数据集上进行了严格测试，例如MSCOCO、LVIS、SA-1B、ADE、Entity、PartImageNet和PACO。如下图1所示，展示了一些值得注意的表现：

• 无监督分割模型与SAM之间的性能差距可以显著缩小：通过在SA-1B中仅1%的未标注图像上使用ResNet50骨干网络进行训练，UnSAM不仅将无监督分割的最新技术水平提高了10%，还达到了与劳动密集的全监督SAM相媲美的性能。
  
• 监督式SAM也可以从我们的自监督标签中受益：将无监督伪mask与SA-1B的真实数据结合并在该组合数据上重新训练UnSAM，使得UnSAM+在AR方面超越SAM 6.7%以上，在AP方面超越3.9%。观察到UnSAM和UnSAM+经常能够发现被SAM遗漏的实体。
  

UnSAM: 无监督SAM

分而治之的分层图像分割

无监督分割任何内容模型从生成尊重视觉场景分层结构的伪masks开始，无需监督。这种方法受到的启发是，“分而治之”的策略是人类视觉系统用来高效处理和分析自然场景中复杂视觉信息的基本组织原则。伪masks生成pipeline “分而治之”，在算法 1 中总结，并在图 2 中说明，包括两个阶段：

分阶段：利用基于归一化切的方法，即 CutLER，从未标注的原始图像中获取语义和实例级别的masks。然而，由 CutLER 预测的粗粒度masks可能存在噪声。为了减少这种影响，过滤掉置信度低于阈值 $\tau$ 的masks。经验上，显著的语义和实例级别实体通常包含更丰富的部分级别实体（例如，一个人具有可识别的部分，如腿、手臂和头部，而背景天空则几乎没有或没有子级实体）。为了提取具有分层结构的这些部分级实体，使用Conquer stage。

Conquer阶段：对于在前一阶段发现的每个实例/语义级别masks，采用迭代合并来将粗粒度masks分解为更简单的部分，形成分层结构。

具体而言，首先使用在分阶段获得的masks裁剪局部patches，并双线性插值局部patches至分辨率为 256×256。然后，将它们送到预训练的 DINO ViT-B/8编码器 f(·) 中，并从最后的注意力层提取“关键”特征 $k_i=f(p_i)$ 作为局部patch $p_i$ 的patch特征。随后，Conquer stage采用迭代合并将patch分组为更大的clusters，具有预定义的余弦相似度阈值 θ ∈ { θ 1 , . . . , θ l } \theta \in \{\theta_1,...,\theta_l\} θ∈{θ1​,...,θl​}，其中 l 是预定义的粒度级别。

在第 t 次迭代中，本文的方法找到两个具有最高余弦相似度的相邻patch ($p_i,p_j$)，这两个patch分别来自于两个不同的clusters ($C_m^t,C_n^t$)，其余弦相似度为 $\frac{k_i^t{k}_j^t}{||k_i^t|{|}_2||k_j^t|{|}_2}$，将它们合并为一个clusters，并更新 $k_i^t$ 和 $k_j^t$ 为 $\frac{a_m{k}_i^t+a_n{k}_j^t}{a_m+a_n}$，其中 $a_m$ 是clusters $C_m^t$ 中的补丁数目（$p_i\in C_m^t$）。Conquer stage重复这一步骤，直到最大余弦相似度小于 ${\theta }_t$，将所有合并的clusters收集为新的部分级别伪masks，并使用较小的阈值 ${\theta }_{t+1}$ 再次进行迭代。每个在分阶段发现的粗粒度masks在Conquer stage后可以形成一个分层结构 H。

$n_t$ 是属于粒度级别 t 的clusters/masks数量，$n_0$ = 1。

masks合并：在Conquer stage发现的新的部分级别伪masks被添加回到分阶段确定的语义和实例级别伪masks中。然后，使用非极大值抑制（NMS）消除重复。在无监督图像分割的先前工作中，还采用诸如条件随机场（CRF）和CascadePSP等现成的masks精化方法，进一步改善伪masks的边缘。最后，在精化前后的交并比（IoU）显著差异的后处理masks被过滤掉。

初步结果：分而治之pipeline在伪masks池中实现了更多实体、更广泛的粒度级别和优越的质量，与之前的工作（如CutLER、U2Seg和SOHES）相比。如下表3所示，它在从SA-1B数据集中随机选择的1000张验证图像上达到了23.9%的语义平均召回率（AR），比现有技术提升了45.7%。

与先前伪masks生成工作的关键区别：UnSAM采用的分而治之策略使其区别于以往的工作：

[39, 28] 仅依赖于自顶向下的聚类方法，提供实例和语义级别的masks，从而错过了复杂图像中存在的分层结构。相反，本文的pipeline通过识别更细粒度的像素clusters捕捉了这种分层结构。

虽然 [6] 通过自底向上的迭代合并确实结合了一些分层结构，但仍然错过了许多细粒度实例和一些大尺度实例masks。此外，[6] 中的迭代合并专注于低于特定masks尺寸阈值的小区域，主要用于改进噪声小masks，限制了其检测各种实体尺寸的能力。本文的实验结果在质量和数量上表现出与先前工作相比的优越性能，特别是在生成更高质量、更详细的伪masks方面，更好地捕捉视觉场景的分层复杂性。

模型学习和自训练

尽管pipeline生成的伪masks在质量和数量上均优于先前的工作，但它们仍可能存在一定的噪声。自监督pipeline在识别某些类型的实例时存在局限性。例如，迭代合并有时无法正确关联同一实体的断开部分。为了解决这个问题，利用自训练策略进一步提升UnSAM模型的性能。UnSAM 使用分而治之策略发现的masks来学习图像分割模型。观察到自训练使模型能够“清洗”伪masks并预测更高质量的masks。一旦准备好伪masks，UnSAM 可以在模型学习或自训练阶段与任意的整体图像或可提示图像分割模型集成。

整体图像分割
为了简化，选择了基本的Masked Attention Mask Transformer（Mask2Former）。Mask2Former的关键创新在于在Transformer的交叉注意力块中引入了mask注意力机制，定义为$softmax(M+Q{K}^T)V$，其中注意力mask M在特征位置(x, y)处的定义如下：$M(x, y) =
\begin{cases}
0, M(x, y) = 1 \
-\infty, otherwise
\end{cases}
\tag{1}
$。这种机制限制了在预测mask区域内的注意力。随后，UnSAM 使用以下mask预测损失进行训练：

基于损失函数 $L_{ce}$ 和 $L_{dice}$，其中 ${\lambda }_{ce}$ 和 ${\lambda }_{dice}$ 是它们各自的权重。在伪masks上进行一轮自训练后，通过合并高置信度的masks预测（置信度大于 ${\tau }_{self-train}$）作为新的“ground-truth”标注来执行第二轮自训练。为避免重复，过滤掉与预测masks的IoU大于0.5的真值masks。

可提示图像分割
与SAM类似，本文的无监督SAM也可以根据输入提示（如点）生成高质量的对象masks。利用Semantic-SAM作为基础模型，从单击中预测多个粒度级别的masks。在学习过程中，在masks的内部圆圈$（半径\le 0.1\cdot min(Mas{k}_{width},Mas{k}_{height})$）内随机采样点，模拟用户的点击操作。

UnSAM+: 通过无监督分割改进监督SAM

监督SAM模型依赖于人工标注数据，这引入了一个显著的偏见，基于标注人员对“什么构成一个实例”的感知，经常会错过图像中的某些实体。相比之下，由于masks生成pipeline不依赖于人工监督，它通常可以识别SA-1B真值标注忽略的有效对象或部分。

受此观察的启发，利用UnSAM来改进监督SAM的性能，实施一种简单而有效的策略：基于IoU，将SA-1B的真值masks DSA-1B 与无监督分割masks DUnSAM 合并，制定如下：

${\tau }_{UnSAM+}$ 是IoU阈值，$Io{U}^{max}$ 是$C_m$ 和$D_{SA-1B}^i$ 中任何masks $C_n$ 之间的最大IoU，$D_{SA-1B}^i$ 和 $D_{UnSAM+}^i$ 分别是图像i中的SA-1B和无监督masks集合。

然后，对DUnSAM+ 进行训练，用于可提示图像分割和整体图像分割。融合方法利用了监督和无监督标注的优势，解决了人工注释数据集固有的限制，同时显著丰富了训练数据的多样性和全面性。这导致了更强大、更具普适性的分割模型 UnSAM+，超越了SAM的性能。

实验

模型训练设置

伪masks生成：在分阶段，设置置信度阈值$\tau =0.3$；在Conquer 阶段，选择阈值${\theta }_{merge}=[0.6,0.5,0.4,0.3,0.2,0.1]$。在将伪masks与真值合并用于训练UnSAM+时，选择${\tau }_{UnSAM+}=0.02$。

整体图像分割：UnSAM选择DINO预训练的ResNet-50作为骨干网络，Mask2former作为mask解码器。默认学习率为$5\times 10^{-5}$，batch大小为16，权重衰减为$5\times 10^{-2}$。训练模型8个epochs。

可提示图像分割：UnSAM使用自监督预训练的Swin-Transformer Tiny模型作为骨干网络，并利用Semantic-SAM 作为基础模型。设置层级数为6，这也是UnSAM在推断期间针对每个提示生成的预测masks数。对于所有实验，使用SA-1B数据集的1∼4%未标注图像进行训练。

评估数据集和指标

整体图像分割：以zero-shot方式在各种数据集上测试本文的模型，以评估从所有粒度级别分割实体的性能。选择 COCO、LVIS、ADE20K、EntitySeg、SA-1B（主要包含语义-/实例级别实体）、PartImageNet 和 PACO（覆盖部分级别实体）。SA-1B 测试集包括随机选择的1000张不包含在训练集中的图像。值得注意的是，每个数据集仅涵盖特定层次的实体和特定预定义的类别，而模型从所有层次和所有类别生成masks。因此，COCO 平均精度（AP）指标可能无法反映模型在分割开放世界中所有实体方面的真实性能。遵循之前的工作，主要考虑平均召回率（AR）来与不同模型进行比较。

基于点提示的可提示图像分割：在 COCO Val2017上评估基于点的交互式模型。根据之前关于可提示图像分割的工作，选择两个指标来评估模型：MaxIoU 和 OracleIoU。对于每个点提示，UnSAM 预测6个masks，代表不同的粒度级别。MaxIoU 计算6个masks中具有最高置信度分数的masks与真值之间的IoU，而OracleIoU 选择6个预测masks与真值之间的最高IoU。对于每个测试图像，选择其中心作为点提示。

评估结果

无监督伪masks：分而治之pipeline生成的无监督伪masks不仅包含精确的粗粒度实例masks，还捕获了SA-1B真值标签常常忽略的细粒度部分，如下图3所示。

整体图像分割：UnSAM在所有评估数据集上均优于现有技术，如前面表1所总结的。UnSAM表现出色，即使只使用SA-1B训练数据的1%，并且使用参数为2300万的ResNet-50骨干网络，而现有技术则使用两倍的训练数据和几乎四倍的参数。这表明UnSAM是一个轻量级、更易于训练且对数据需求较少的模型，在分割开放世界中的实体时具有更好的zero-shot性能，如下图4和5所示。

平均而言，UnSAM在AR方面超过先前的SOTA达11.0%。在PartImageNet 和 PACO 上评估时，UnSAM分别超过先前的SOTA达16.6%和12.6%。与监督SAM相比，UnSAM在所有数据集上的AR差距仅为1%。在PartImageNet 和 PACO上，UnSAM分别比SAM超过24.4%和11.6%。这进一步展示了分而治之pipeline在发现人类标注者往往会忽略的细节方面的出色能力。

此外，UnSAM+，通过整合无监督伪masks和SA-1B真值进行训练，表现优于SAM的AR超过6.7%，AP超过3.9%，如前面表2和4所示。UnSAM+ 在所有评估数据集上的平均召回率优于SAM，除了ADE20K，这是由语义级别标注主导的数据集。UnSAM+ 在小实体上显著比SAM高出16.2%的AR进一步证实，伪masks可以有效地补充SA-1B数据集忽略的细节，而UnSAM+ 可以常常发现SAM错过的实体，如上面图4所示。

基于点提示的可提示图像分割：如下表5所示，使用伪masks训练的UnSAM在COCO上达到了40.3%的MaxIoU和59.5%的OracleIoU。值得注意的是，只使用了SAM 使用数据的1%以及一个具有4倍少参数的骨干网络进行训练。此外，使用整合的伪masks和SA-1B真值进行训练的UnSAM+ 在MaxIoU和OracleIoU上分别比SAM优于0.3%和1.3%。定性结果见下图6。

总结

图像分割是计算机视觉中的基本任务，传统上依赖于密集的人工标注来详细理解视觉场景。本文提出了UnSAM，一个无监督分割模型，在无监督图像分割的性能上显著超越了先前的最新技术。此外，无监督UnSAM模型取得了令人印象深刻的结果，与领先的监督SAM的性能相媲美，并在某些半监督设置下超越它。

参考文献

[1] Segment Anything without Supervision