Segment Anything
paper/code: https://segment-anything.com/
Abstract
我们介绍了 Segment Anything (SA) 项目：一个用于图像分割的新任务、模型和数据集。通过在数据收集循环中使用我们的高效模型，我们建立了迄今为止最大的分割数据集（迄今为止），其中包含 1100 万张授权图像上的 10 亿多个掩码，并且尊重隐私。该模型的设计和训练具有可提示性，因此它可以在新的图像分布和任务中进行零转移。我们在大量任务中评估了它的能力，发现它的零镜头性能令人印象深刻–通常可与之前的完全监督结果相媲美，甚至更胜一筹。我们在 Segment Anything 上发布了 Segment Anything Model（SAM）和相应的数据集（SA-1B），其中包括 1B 个遮罩和 1,100 万张图像，以促进计算机视觉基础模型的研究。
1.Introduction
通过NLP中大语言模型，引出CV中的也可以实现大模型：
在网络规模的数据集上预先训练的大型语言模型正以强大的zero-shot泛化和few-shot泛化能力为NLP带来革命性的变化。这些 "基础模型"可以泛化到训练期间所见之外的任务和数据分布。这种能力通常是通过提示工程来实现的，即使用手工制作的文本来提示语言模型为当前任务生成有效的文本响应。当使用来自网络的大量文本语料进行扩展和训练时，这些模型的零点和少点性能与经过微调的模型相比（在某些情况下甚至不相上下）令人惊讶。经验趋势表明，随着模型规模、数据集大小和总训练计算量的增加，这种性能也在提高。
计算机视觉领域也对基础模型进行了探索，尽管程度较低。例如，CLIP和 ALIGN使用对比学习来训练文本和图像编码器，使两种模式对齐。一旦经过训练，工程化文本提示就能实现对新的视觉概念和数据分布的零点泛化。这种编码器还能与其他模块有效组合，以完成图像生成等下游任务（如 DALL-E）。虽然在视觉和语言编码器方面已经取得了很大进展，但计算机视觉还包括许多超出这一范围的问题，而且其中许多问题并不存在丰富的训练数据。
在这项工作中，我们的目标是为图像分割建立一个基础模型。也就是说，我们试图开发一个可提示的模型，并在一个广泛的数据集上使用可实现强大泛化的任务对其进行预训练。有了这个模型，我们就能在新的数据分布上利用提示工程解决一系列下游分割问题。 这一计划的成功取决于三个要素：任务、模型和数据。为了开发这三个部分，我们要解决以下有关图像分割的问题：
什么任务能实现zero-shot泛化？相应的模型架构是什么？哪些数据可以支持这项任务和模型？
这些问题错综复杂，需要全面的解决方案。我们首先定义了一个可提示的分割任务，它的通用性足以提供强大的预训练目标，并支持广泛的下游应用。这项任务需要一个支持灵活提示的模型，并能在提示时实时输出分割掩码，以便交互使用。为了训练我们的模型，我们需要多样化、大规模的数据源。 不幸的是，目前还没有网络规模的分割数据源；为了解决这个问题，我们建立了一个 “数据引擎”，即在使用我们的高效模型来协助数据收集和使用新收集的数据来改进模型之间进行迭代。接下来，我们将介绍每个相互关联的组件，然后是我们创建的数据集和证明我们方法有效性的实验。
介绍论文每一节详细讲解的内容：
任务（§2）：在 NLP 以及最近的计算机视觉领域，基础模型是一项前景广阔的发展，它通常通过使用 "提示 "技术，对新数据集和任务进行 zero-shot或 few-shot学习。受此启发，我们提出了可提示分割任务，其目标是在给出任何分割提示的情况下返回有效的分割掩码（见图 1a）。提示只需指定图像中要分割的内容，例如，提示可以包括识别物体的空间或文本信息。对有效输出掩码的要求意味着，即使提示模棱两可，可能指向多个对象（例如，衬衫上的一个点既可能指向衬衫，也可能指向穿着衬衫的人），输出也应至少是其中一个对象的合理掩码。我们将可提示分割任务作为预训练目标，并通过提示工程解决一般的下游分割任务。
模型（§3）：可提示的分割任务和实际使用的目标对模型结构提出了限制。特别是，模型必须支持灵活的提示，需要实时计算掩码以允许交互式使用，并且必须具有模糊感知能力。令人惊讶的是，我们发现一个简单的设计就能满足所有三个限制条件：一个功能强大的图像编码器计算图像嵌入，一个提示编码器嵌入提示，然后将这两个信息源结合到一个轻量级掩码解码器中，预测分割掩码。我们将这一模型称为 “任何分割模型”（Segment Anything Model）或 SAM（见图 1b）。通过将 SAM 分离为图像编码器和快速提示编码器/掩码解码器，相同的图像嵌入可以在不同的提示中重复使用（并摊销其成本）。 给定图像嵌入后，提示编码器和掩码解码器可以在网络浏览器中根据提示在 50ms 内预测出掩码。我们的重点是点、方框和掩码提示，同时也展示了自由格式文本提示的初步结果。为了使 SAM 具有歧义感知能力，我们设计了它来预测单个提示的多个掩码，使 SAM 能够自然地处理歧义，例如衬衫与人的例子。
数据引擎（§4）:为了实现对新数据分布的强泛化，我们发现有必要在大量不同的掩码集上对 SAM 进行训练，而不局限于已有的任何分割数据集。虽然基础模型的典型方法是在线获取数据，但掩码并不天然丰富，因此我们需要一种替代策略。我们的解决方案是建立一个 “数据引擎”，即通过模型在环数据集注释共同开发我们的模型（见图 1c）。我们的数据引擎分为三个阶段：辅助手动、半自动和全自动。在第一阶段，SAM 协助注释者注释掩码，这与经典的交互式分割设置类似。在第二阶段，SAM 可以通过提示可能的对象位置来自动生成对象子集的mask，而注释者则专注于注释其余的对象，这有助于提高mask的多样性。在最后阶段，我们向 SAM 提示前景点的规则网格，平均每幅图像可生成 100 个高质量的mask。
数据集（§5）:我们的最终数据集 SA-1B 包括来自 1100 万张授权和隐私保护图像的超过 1B 个掩码。SA-1B 是利用我们数据引擎的最后阶段全自动收集的，比任何现有的分割数据集都多 400 倍的掩码，而且我们广泛验证了这些掩码的高质量和多样性。除了用于训练 SAM 的鲁棒性和通用性外，我们还希望 SA-1B 成为旨在建立新基础模型的研究的宝贵资源。
Responsible AI（§6）：
实验（§5）：我们对 SAM 进行了广泛的评估。首先，我们使用一套新的 23 个分割数据集，发现 SAM 能从单个前景点生成高质量的遮罩，通常仅略低于人工标注的地面实况。其次，我们发现，在zero-shot传输协议下，SAM 在各种下游任务（包括边缘检测、对象建议生成、实例分割以及文本到掩码预测的初步探索）上都取得了一致的优异定量和定性结果。这些结果表明，除了 SAM 的训练数据之外，SAM 还能与即时工程一起用于解决涉及对象和图像分布的各种任务。尽管如此，仍有改进的余地，我们将在第8节中讨论。
2. Segment Anything Task
讲解SAM中的提示工程的由来，并引出和介绍SAM的一系列任务：
我们从 NLP 中汲取灵感，在 NLP 中，下一个标记预测任务被用于基础模型的预训练，并通过提示工程（prompt engineering）解决各种下游任务。为了建立一个用于分割的基础模型，我们的目标是定义一个具有类似能力的任务。
Task：我们首先将提示的概念从 NLP 转化为分割，提示可以是一组前景/背景点、一个粗略的方框或mask、自由形式的文本，或者，一般来说，指示图像中要分割的内容的任何信息。因此，可提示分割任务就是在给出任何提示的情况下，返回一个有效的分割掩码。对 "有效 "掩码的要求简单来说就是，即使提示模棱两可，可能指向多个对象（例如，回顾衬衫与人的对比示例，见图 3），输出也应该是至少其中一个对象的合理掩码。这一要求类似于期望语言模型对模棱两可的提示输出连贯的反应。我们之所以选择这项任务，是因为它带来了一种自然的预训练算法和一种通过提示将zero-shot转移到下游分割任务的通用方法。
Pre-training：可提示分割任务提出了一种自然的预训练算法，即为每个训练样本模拟一系列提示（如点、方框、掩码），并将模型的掩码预测与地面实况进行比较。我们从交互式分割法中借鉴了这一方法，但与交互式分割法不同的是，交互式分割法的目的是在足够多的用户输入后最终预测出有效的掩码，而我们的目的则是始终为任何提示预测出有效的掩码，即使该提示是模棱两可的。这就确保了预训练模型在涉及模棱两可的用例（包括我们的数据引擎 §4 所要求的自动注释）中的有效性。我们注意到，完成这项任务具有挑战性，需要选择专门的建模和训练损失，我们将在§3中讨论这一点。
Zero-shot transfer：直观地说，我们的预训练任务赋予了模型在推理时对任何提示做出适当反应的能力，因此下游任务可以通过设计适当的提示来解决。例如，如果我们有一个猫的边界框检测器，就可以通过将检测器的边界框输出作为对模型的提示来解决猫的实例分割问题。一般来说，大量的实际分割任务都可以作为提示任务。除了自动数据集标注，我们还在第 7 节的实验中探索了五种不同的示例任务。
Related tasks：分割是一个广泛的领域：有交互式分割、边缘检测、超级像素化、对象建议生成、前景分割、语义分割、实例分割、全景分割等。我们的可提示分割任务的目标是生成一个能力广泛的模型，可以通过提示工程适应许多（尽管不是全部）现有的和新的细分任务。这种能力是任务泛化的一种形式。请注意，这与之前的多任务分割系统工作不同。在多任务系统中，单个模型执行一组固定的任务，例如联合语义分割、实例分割和全视角分割，但训练和测试任务是相同的。我们工作中的一个重要区别是，为可提示分割训练的模型在推理时可以通过作为更大系统中的一个组件来执行新的、不同的任务，例如，为了执行实例分割，可提示分割模型与现有的对象检测器相结合。
Discussion：提示和组合是强大的工具，可以使单一模型以可扩展的方式使用，从而完成模型设计时未知的任务。这种方法类似于其他基础模型的使用方式，例如 CLIP 是 DALL-E图像生成系统的文本-图像对齐组件。我们预计，与专门针对固定任务集进行训练的系统相比，由提示工程等技术驱动的可组合系统设计将带来更广泛的应用。从组合的角度来比较可提示分割和交互式分割也很有趣：交互式分割模型是以人类用户为中心设计的，而为可提示分割而训练的模型也可以组合成一个更大的算法系统，正如我们将演示的那样。
3. Segment Anything Model
接下来，我们将介绍用于可提示分割的 “任意分割模型”（SAM）。SAM 有三个组成部分，如图 4 所示：图像编码器（Image encoder）、灵活的提示编码器(flexible prompt encoder)和快速掩码解码器(fask mask decoder)。我们以 Transformer 视觉模型为基础，对（摊销后的）实时性能进行了特定的权衡。在此，我们将对这些组件进行高层次的描述。
3.1.Image encoder
在可扩展性和强大的预训练方法的激励下，我们使用 MAE预训练的视觉转换器（ViT），并对其进行了最小化调整，以处理高分辨率输入。图像编码器每张图像运行一次，可在提示模型之前应用。
3.2.Prompt encoder
我们考虑了两组提示：稀疏（点、方框、文本）和密集（mask）。我们使用位置编码表示点和方框，并与每种提示类型的学习嵌入相加；使用 CLIP的现成文本编码器表示自由格式文本。密集提示（即masks）则使用卷积法嵌入，并与图像嵌入相加。
3.3.Mask decoder
掩码解码器能有效地将图像嵌入、提示嵌入和输出标记映射到掩码上。这一设计受的启发，采用了变形解码器块的修改，然后是动态掩码预测头。我们修改后的解码器块使用提示自注意和双向交叉注意（提示到图像嵌入，反之亦然）来更新所有嵌入。运行两个区块后，我们对图像嵌入进行上采样，然后用 MLP 将输出标记映射到动态线性分类器，再计算每个图像位置的掩码前景概率。
3.4.Resolving ambiguity
为了解决这个问题，我们对模型进行了修改，以预测单个提示的多个输出掩码（见图 3）。我们发现3个掩码输出足以应对大多数常见情况（嵌套掩码通常最多有三个深度：整体、部分和子部分）。在训练过程中，我们只对掩码进行最小损失反推。为了对掩码进行排序，模型会预测每个掩码的置信度得分（即估计的 IoU）。
3.5.Efficiency
整个模型的设计主要是为了提高效率。给定一个预先计算的图像嵌入，提示编码器和掩码解码器在网络浏览器的 CPU 上运行，只需 50 毫秒。这种运行性能使我们的模型能够实现无缝、实时的交互式提示。
3.6.Losses and training
我们使用中使用的focal loss和dice loss的线性组合来监督掩码预测。我们使用混合几何提示对可提示分割任务进行训练。我们通过在每个掩码中随机抽取 11 轮提示来模拟交互式设置，从而使 SAM 可以无缝集成到我们的数据引擎中。
4. Segment Anything Data Engine
由于互联网上的分割掩码并不丰富，我们建立了一个数据引擎，以便收集 1.1B 掩码数据集 SA-1B。数据引擎分为三个阶段：(1) 模型辅助人工标注阶段；(2) 半自动阶段，包括自动预测的掩码和模型辅助标注；(3) 全自动阶段，我们的模型无需标注者输入即可生成掩码。接下来我们将详细介绍每个阶段。
4.1.Assisted-manual stage
第一阶段类似于经典的交互式分割，由一组专业注释人员使用由 SAM 支持的基于浏览器的交互式分割工具，通过点击前景/背景对象点来标记掩码。可以使用像素精确的 "画笔 "和 "橡皮擦 "工具对掩码进行细化。我们的模型辅助标注直接在浏览器中实时运行（使用预先计算的图像嵌入），实现了真正的交互式体验。我们没有对标注对象施加语义限制，标注者可以自由标注 "东西 "和 “事物”。我们建议标注者标注他们可以命名或描述的对象，但不收集这些名称或描述。我们要求标注者按突出顺序标注对象，并鼓励标注者在一个mask的标注时间超过30 秒后继续标注下一幅图像。
在这一阶段开始时，使用常见的公共分割数据集对 SAM 进行训练。在获得足够的数据注释后，只使用新注释的掩码对SAM进行重新训练。随着收集的掩码越来越多，图像编码器也从ViT-B 升级到 ViT-H，其他架构细节也在不断变化；我们总共对模型进行了 6 次重新训练。随着模型的改进，每个掩码的平均标注时间从 34 秒减少到 14 秒。我们注意到，14 秒的注释时间比 COCO 的掩膜注释时间快 6.5 倍，只比用极值点进行边界框标注慢 2 倍。随着 SAM 的改进，每幅图像的平均掩码数量从 20 个增加到 44 个。总体而言，我们在这一阶段从 120k 幅图像中收集了 430 万个掩码。
4.2.Semi-automatic stage
在这一阶段，我们的目标是增加mask的多样性，以提高我们的模型分割任何物体的能力。为了让注释者专注于不太突出的物体，我们首先自动检测出有信心的mask。然后，我们向注释者展示预填充了这些掩码的图像，并要求他们注释任何其他未注释的物体。为了检测有把握的mask，我们使用通用的 "物体 "类别对所有第一阶段的mask进行了边界框检测器训练。在这一阶段，我们又在 180k 幅图像中收集了 590 万个mask（共计 1020 万个mask）。与第一阶段一样，我们定期在新收集的数据上重新训练模型（5 次）。每个掩码的平均标注时间回升到 34 秒（不包括自动掩码），因为这些对象的标注更具挑战性。每幅图像的平均mask数从 44 个增加到 72 个（包括自动mask）。
4.3.Fully automatic stage
在最后阶段，注释工作完全自动化。之所以能够做到这一点，是因为我们对模型进行了两大改进。首先，在这一阶段开始时，我们已经收集了足够多的掩码来大大改进模型，包括上一阶段的各种掩码。其次，在这一阶段，我们已经开发出了模棱两可的感知模型，即使在模棱两可的情况下，我们也能预测出有效的掩码。具体来说，我们用 32×32 的规则点网格来提示模型，并为每个点预测一组可能与有效对象相对应的遮罩。有了模糊感知模型，如果一个点位于一个部件或子部件上，我们的模型就会返回子部件、部件和整个对象。我们模型的 IoU 预测模块用于选择有把握的掩码；此外，我们只识别和选择稳定的掩码（如果在 0.5 - δ 和 0.5 + δ 时对概率图进行阈值处理会得到相似的掩码，我们就认为该掩码是稳定的）。最后，在筛选出可信和稳定的掩码后，我们使用非最大抑制（NMS）来过滤重复的掩码。为了进一步提高较小掩码的质量，我们还处理了多个重叠的放大图像裁剪。有关这一阶段的更多详情。我们对数据集中的所有 1,100 万张图像进行了全自动掩码生成，总共生成了 1.1B 个高质量掩码。 接下来，我们将对生成的数据集 SA-1B 进行描述和分析。
5. Segment Anything Dataset
我们的数据集 SA-1B 由 1,100 万张不同的高分辨率授权隐私保护图像和我们的数据引擎收集的 11B 个高质量分割掩码组成。我们将 SA-1B 与现有数据集进行了比较，并分析了掩码的质量和属性。我们发布 SA-1B 是为了帮助未来计算机视觉基础模型的开发。我们注意到，SA-1B 将根据有利的许可协议发布，用于某些研究用途，并为研究人员提供保护。
6. Segment Anything RAI Analysis
7. Zero-Shot Transfer Experiments
Appendix （Model细节）
A. Segment Anything Model and Task Details
A.1.Image encoder
原文：In general, the image encoder can be any network that outputs a C×H×W image embedding. Motivated by scalability and access to strong pre-training, we use an MAE [47] pre-trained Vision Transformer (ViT) [33] with minimal adaptations to process high resolution inputs, specifically a ViT-H/16 with 14×14 windowed attention and four equally-spaced global attention blocks, following [62]. The image encoder’s output is a 16× downscaled embedding of the input image. Since our runtime goal is to process each prompt in real-time, we can afford a high number of image encoder FLOPs because they are computed only once per image, not per prompt.
Following standard practices (e.g., [40]), we use an input resolution of 1024×1024 obtained by rescaling the image and padding the shorter side. The image embedding is therefore 64×64. To reduce the channel dimension, following [62], we use a 1×1 convolution to get to 256 channels, followed by a 3×3 convolution also with 256 channels. Each convolution is followed by a layer normalization [4].
翻译：一般来说，图像编码器可以是任何能输出 C×H×W 图像嵌入的网络。出于可扩展性和强大预训练能力的考虑，我们使用了经过 MAE预训练的视觉转换器（ViT），并对其进行了最小化调整，以处理高分辨率输入，具体来说，我们使用的是ViT-H/16，它具有 14×14 窗口注意力和四个等间距的全局注意力区块。图像编码器的输出是输入图像的16倍缩放嵌入。由于我们的运行目标是实时处理每个提示，因此我们可以承受较高的图像编码器 FLOPs 数量，因为它们只在每个图像上计算一次，而不是每个提示。
按照标准做法，我们使用 1024×1024 的输入分辨率，这是通过重新缩放图像并填充较短边而获得的。因此，图像嵌入为 64×64。为了减少通道维度，我们按照文献的方法，先进行 1×1 卷积，得到 256 个通道，然后再进行 3×3 卷积，也得到 256 个通道。 每次卷积后都要进行层归一化处理。
A.2.Prompt encoder
原文：Sparse prompts are mapped to 256dimensional vectorial embeddings as follows. A point is represented as the sum of a positional encoding [95] of the point’s location and one of two learned embeddings that indicate if the point is either in the foreground or background. A box is represented by an embedding pair: (1) the positional encoding of its top-left corner summed with a learned embedding representing “top-left corner” and (2) the same structure but using a learned embedding indicating “bottom-right corner”. Finally, to represent free-form text we use the text encoder from CLIP [82] (any text encoder is possible in general). We focus on geometric prompts for the remainder of this section and discuss text prompts in depth in §D.5.
Dense prompts (i.e., masks) have a spatial correspondence with the image. We input masks at a 4× lower resolution than the input image, then downscale an additional 4× using two 2×2, stride-2 convolutions with output channels 4 and 16, respectively. A final 1×1 convolution maps the channel dimension to 256. Each layer is separated by GELU activations [50] and layer normalization. The mask and image embedding are then added element-wise. If there is no mask prompt, a learned embedding representing “no mask” is added to each image embedding location.
翻译：稀疏提示映射到 256 维向量嵌入的方法如下：一个点表示为该点的位置编码和两个学习到的嵌入式编码的总和，嵌入式编码表示该点是在前景还是背景中。 一个方框用一对嵌入表示：(1) 其左上角的位置编码与表示 "左上角 "的学习嵌入总和；(2) 同样的结构，但使用表示 "右下角 "的学习嵌入。最后，我们使用 CLIP的文本编码器来表示自由格式文本（一般来说，任何文本编码器都可以）。在本节的其余部分，我们将重点讨论几何提示。
密集提示（即掩码）与图像具有空间对应关系。我们以比输入图像低 4 倍的分辨率输入掩码，然后使用两个 2×2 的stride-2 卷积将输出通道分别为 4 和 16 的掩码再降低 4 倍。最后的 1×1 卷积将通道维度映射为 256。每一层都通过 GELU 激活和层归一化进行分隔。然后逐元素添加掩码和图像嵌入。如果没有掩码提示，则在每个图像嵌入位置添加代表 "无掩码 "的学习嵌入。
A.3.Lightweight mask decoder
原文：This module efficiently maps the image embedding and a set of prompt embeddings to an output mask. To combine these inputs, we take inspiration from Transformer segmentation models [14, 20] and modify a standard Transformer decoder [103]. Before applying our decoder, we first insert into the set of prompt embeddings a learned output token embedding that will be used at the decoder’s output, analogous to the [class] token in [33]. For simplicity, we refer to these embeddings (not including the image embedding) collectively as “tokens”.
Our decoder design is shown in Fig. 14. Each decoder layer performs 4 steps: (1) self-attention on the tokens, (2) cross-attention from tokens (as queries) to the image embedding, (3) a point-wise MLP updates each token, and (4) cross-attention from the image embedding (as queries) to tokens. This last step updates the image embedding with prompt information. During cross-attention, the image embedding is treated as a set of 642 256-dimensional vectors. Each self/cross-attention and MLP has a residual connection [49], layer normalization, and a dropout [93] of 0.1 at training. The next decoder layer takes the updated tokens and the updated image embedding from the previous layer. We use a two-layer decoder.
To ensure the decoder has access to critical geometric information the positional encodings are added to the image embedding whenever they participate in an attention layer. Additionally, the entire original prompt tokens (including their positional encodings) are re-added to the updated tokens whenever they participate in an attention layer. This allows for a strong dependence on both the prompt token’s geometric location and type.
After running the decoder, we upsample the updated image embedding by 4× with two transposed convolutional layers (now it’s downscaled 4× relative to the input image). Then, the tokens attend once more to the image embedding and we pass the updated output token embedding to a small 3-layer MLP that outputs a vector matching the channel dimension of the upscaled image embedding. Finally, we predict a mask with a spatially point-wise product between the upscaled image embedding and the MLP’s output.
The transformer uses an embedding dimension of 256. The transformer MLP blocks have a large internal dimension of 2048, but the MLP is applied only to the prompt tokens for which there are relatively few (rarely greater than 20). However, in cross-attention layers where we have a 64×64 image embedding, we reduce the channel dimension of the queries, keys, and values by 2× to 128 for computational efficiency. All attention layers use 8 heads.
The transposed convolutions used to upscale the output image embedding are 2×2, stride 2 with output channel dimensions of 64 and 32 and have GELU activations. They are separated by layer normalization.
翻译：该模块能有效地将图像嵌入和一组提示嵌入映射到输出掩码。为了结合这些输入，我们从Transformer分割模型中获得灵感，并修改了标准Transformer decoder。在应用decoder之前，我们首先在提示嵌入集合中插入一个学习到的输出标记嵌入，该嵌入将用于解码器的输出，类似于中的 [class] 标记。 为简单起见，我们将这些嵌入（不包括图像嵌入）统称为 “标记”。
我们的解码器设计如图 14 所示。每个解码器层执行 4 个步骤：(1) 对标记进行自关注，(2) 从标记（作为查询）到图像嵌入进行交叉关注，(3) 由一个点向 MLP 更新每个标记，(4) 从图像嵌入（作为查询）到标记进行交叉关注。最后一步用提示信息更新图像嵌入。在交叉注意过程中，图像嵌入被视为一组 642 个 256 维向量。 每个自注意/交叉注意和 MLP 都有残差连接、层归一化和训练时 0.1 的 dropout。下一层解码器接收上一层更新的标记和更新的图像嵌入。 我们使用双层解码器。
为了确保解码器能够获得关键的几何信息，只要它们参与了注意力层，位置编码就会被添加到图像嵌入中。 此外，只要它们参与了注意力层，整个原始提示标记（包括其位置编码）就会被重新添加到更新的标记中。这使得提示标记的几何位置和类型都具有很强的依赖性。
运行解码器后，我们用两个转置卷积层对更新后的图像嵌入进行 4 倍增样（现在相对于输入图像缩小了 4 倍）。 然后，令牌再次参加图像嵌入，我们将更新后的输出令牌嵌入传递给一个小型 3 层 MLP，该 MLP 输出一个与放大图像嵌入的通道维度相匹配的向量。最后，我们用放大图像嵌入和 MLP 输出之间的空间点积来预测掩码。
转换器使用的嵌入维度为 256。 转换器 MLP 块的内部维度较大，为 2048，但 MLP 仅应用于提示令牌，而提示令牌的数量相对较少（很少超过 20 个）。不过，在交叉注意层中，我们有一个 64×64 的图像嵌入，为了提高计算效率，我们将查询、键和值的通道维度减少了 2 倍，为 128。所有注意力层都使用 8 个头，用于放大输出图像嵌入的转置卷积为 2×2，步长 2，输出通道维度分别为 64 和 32，并具有 GELU 激活。它们通过层归一化分开。
A.4.Making the model ambiguity-aware
原文：As described, a single input prompt may be ambiguous in the sense that it corresponds to multiple valid masks, and the model will learn to average over these masks. We eliminate this problem with a simple modification: instead of predicting a single mask, we use a small number of output tokens and predict multiple masks simultaneously. By default we predict three masks, since we observe that three layers (whole, part, and subpart) are often enough to describe nested masks. During training, we compute the loss (described shortly) between the ground truth and each of the predicted masks, but only backpropagate from the lowest loss. This is a common technique used for models with multiple outputs [15, 45, 64]. For use in applications, we’d like to rank predicted masks, so we add a small head (operating on an additional output token) that estimates the IoU between each predicted mask and the object it covers.
Ambiguity is much rarer with multiple prompts and the three output masks will usually become similar. To minimize computation of degenerate losses at training and ensure the single unambiguous mask receives a regular gradient signal, we only predict a single mask when more than one prompt is given. This is accomplished by adding a fourth output token for an additional mask prediction. This fourth mask is never returned for a single prompt and is the only mask returned for multiple prompts.
翻译：如前所述，单个输入提示可能是模棱两可的，因为它对应多个有效掩码，而模型将学习对这些掩码进行平均处理。我们通过一个简单的修改就可以解决这个问题：我们不预测单个掩码，而是使用少量的输出标记同时预测多个掩码。默认情况下，我们预测三个掩码，因为我们观察到三个层（整体、部分和子部分）往往足以描述嵌套掩码。在训练过程中，我们会计算地面实况和每个预测掩码之间的损失（稍后会介绍），但只会从损失最小的掩码开始反向传播。这是有多个输出的模型常用的技术。 在应用中，我们希望对预测的掩码进行排序，因此我们添加了一个小头（在额外的输出标记上运行）来估计每个预测的掩码与其覆盖的对象之间的 IoU。
如果有多个提示，模棱两可的情况就会少得多，而且三个输出掩码通常会变得相似。为了在训练时尽量减少退化损失的计算，并确保单一的无歧义掩码能接收到有规律的梯度信号，我们在给出多个提示时只预测一个掩码。为此，我们增加了第四个输出标记，用于额外的掩码预测。对于单个提示，第四个掩码不会返回，而对于多个提示，第四个掩码是唯一返回的掩码。
对Model结构的详解
image decoder:
image encoder的作用是把图像映射到特征空间，本质上这个encoder可以是任何网络结构，作者这里使用的是微调的detectron的ViT。输入图像经过ViT的流程，只有4个步骤：
- 原始图像被等比和padding的缩放到1024大小，输入图像进入网络，先经过一个卷积base的patch_embedding：取16*16为一个patch，步长也是16，这样feature map的尺寸就缩小了16倍，同时channel从3映射到768，即64x64X768（W，H，C)的向量。
- patch_embed过后加positional_embedding：positional_embedding是个可学习的参数矩阵，初始化是0。
- 加了positional_embedding后的feature map过16个transformer block，其中12个transformer是基于window partition（就是把特征图分成14*14的windows做局部的attention）的attn模块，和4个全局attn，这4个全局attn是穿插在windowed attention中的。
- 最后过两层卷积（neck）（kernel分别为1和3，每层输出接入layer norm2d）把channel数降到256，这就是最终的image embedding的结果。
prompt encoder：
prompt encoder负责prompt到prompt的特征空间，映射出的特征的channel和image embedding的channel一致，因为这两个后边要用attention进行融合。SAM 支持的prompt可以分为以下两类：
稀疏类（sparse prompt）：包含 point，bbox，free text
Point: prompt是point，那么它的映射由两个部分相加组成，一个是位置编码，这里的位置编码使用的是Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains的编码方式，另一个部分是一个描述当前点是前景还是背景特征的可学习的一维向量。
Box：box的映射也是由两个部分相加组成，第一部分是左上和右下两个点的位置编码，第二部分是一组一维向量用来描述这个点是“左上”还是“右下”。对于左上的点，他的映射就是位置编码+“左上”这个特征的描述向量，右下的点，就是位置编码+“右下”这个特征的描述向量。
Text：分别采用clip预训练好的text encoder 作为文本编码器，文中没有说明image encoder 是采用clip 的image encoder，还是采用前边MAE 预训练的vit，而point和bbox的特征维度为256，所以还存在全连接进行特征维度对齐），将文本特征向量和图像特征向量维度从768变为256作为下一步使用。
稠密类（dense prompt）
Mask：在输入prompt encoder之前，先要把mask降采样到4x，再过两个2x2，stride=2的卷积，这样尺寸又降了4x，就和降了16x的图像特征图尺寸一致了，再过一个1*1的卷积，把channel也升到256。果不提供粗略分割的输入，会用默认可学习的embedding代表空分割prompt的特征。
mask decoder:
mask decoder 架构图如下，但是这张图没有说明太多细节不容易看懂
decoder的结构之所以看起来复杂，主要原因是prompt embedding和image embedding在这个结构中反复融合并且反复更新。
1.在prompt embedding进入decoder之前，先在它上面concat了一组可学习的output tokens，output tokens由两个部分构成：
- 一个是iou token，它会在后面被分离出来用于预测iou的可靠性（对应结构图右侧的IoU output token），它受到模型计算出的iou与模型计算出的mask与GT实际的iou之间的MSE loss监督；
- 另一个是mask token，它也会在后面被分离出来参与预测最终的mask（对应结构图右侧的output token per mask），mask受到focal loss和dice loss 20:1的加权组合监督。
- 这两个token的意义我感觉比较抽象，因为理论来说进入decoder的变量应该是由模型的输入，也就是prompt和image的映射构成，但这两个token的定义与prompt和image完全没有关系，而是凭空出现的。从结果反推原因，只能把它们理解成对模型的额外约束，因为它们两个参与构成了模型的两个输出并且有loss对他们进行监督。
- 最终prompt embedding（这一步改名叫prompt token）和刚才提到这两个token concat到一起统称为tokens进入decoder。
2.image embedding在进入decoder之前也要进行一步操作：dense prompt由于包含密集的空间信息，与image embedding所在的特征空间一致性更高，所以直接与image embedding相加融合。因为后面要与prompt做cross attention融合，这里还要先算一下image embedding的位置编码。
3.接下来{image embedding，image embedding的位置编码，tokens}进入一个两层transformer结构的decoder做融合。值得注意的是，在transformer结构中，为了保持位置信息始终不丢失，每做一次attention运算，不管是self-attention还是cross-attention，tokens都叠加一次初始的tokens，image embedding都叠加一次它自己的位置编码，并且每个attention后边都接一个layer_norm。
- tokens先过一个self-attention。
- tokens作为q，对image embedding做cross attention，更新tokens。
- tokens再过两层的mlp做特征变换。
- image embedding作为q，对tokens做cross attention，更新image embedding。
4.更新后的tokens作为q，再对更新后的image embedding做cross attention，产生最终的tokens。
5.更新后的image embedding过两层kernel_size=2, stride=2的转置卷积，升采样到4x大小（依然是4x降采样原图的大小），产生最终的image embedding。
6.接下来兵分两路：
- mask token被从tokens中分离出来（因为他一开始就是concat上去的，可以直接按维度摘出来），过一个三层的mlp调整channel数与最终的image embedding一致，并且他们两个做矩阵乘法生成mask的预测。
- iou token被从tokens中分离出来，也过一个三层的mlp生成最终的iou预测。
7.最后，如前文所述，分别对mask的预测和iou预测进行监督，反向传播，更新参数。
272
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
