Segment Anything Model(SAM)论文解读

一、引言

在这项工作中，作者的目标是建立一个图像分割的基础模型。也就是说，寻求开发一个提示模型，并使用一个能够实现强大泛化的任务在广泛的数据集上对其进行预训练。有了这个模型，使用即时工程解决新数据分布上的一系列下游分割问题。
该计划的成功取决于三个组成部分:任务、模型和数据。为了开发它们，作者解决了以下关于图像分割的问题:
1、什么任务可以实现zero-shot泛化?
2、相应的模型体系结构是什么?
3、哪些数据可以为这项任务和模型提供支持?
这些问题纠缠在一起，需要综合解决。作者首先定义了一个提示的分割任务，它足够通用，可以提供一个强大的预训练目标，并支持广泛的下游应用程序。此任务需要一个支持灵活提示的模型，并且可以在提示时实时输出分段掩码以允许交互使用。为了训练模型，需要一个多样化的、大规模的数据源。不幸的是，没有网络规模的数据来源的分割;为了解决这个问题，作者构建了一个“数据引擎”，也就是说，在使用高效模型来协助数据收集和使用新收集的数据来改进模型之间进行迭代。
**任务：**在NLP和最近的计算机视觉中，基础模型是一个很有前途的发展，它可以通过使用“提示”技术对新的数据集和任务执行zero-shot和few-shot学习。受这一行工作的启发，作者提出了提示分割任务，其目标是给定任何分割提示返回有效的分割掩码(见图1a)。提示符只是指定图像中要分割的内容，例如，提示符可以包括识别对象的空间或文本信息。即使提示是模糊的，并且可能引用多个对象(例如，衬衫上的一个点可能表示衬衫或穿着它的人)，输出也应该是这些对象中至少一个对象的合理掩码。使用提示分割任务作为预训练目标，并通过提示工程解决一般的下游分割任务。
**模型：**可提示的分割任务和实际使用的目标对模型体系结构施加了约束。特别是，模型必须支持灵活的提示、需要实时计算掩码，以实现交互式使用。
而且必须具有模糊感知能力。令人惊讶的是，作者发现一个简单的设计满足所有三个约束:一个强大的图像编码器计算图像嵌入，一个提示编码器嵌入提示，然后将两个信息源组合在一个轻量级的掩码解码器中，该解码器预测分割掩码。作者将此模型称为分段任意模型(Segment Anything model，简称SAM)(见图1b)。通过将SAM分为图像编码器和快速提示编码器/掩码解码器，可以使用不同的提示重复使用相同的图像嵌入(并平摊其成本)。给定图像嵌入，提示编码器和掩码解码器在50ms内从web浏览器中的提示预测掩码。将重点放在点、框和掩码提示上，并使用自由格式的文本提示来呈现初始结果。为了使SAM能够感知歧义，将其设计为预测单个提示的多个掩码，从而允许SAM自然地处理歧义，例如衬衫与人的例子。

**数据引擎：**为了实现对新数据分布的强泛化，有必要在一个大而多样的掩码集上训练SAM，而不仅仅是已经存在的任何分割数据集。虽然基础模型的典型方法是在线获取数据，但掩模并不自然丰富，因此需要一种替代策略。作者的解决方案是建立一个“数据引擎”，也就是说，与模型在循环数据集注释共同开发模型(见图1c)。数据引擎有三个阶段:辅助手动、半自动和全自动。在第一阶段，SAM帮助注释者注释掩码，类似于经典的交互式分段设置。在第二阶段，SAM可以通过提示可能的对象位置来自动为对象子集生成掩码，而注释器则专注于注释剩余的对象，从而帮助增加掩码的多样性。在最后阶段，用前景点的规则网格提示SAM，平均每张图像产生100个高质量掩模。

**数据集：**最终数据集SA-1B包括来自11M张授权和隐私保护图像的超过1B个掩码(见上图)。SA-1B是使用数据引擎的最后阶段完全自动收集的，比任何现有的分割数据集都多400个掩码，并且经过验证，掩码具有高质量和多样性。除了将SA-1B用于训练SAM具有鲁棒性和通用性之外，作者希望SA-1B成为旨在建立新基础模型的研究的宝贵资源。
实验： 作者广泛评估SAM。首先，使用23个不同的新分割数据集，发现SAM从单个前景点产生高质量的掩模，通常仅略低于手动注释的真值。其次，在zero-shot传输协议下使用即时工程的各种下游任务上发现了一致的强定量和定性结果，包括边缘检测，目标提案生成，实例分割以及文本到掩码预测的初步探索。这些结果表明，SAM可以使用开箱即用的快速工程来解决涉及SAM训练数据之外的对象和图像分布的各种任务

二、Segment Anything Task

作者从NLP中获得灵感，其中令牌预测下一个任务用于基础模型预训练，并通过提示工程解决各种下游任务。为了建立分割的基础模型，作者的目标是定义一个具有类似功能的任务。
任务首先将提示的概念从NLP翻译到分割，其中提示可以是一组前景/背景点，一个粗略的框或蒙版，自由格式的文本，或者一般情况下，任何指示图像中要分割的信息。因此，提示分段任务是在给定任何提示的情况下返回一个有效的分段掩码。“有效”掩码的要求仅仅意味着，即使提示是模糊的，并且可以引用多个对象(例如，回想一下衬衫与人的例子，参见下图)，输出也应该是这些对象中至少一个的合理掩码。这个需求类似于期望语言模型对不明确的提示输出一致的响应。之所以选择这个任务，是因为它带来了一种自然的预训练算法和一种通过提示将zero-shot转移到下游分割任务的通用方法。

预训练提示分割任务提出了一种自然的预训练算法，该算法为每个训练样本模拟一系列提示(例如，点、框、掩码)，并将模型的掩码预测与基本事实进行比较。从交互式分割中采用了这种方法，尽管与交互式分割不同，交互式分割的目的是在足够的用户输入后最终预测一个有效的掩码，但作者的目标是始终预测任何提示的有效掩码，即使提示是模糊的。这确保了预训练模型在涉及歧义的用例中是有效的，包括数据引擎所要求的自动注释。在这个任务中表现良好是具有挑战性的，需要专门的建模和训练损失选择。
Zero-shot转移直观地说，预训练任务赋予了模型在推理时对任何提示作出适当响应的能力，因此下游任务可以通过设计适当的提示来解决。例如，如果有一个猫的边界框检测器，猫实例分割可以通过提供检测器的框输出作为提示给模型来解决。一般来说，许多实际的分割任务都可以作为提示。
相关的任务分割是一个很广阔的领域，有交互式分割、边缘检测、超像素化、目标建议生成、前景分割、语义分割、实例分割、全视分割等。提示分割任务的目标是产生一个功能广泛的模型，可以通过快速工程适应许多(尽管不是全部)现有的和新的分割任务。这种能力是任务泛化的一种形式。请注意，这与之前在多任务分割系统上的工作不同。在多任务系统中，单个模型执行一组固定的任务，例如联合语义分割、实例分割和全视分割，但训练和测试任务是相同的。本文是训练用于提示分割的模型可以作为更大系统中的组件在推理时间执行新的不同任务，例如，执行实例分割，提示分割模型与现有的对象检测器相结合。

三、Segment Anything Model

SAM有三个组件，如下图所示:一个图像编码器，一个灵活的提示编码器和一个快速掩码解码器。建立在Transformer视觉模型的基础上，对(平摊)实时性能进行了特定的权衡。
图像编码器 在可扩展性和强大的预训练方法的激励下，使用了MAE预训练的视觉变压器(ViT)，以最小程度适应处理高分辨率输入。图像编码器每个图像运行一次，可以在提示模型之前应用。
提示编码器考虑两组提示:稀疏(点、框、文本)和密集(掩码)。通过位置编码来表示点和框，并对每个提示类型和使用CLIP的现成文本编码器的自由格式文本进行学习嵌入求和。密集提示(即掩码)使用卷积嵌入，并在图像嵌入中按元素求和。
掩膜解码器 掩膜解码器有效地将图像嵌入、提示嵌入和输出令牌映射到掩码。该设计采用了对Transformer解码器块的修改，然后是动态掩码预测头。改进的解码器块在两个方向上使用提示自注意和交叉注意(提示到图像嵌入，再由嵌入到图像)来更新所有嵌入。在运行两个块之后，对图像嵌入进行上采样，MLP将输出标记映射到动态线性分类器，然后该分类器计算每个图像位置的掩码前景概率。
解决歧义对于一个输出，如果给出一个模糊的提示，该模型将平均多个有效掩码。为了解决这个问题，修改了模型，以预测单个提示符的多个输出掩码(见下图)。3个掩码输出足以解决大多数常见情况(嵌套掩码通常最多有三个深度:整体、部分和子部分)。在训练中，为了对掩码进行排序，该模型预测每个掩码的置信度得分（即估计的IoU）。

效率 整个模型的设计很大程度上是出于效率的考虑。给定预先计算的图像嵌入，提示编码器和掩码解码器在网络浏览器中运行，在CPU上，大约50ms。这种运行时性能使模型能够无缝、实时地交互提示。
损失和训练 使用的focal loss和dice loss的线性组合来监督掩模预测。使用几何提示的混合来训练可提示的分割任务。

四、 Segment Anything Data Engine

由于互联网上的分割掩码并不丰富，作者建立了一个数据引擎来收集1.1亿掩码数据集SA-1B。数据引擎有三个阶段：（1）模型辅助的手动注释阶段（2）混合了自动预测掩码和模型辅助注释的半自动阶段，以及（3）模型在没有注释器输入的情况下生成掩码的全自动阶段。
辅助手动阶段在第一阶段，类似于经典的交互式分割，一组专业注释人员通过使用SAM提供的基于浏览器的交互式分割工具点击前景/背景对象点来标记mask。mask可以使用像素级“画笔”和“橡皮擦”工具进行细化。模型辅助注释直接在浏览器中实时运行（使用预先计算的图像嵌入），从而实现真正的交互式体验。作者没有对标记对象施加语义约束，注释器可以自由地标记“东西”和“事物”。作者建议注释器标记他们可以命名或描述的对象，但没有收集这些名称或描述。注释者被要求按照突出的顺序标记对象，并被鼓励在mask注释超过30秒后继续下一张图像。
在这个阶段开始时，SAM是使用公共分割数据集进行训练的。在充分的数据注释之后，仅使用新注释的掩码对SAM进行再训练。随着更多掩模的收集，图像编码器从ViT-B扩展到ViT-H，其他架构细节也在发展；作者总共对模型进行了6次训练。随着模型的改进，每个掩码的平均注释时间从34秒减少到14秒。14秒比COCO的掩码注释快6.5倍，仅比使用极值点的边界框标记慢2倍。随着SAM的改进，每张图像的平均掩模数量从20个增加到44个。总的来说，在这个阶段从120k张图像中收集了430万个mask。
半自动阶段 在这个阶段，作者的目标是增加mask的多样性，以提高模型分割任何东西的能力。为了将注释器集中在不太突出的对象上，首先自动检测到自信的掩码。然后，向注释器展示了预先填充了这些掩码的图像，并要求他们注释任何其他未注释的对象。为了检测有信心的掩码，使用通用的“对象”类别在所有第一阶段掩码上训练了一个边界框检测器。在此阶段，在180k张图像中额外收集了590万个掩模（总共1020万个掩膜）。与第一阶段一样，定期根据新收集的数据对模型进行再训练（5次）。每个掩码的平均注释时间回到了34秒（不包括自动掩码），因为这些对象更难标记。每张图像的平均mask数量从44个增加到72个（包括自动mask）。
全自动阶段 在最后阶段，注释是完全自动的。这是可行的，因为模型有两个主要的增强。首先，在这个阶段开始时，作者收集了足够的mask来大大改进模型，包括前一阶段的各种mask。其次，到了这个阶段，已经开发了模糊感知模型，它能够预测有效的掩码，即使在模糊的情况下也是如此。具体来说，用32×32的规则网格提示模型，并为每个点预测一组可能对应于有效对象的掩码。使用模糊感知模型，如果一个点位于部分或子部分上，模型将返回子部分、部分和整个对象。模型的IoU预测模块用于选择置信掩码；此外，只识别并选择了稳定的掩码（如果在0.5−δ和0.5+δ处对概率图进行阈值处理会导致类似的掩码，则认为掩码是稳定的）。最后，在选择了置信和稳定的掩码后，应用非最大抑制（NMS）来过滤重复。为了进一步提高较小mask的质量，还处理了多个重叠的放大图像裁剪。

五、网络结构细节：

图像编码器 通常，图像编码器可以是输出C×H×W图像嵌入块的任何网络。受可扩展性和强大的预训练的启发，作者使用MAE预训练的视觉transformer（ViT），具有最小的适应能力来处理高分辨率输入，特别是具有14×14窗口注意力和四个等距全局注意力块的ViT-H/16。图像编码器的输出是输入图像的16倍缩小的嵌入。由于运行时目标是实时处理每个提示，因此可以提供大量的图像编码器FLOP，因为每个图像只计算一次FLOP，而不是每个提示。
根据标准实践，使用1024×1024的输入分辨率，该分辨率是通过重新缩放图像并填充短边获得的。因此，图像嵌入是64×64。为了降低通道维度，使用1×1卷积来获得256个通道，然后使用3×3卷积来获得同样的256个通道。每个卷积后面都有一个层归一化。

提示编码器 稀疏提示被映射到256维矢量嵌入。一个点被表示为该点的位置的位置编码和指示该点是在前景中还是在背景中的两个学习嵌入之一的总和。方框由嵌入对表示：（1）其左上角的位置编码与表示“左上角”的学习嵌入相加；（2）相同的结构，但使用表示“右下角”的习得嵌入。最后，为了表示自由形式的文本，使用CLIP中的文本编码器（通常可以使用任何文本编码器）。密集提示（即掩码）与图像具有空间对应关系。以比输入图像低4倍的分辨率输入掩码，然后使用两个分别具有输出通道4和16的2×2，跨步-2卷积来缩小另外的4倍。最后的1×1卷积将通道维度映射到256。每一层通过GELU激活和层标准化进行分离。然后将掩码逐元素添加图像嵌入。如果没有掩码提示，则向每个图像嵌入位置添加表示“无掩码”的学习嵌入。

轻量级掩码解码器 该模块有效地将图像嵌入和一组提示嵌入映射到输出掩码。为了组合这些输入，从Transformer分割模型中获得灵感，并修改了标准Transformer解码器。在应用解码器之前，首先在提示嵌入集中插入一个学习的输出令牌嵌入，该嵌入将在解码器的输出中使用。为了简单起见，我们将这些嵌入（不包括图像嵌入）统称为“令牌”

我们的解码器设计如上图所示。每个解码器层执行4个步骤：（1）对令牌的自注意力（2）从令牌（作为查询）到图像嵌入的交叉注意力（3）逐点MLP更新每个令牌，以及（4）从图像嵌入（作为查询的）到令牌的交叉关注。最后一步使用提示信息更新图像嵌入。在交叉关注期间，图像嵌入被视为642个256维向量的集合。每个自/交叉注意力和MLP在训练时都有残差连接、层归一化和dropout。下一解码器层从上一层获取更新的令牌和更新的图像嵌入。使用两层解码器。
为了确保解码器能够访问关键的几何信息，每当位置编码参与注意力层时，都会将其添加到图像嵌入中。此外，当整个原始提示标记（包括它们的位置编码）参与注意力层时，它们都会被重新添加到更新的标记中。这允许强烈依赖提示标记的几何位置和类型。
在运行解码器之后，用两个转置卷积将更新的图像嵌入上采样4×（现在它相对于输入图像缩小了4倍）。然后，令牌再次参与图像嵌入，并且将更新的输出令牌嵌入传递给小的3层MLP，该MLP输出与放大图像嵌入的通道维度匹配的向量。最后，预测了在放大图像嵌入和MLP的输出之间具有空间逐点乘积的掩模。转换器使用256的嵌入尺寸。转换器MLP块具有2048的大尺寸，但MLP仅应用于相对较少（很少大于20）的提示令牌。然而，在具有64×64图像嵌入的交叉关注层中，为了提高计算效率，将查询、键和值的通道维度降低了2×至128。所有注意力层使用8个头。用于升级输出图像嵌入的转置卷积是2×2，步长2，输出通道尺寸为64和32，并且具有GELU激活。它们通过层规范化来分隔。
Making the model ambiguity-aware 如上所述，单个输入提示可能是模糊的，因为它对应于多个有效掩码，并且模型将学习在这些掩码上求平均值。作者通过一个简单的修改来消除这个问题：使用少量输出令牌并同时预测多个掩码，而不是预测单个掩码。默认情况下，预测三个掩码，因为作者观察到三个层（整体、部分和子部分）通常足以描述嵌套掩码。在训练过程中，计算真值和每个预测掩码之间的损失，但仅从最低损失进行反向传播。这是一种常见的技术，用于具有多个输出的模型。为了在应用程序中使用，希望对预测掩码进行排序，因此我们添加了一个小头（对额外的输出令牌进行操作），用于估计每个预测掩码与其覆盖的对象之间的IoU。
多个提示的歧义要少得多，三个输出掩码通常会变得相似。为了最大限度地减少训练中退化损失的计算，并确保单个无模糊掩码接收规则梯度信号，只在给出多个提示时预测单个掩码。这是通过添加用于附加掩码预测的第四输出令牌来实现的。第四个掩码从不为单个提示返回，也是为多个提示返回的唯一掩码。