摘要
最近的 Segment Anything Model(SAM)在扩展分割模型规模方面取得了重大突破,具备强大的零样本能力和灵活的提示机制。尽管 SAM 在训练时使用了 11 亿个掩码,其掩码预测质量在许多情况下仍不理想,尤其是对于结构复杂的目标。我们提出了 HQ-SAM,使 SAM 能够精确地分割任意目标,同时保留其原有的可提示设计、高效性和零样本泛化能力。我们的设计充分复用并保留了 SAM 预训练的模型权重,仅引入极少量的附加参数和计算。我们设计了一个可学习的“高质量输出令牌”(High-Quality Output Token),将其注入 SAM 的掩码解码器,用于预测高质量掩码。为了提升掩码细节,我们不仅在解码器特征上使用该令牌,还将其与 ViT 编码器的早期和最终特征进行融合。为训练新增的可学习参数,我们收集并整合了来自多个来源的 44K 张精细掩码数据集。HQ-SAM 仅在这 44K 张掩码数据上训练,使用 8 块 GPU 仅耗时 4 小时。我们在 10 个多样化的分割数据集上评估了 HQ-SAM 的性能,其中 8 个数据集采用零样本转移协议。实验证明,HQ-SAM 在保持零样本能力的同时,显著提升了掩码质量。代码和预训练模型可在 https://github.com/SysCV/SAM-HQ 获取。
为了让 HQ‑Output Token 能够生成更精细的掩码边缘与细节,我们在掩码解码器的特征之上,又额外引入了 ViT 编码器的“早期特征”和“后期特征”两种信息:
早期特征(High‑Resolution 细节信息)
ViT 在最开始几层处理的 patch 嵌入保留了较高的空间分辨率(相当于 CNN 中浅层的 feature map),能够捕捉到物体边缘、纹理、细线等局部细节。
如果只用解码器输出的特征,往往因多次下采样/Transformer 自注意力而丢失这些精细结构。
因此,我们取自 ViT 编码器前几层的高分辨率 feature map,经一次小型映射(如 1×1 卷积或线性层)降维后,与 HQ‑Output Token 一起送入特征融合模块。
后期特征(全局语义信息)
ViT 编码器最后几层的输出具有强大的语义表达能力,能全局感知“这块区域是哪个物体,属于哪个类别”,但空间分辨率相对较低。
将这些深层特征融入,可以帮助 HQ‑Output Token 准确判断哪些细节应当被保留、哪些应被忽略,从而避免误分割噪声或背景纹理。
特征融合机制
我们设计了一个轻量级的 “Fusion Block”:首先对三组特征(早期、解码器中期、后期)分别做投影到相同维度,然后在 token 维度上进行拼接或加权相加,最后通过一层小型 MLP(或卷积)融合出一组既有高分辨率细节又有全局语义的特征表示。
HQ‑Output Token 在 mask 解码器中,就基于这组融合特征进行自注意力操作和 MLP 预测,从而在原来粗糙的掩码边界基础上,补上细线、缝隙、曲面等复杂结构。
通过以上“多尺度、跨阶段”的特征融合,HQ‑SAM 能在保持 SAM 原有 promptable 设计与零样本泛化能力的同时,显著改善对细小结构和边缘细节的分割质量。
1 引言
准确分割多样化目标对于图像/视频编辑、机器人感知和 AR/VR 等场景理解应用至关重要。SAM [21] 作为一款基于海量掩码标签训练的通用图像分割基础模型,以点、边框或粗略掩码作为输入提示,能够在多种场景下分割多样化目标、部件和视觉结构,其零样本分割能力正引发范式转变。
尽管 SAM 在多项任务中表现优异,其分割结果在许多场景下仍不尽如人意,主要存在两大问题:
掩码边界粗糙,甚至遗漏细小结构(如图 1 中的风筝线);
分割错误或遮断掩码,在挑战性场景下出现较大偏差。
这些问题严重限制了基础分割模型在自动标注和图像/视频编辑等应用中的实用性,因为此类任务对掩码的高精度有严格需求。
我们提出 HQ-SAM,在不损害 SAM 零样本性能和提示灵活性的前提下,实现高质量分割。直接微调 SAM 解码器或新增完整解码器都会显著削弱其零样本泛化能力;为此,我们设计了一种与原 SAM 紧密集成的轻量化架构。首先,新增一个可学习的 HQ-输出令牌(HQ-Output Token),与原提示令牌一并输入到 SAM 的掩码解码器;该令牌及其附属的 MLP 层专门用于生成高质量掩码。其次,该令牌不仅作用于解码器特征,还通过特征融合模块,与 ViT 编码器的初始和末端特征相结合,以兼顾全局语义和局部细节。在训练阶段,我们冻结所有 SAM 预训练参数,仅更新 HQ-输出令牌、三层 MLP 以及小规模特征融合模块。
训练优质分割模型需要多样化且精确的掩码标注。SAM 原训练集 SA-1B 包含 1100 万张图像与 11 亿个自动生成的掩码,但数据规模巨大且标注质量不足以满足我们对高精度掩码的需求。因此,我们构建了 HQSeg‑44K 数据集,汇集了 44K 张极致精细的掩码标注,覆盖 1000 多个多样化语义类别,来源于六个现有高质量掩码数据集 [