SAM模型细节分析 (附录内容)

Image Encoder

ViTDet:arxiv.org/pdf/2203.16527

一般来说，图像编码器可以是任何输出大小为 C×H×W 的图像嵌入（image embedding）的网络。出于可扩展性和强预训练模型的考虑，我们使用了一个在 MAE 上预训练的视觉变换器（Vision Transformer, ViT），并进行了最小的调整以处理高分辨率输入。

具体来说，我们采用了 ViT-H/16，使用 14×14 的窗口注意力机制，并在四个等距位置加入了全局注意力模块，结构参考 VITDet。图像编码器的输出是输入图像下采样 16 倍后的嵌入表示。

由于我们的运行时目标是实时处理每个提示（prompt），因此我们可以在图像编码器上使用大量计算资源，因为它每个图像只需计算一次，而不是每个提示都要重新计算。

按照标准做法（例如 [Copy-Paste]），我们将图像的输入分辨率统一调整为 1024×1024，通过缩放图像并填充短边。因此图像嵌入的大小为 64×64。

为了降低通道维度，参考VITDet的做法，我们使用一个 1×1 的卷积将通道数降到 256，接着使用一个 3×3 的卷积，输出也为 256 通道。每个卷积层之后都跟一个层归一化(Layer Normalization)。

Prompt Encoder

稀疏提示被映射为 256 维的向量化嵌入，具体方式如下：

• 一个点由该点位置的位置编码 [Fourier features let networks learn high frequency functions in low dimensional domains] 与两个可学习嵌入中的一个相加得到，这两个可学习嵌入分别表示该点是前景还是背景。
• 一个框由一对嵌入表示：(1) 左上角的位置编码与一个表示“左上角”的可学习嵌入相加；(2) 同样的结构，但使用表示“右下角”的可学习嵌入。
• 最后，为了表示自由形式的文本，我们使用 CLIP 的文本编码器（原则上也可以使用任何文本编码器）。

密集提示（例如掩码）在空间上与图像相对应。我们将掩码的输入分辨率设为原始图像的 1/4，然后通过两个 2×2、步长为 2 的卷积层进一步下采样 4 倍，输出通道数分别为 4 和 16。最后通过一个 1×1 卷积将通道维度映射到 256。每一层之间使用 GELU 激活函数和层归一化（Layer Normalization）。然后将掩码嵌入与图像嵌入逐元素相加。如果没有掩码提示，则在每个图像嵌入位置添加一个表示“无掩码”的可学习嵌入。

Lightweight Mask Decoder

该模块高效地将图像嵌入和一组提示嵌入映射为输出掩码。为了融合这些输入，我们借鉴了 Transformer 分割模型 [DETR, MaskFormer] 的设计，并对标准 Transformer 解码器进行了修改。

在应用解码器之前，我们在提示嵌入集合中插入一个可学习的输出标记嵌入（output token embedding），该标记将在解码器输出时使用，类似于 ViT 中的 [class] 标记。为简便起见，我们将这些嵌入（不包括图像嵌入）统称为“标记（tokens）”。我们的解码器设计如图 14 所示。

每个解码器层执行以下四个步骤：

1. 对标记进行自注意力（self-attention）；
2. 使用标记作为查询，对图像嵌入进行交叉注意力（cross-attention）；
3. 逐点 MLP（多层感知机）更新每个标记；
4. 使用图像嵌入作为查询，对标记进行交叉注意力，这一步更新图像嵌入以融合提示信息。

在交叉注意力过程中，图像嵌入被视为一组 64² 个 256 维向量。每个自注意力/交叉注意力模块和 MLP 模块都包含残差连接 、层归一化和训练时 0.1 的 dropout 。下一个解码器层接收上一层更新后的标记和图像嵌入。我们使用了一个两层的解码器。

为了确保解码器能够获取关键的几何信息，每当图像嵌入参与注意力层时，都会加入位置编码。此外，每当提示标记参与注意力层时，都会重新加入原始的完整提示标记（包括它们的位置编码）。这使得输出对提示标记的几何位置和类型具有强依赖性。

运行解码器之后，我们通过两个转置卷积层将更新后的图像嵌入上采样 4 倍（现在其分辨率是输入图像的 1/4）。然后，让标记再次对图像嵌入进行注意力操作，并将更新后的输出标记嵌入传入一个小型的三层 MLP，该 MLP 输出一个向量，其通道维度与上采样后的图像嵌入匹配。最后，我们通过图像嵌入与 MLP 输出之间的逐点空间乘积来预测掩码。

Transformer 使用的嵌入维度为 256。Transformer 中的 MLP 块具有较大的内部维度 2048，但 MLP 仅作用于提示标记（数量通常较少，很少超过 20）。然而，在交叉注意力层中，由于我们处理的是 64×64 的图像嵌入，为了计算效率，我们将查询（queries）、键（keys）和值（values）的通道维度减少一半至 128。所有注意力层均使用 8 个注意力头。

用于上采样输出图像嵌入的转置卷积层为 2×2、步长为 2，输出通道分别为 64 和 32，并使用 GELU 激活函数，层之间通过层归一化分隔开。

Making the model ambiguity-aware使模型具备识别歧义的能力

如前所述，一个单一的输入提示可能会存在歧义，即它可以对应多个有效的掩码，而模型会倾向于在这些掩码之间进行平均。我们通过一个简单的修改来解决这个问题：不是只预测一个掩码，而是使用少量的输出标记（output tokens），同时预测多个掩码。默认情况下，我们会预测三个掩码，因为我们发现三层（整体、部分、子部分）通常足以描述嵌套的掩码结构。

在训练过程中，我们在真实掩码与每个预测掩码之间分别计算损失，但仅从损失最小的那个掩码进行反向传播。这是一种在具有多个输出的模型中常用的技巧。

在实际应用中，我们希望对预测的掩码进行排序，因此我们增加了一个小型的头部网络（head）（作用在额外的一个输出标记上），用于估计每个预测掩码与其所覆盖对象之间的 IoU（交并比）。在使用多个提示时，歧义情况会少得多，此时三个输出掩码通常会变得相似。

为了最小化训练过程中无效损失的计算，并确保唯一一个无歧义的掩码能获得正常的梯度信号，当给出多于一个提示时，我们只预测一个掩码。这是通过添加第四个输出标记来实现的，用于额外的掩码预测。

第四个掩码在单个提示时从不返回，而在多个提示时是唯一返回的掩码。

Loss

我们使用 focal loss（焦点损失） 和 dice loss（骰子损失） 的线性组合作为监督信号，其中 focal loss 与 dice loss 的比例为 20:1，这一设置参考了 [MaskFormer和DETR] 的做法。

与 [MaskFormer和DETR] 不同的是，我们发现在每个解码器层之后加入辅助的深度监督（deep supervision）并没有帮助。

用于预测 IoU 的头部网络是通过预测的 IoU 与预测掩码与真实掩码之间的 IoU 的均方误差（MSE）进行训练的。该损失项通过一个固定的比例因子 1.0 被加入到掩码损失中。

算法训练

遵循近期的研究方法 [RITM]，我们在训练过程中模拟了一个交互式分割（interactive segmentation）的设置。

首先，以相等的概率，从目标掩码中随机选择一个前景点或边界框作为提示。

• 点是从真实掩码中均匀采样得到的。
• 边界框是取自真实掩码的边界框，并在每个坐标上添加标准差为框边长10%的随机噪声（最大不超过20像素），以模拟不同应用场景下的框输入。这种噪声设定在实例分割（通常生成紧贴目标的框）和交互式分割（用户可能画出较松散的框）之间取得了一个合理的折中。

在使用第一个提示进行预测之后，后续的点将从上一次预测的掩码与真实掩码之间的误差区域中均匀采样。

• 如果误差区域是假阴性区域（false negative），则采样一个前景点。
• 如果是假阳性区域（false positive），则采样一个背景点。

此外，我们还将上一次迭代的掩码预测结果作为额外提示输入给模型。为了在下一轮迭代中提供尽可能多的信息，我们输入的是未经过阈值处理的掩码logits，而不是二值化的掩码。

当模型返回多个掩码时，我们选择预测IoU最高的那个掩码用于下一轮迭代和采样下一个点。

我们发现，在采样8个点之后收益递减（我们测试过最多16个点）。为了鼓励模型利用提供的掩码信息，我们还额外进行了两次不采样新点的迭代。其中一次被随机插入到8个采样点之间，另一次则固定在最后。

因此，总共有11轮迭代：

1. 第一轮：初始提示采样
2. 接下来的8轮：依次采样新点
3. 最后2轮：无新点输入，仅优化掩码预测

我们注意到，之所以可以使用较多的迭代次数，是因为我们的轻量级掩码解码器仅消耗图像编码器不到1%的计算量，每轮迭代带来的额外开销非常小。这与以往的交互式方法不同，那些方法通常每个优化步骤只执行一到几次交互。

训练策略（Training recipe）

我们使用 AdamW 优化器（β1 = 0.9，β2 = 0.999），并采用 250次迭代的线性学习率预热（linear learning rate warmup）以及分阶段学习率衰减策略。

• 预热后初始学习率（lr）设为 8e−4。
• 总共训练 90,000次迭代（大约2个SA-1B数据集的完整训练周期）。
• 在第60,000次和第86,666次迭代时，学习率分别降低为原来的1/10。
• 每批训练图像数量为256张。

为了防止过拟合，我们对SAM进行了以下正则化设置：

• 权重衰减（weight decay, wd）设为 0.1
• 使用路径丢弃（drop path, dp）技术，丢弃率设为 0.4
• 应用了分层学习率衰减（layer-wise learning rate decay, ld），衰减因子为 0.8

不使用数据增强（no data augmentation）。
我们从MAE  预训练的ViT-H模型初始化SAM。由于图像编码器较大且输入尺寸为1024×1024，我们在256块GPU上进行分布式训练。

• 为了限制GPU内存使用，我们每块GPU上最多训练64个随机采样的掩码。
• 此外，我们发现对SA-1B数据集中的掩码进行轻微过滤（例如去除覆盖超过图像90%的掩码）可以在定性上提升模型表现。

消融实验与训练变体（Ablations and training variations）

在进行消融实验或其他训练变体（如文本到掩码任务 ）时，我们会对上述默认训练策略进行如下调整：

• 当仅使用第一阶段和第二阶段的数据引擎生成的数据进行训练时，我们对输入图像进行大规模抖动增强（large-scale jitter [40]），缩放范围设为 [0.1, 2.0]。直觉上，在训练数据有限的情况下，数据增强可能有助于提升性能。

对于训练 ViT-B 和 ViT-L 模型，我们调整训练策略如下：

• 训练迭代次数为180,000次
• 批量大小为128张图像，分布在128块GPU上
• 学习率分别设为 8e−4（ViT-B） / 4e−4（ViT-L）
• 分层学习率衰减设为 0.6（ViT-B） / 0.8（ViT-L）
• 权重衰减统一设为 0.1
• 路径丢弃率设为 0.6（ViT-B） / 0.4（ViT-L）