点击下方卡片，关注 「AI视界引擎」 公众号
在各种下游任务中，包括图像编辑、图像对应和3D形状生成等，已经取得了显著的进展，其中使用了大型的视觉语言模型，例如Stable Diffusion（SD）。受到这些进展的启发，作者通过提出SLiMe来探索如何利用这些强大的视觉语言模型来实现在所需粒度上分割图像，仅使用一个标记样本即可。SLiMe将这个问题框架化为一个优化任务。
具体而言，给定一张训练图像及其分割Mask，作者首先从SD先验中提取注意力图，包括作者的新颖的“加权累积自注意力图”。然后，使用提取的注意力图，作者对Stable Diffusion的文本嵌入进行优化，使每个文本嵌入都能学习来自训练图像的单个分割区域。然后，这些学习到的嵌入会在注意力图中突出显示分割区域，进而可用于生成分割图。这使得SLiMe能够在推断过程中以训练图像中分割区域的粒度对任何现实世界的图像进行分割，仅使用一个示例。
此外，在可用时利用额外的训练数据，即fewshot，可以提高SLiMe的性能。作者进行了一系列知识丰富的实验，研究了各种设计因素，并表明SLiMe在超过其他现有的One-Shot和Few-Shot分割方法上表现出色。
1、简介
分割是一个多层次的问题，解决方案存在于不同的粒度级别。因此，从最终用户的角度来看，引导分割方法的一个简单而有效的方法是确定需要分割的内容和所需的细节级别，只需提供一两个示例供方法在训练期间使用作为参考。用户不应该需要整理带有地面真实分割的大量图像数据集，训练大型分割模型，或者编码目标对象的复杂属性。这种以示例为驱动的指导可以使用户能够根据其特定需求直观地定义和优化目标分割。
例如，在表情识别或面部对齐的应用中，将面部图像分割为鼻子、眼睛和头发等基本区域可能足够了。但是，在视觉效果应用中，诸如眼袋、额头、下巴和嘴唇等更详细的分割是必要的，例如皱纹去除等任务。因此，对于不同粒度级别的灵活可定制分割技术变得至关重要。然而，对于所有可能的预期粒度级别拥有大量的地面真实标注是非常不切实际的，因此需要一种一次或少样本学习来解决这一挑战。
最近的研究已经探索了少样本学习，提出了一些有前途的方法，如ReGAN，它使用GAN生成数据进行手动注释和模型训练。相比之下，SegDDPM从预训练扩散模型（DM）中提取特征，并使用少量标签对分割进行多MLP训练。这两者在使用10-50个示例进行分割时表现出色，但在样本非常有限的情况下表现不佳。而SegGPT采用一次学习，训练包括实例和部分级别细节的颜色随机化分割数据。在推断过程中，它仅使用参考图像及其二进制分割图来分割目标图像中的一个区域。尽管SegGPT非常有效，但它需要大量的分割数据进行预训练，仍然无法解决仅依赖单个注释进行有效训练的挑战。
在这项工作中，作者提出了Segment-Like-Me（SLiMe），它可以以一次学习的方式从同一类别的给定模板中分割任何对象/部分，避免了需要大量标记的分割数据（请参见图1和图2中的一些示例）。
为此，作者利用了现有的大规模预训练视觉/语言模型的丰富知识，例如Stable Diffusion（SD）。最近的研究，例如Hertz等强调了像SD这样的模型的交叉注意力图在相应的文本发生变化时如何突出显示图像的不同区域。这一显著特性最近已被用于修改生成的图像和实现图像对应。在这一思想的基础上，作者提出了两个关键见解。首先，作者观察到独立的交叉注意力图可能存在噪音，导致分割不精确，如图3所示。为了纠正这一问题，作者提出了一种新颖的加权累积自注意力图。
这个注意力图包含关键的语义边界信息，并使用高分辨率的自注意力图，确保提高了分割精度。其次，SD中的文本嵌入可以进行微调，以捕捉语义细节，如分割区域，其中每个嵌入对应于一个不同的分割区域，由参考图像及其分割图进行引导。
基于这些见解，SLiMe利用单个训练图像及其关联的分割图来通过注意力图微调SD的文本嵌入。这些精炼的嵌入强调了注意力图中的分割区域，可以转化为分割图。一旦学会了这些嵌入，它们就可以用于在推断过程中分割多个现实世界的图像，反映了训练图像中分割区域的粒度。
通过各种定量和定性实验，作者突出了作者方法的有效性。即使仅依赖一个或少数示例，作者的方法在性能上表现出色，可以与需要大量训练的监督方法媲美。此外，与少样本方法相比，作者的结果显示出显著的性能优势。在10个样本的情况下，作者的性能超过ReGAN约8%，在1个样本的情况下，作者超过SegGPT约7%。在20个样本的情况下，SLiMe在mIoU方面也优于SegDDPM。这些成就确立了作者的方法作为最先进技术领域的一流竞争者。
2、背景
2.1 潜在扩散模型
潜在扩散模型（LDMs）属于生成模型的一类，通过将输入随机变量高效地压缩成自动编码器的潜在空间并利用DM来表征这个潜在空间的分布，来建模输入随机变量的数据分布。自动编码器在大量图像集合上进行训练，包括两个主要组件：编码器（E）和解码器（D）。
编码器为给定图像I生成图像嵌入，使得。与此同时，解码器的目标是将I映射回图像空间，确保。自动编码器可以被公式化为离散生成模型或连续生成模型，并使用KL散度损失进行优化。另一方面，DM用于对由自动编码器生成的图像嵌入的分布进行建模。
DM的一个吸引人的特点是它可以轻松扩展以表示条件分布，即在文本或类别的条件下。在本研究中，作者关注的是文本条件的LDMs，其中其DM是条件于文本的。在文本条件的LDM中，作者从自然语言提示开始，记为B，并对其进行分词以获取P，其中包含一个由T个Token组成的序列，每个Token索引为，的范围在0到T - 1之间。分词的提示通过文本编码器θ进行变换，得到相应的文本嵌入θ。这个文本嵌入随后作为DM的一个输入，表示为ϵθ。
对于训练文本条件的LDM，作者将时间相关的噪声添加到输入图像的嵌入中。然后，作者使用以下目标来优化θ和ϵθ的参数，旨在使模型具有预测和消除添加噪声ϵ的能力：
其中是在步骤（即ϵ相对于添加到中）的噪声样本。
在这项工作中，作者使用SD，一个文本条件的LDM，有两个主要原因。首先，SD的交叉注意力模块展示了文本和图像嵌入之间丰富的语义连接。其次，SD的自动编码器将输入图像转换为具有语义含义的潜在向量。这个潜在向量包括4个通道，具有全局和局部信息，如图4所示，增强了文本和图像之间的相互关系。
2.2 注意力模块
SD中使用的DM采用UNet结构，包括ResNet块和两个注意力模块的变种：自注意力和交叉注意力。自注意力块计算图像嵌入之间的注意力，捕获特定元素与同一图像嵌入中的其他元素之间的关系。另一方面，交叉注意力块计算两种不同模态的潜在表示之间的关系，例如文本和图像。
注意力块包括三个组件：Query、Key 和Value。注意力模块是一个将输入Query通过key-value对变换为输出的函数。因此，给定维度为d的Query Q、Key K和Value V向量，输出O定义如下：
对于自注意力模块，所有Query、键和值向量都来自图像嵌入，而在交叉注意力模块中，Query向量来自图像的潜在表示，Key和Value向量来自文本嵌入。
在作者的情况下，作者提取规范化的注意力图S，表示为，适用于自注意力和交叉注意力模块，并分别表示为和。在这个背景下，和分别代表了自注意力和交叉注意力模块中使用的图像嵌入的高度和宽度。
在自注意力模块中，输入的每个像素p都与一个激活图关联，有效地识别与p共享相似语境的像素。此外，随着作者离开像素p，这个激活图的强度逐渐减小。本质上，这意味着如果作者选择图像上的一个像素，例如鸟嘴上的一个像素，并检查自注意力模块中对应的激活图，作者将观察到它突出显示鸟嘴上的其他像素，同时排除语义无关的像素。
此外，对于突出显示的像素，随着它们离开p越远，其亮度逐渐增加。图5显示了一只鸟的图像以及SD的自注意力的6个通道，其中每个通道对应于鸟嘴上随机选择的一个像素。
另一方面，对于每个输入的文本标记，交叉注意力模块生成一个注意力图，表示为。该图有效地突出显示了与模型语义空间中的Pi语义相符的图像嵌入像素。例如，如果作者指示作者的模型通过输入文本B = “一只熊”来生成一张熊的图像，并检查与“熊”标记相关联的中的注意力图，作者将看到在生成的图像中与熊对象对应的像素上有更强的强调。
3 方法
先前的研究已经证明了在生成过程中使用交叉注意力图来检测粗略分段的对象，但仍然存在关于这些方法是否适用于在真实世界图像中进行更精细分段的不确定性。作者引入了SLiMe，这是一种方法，使作者能够在真实世界和生成的图像中以各种粒度级别执行分段，只需要一个标记的图像作为参考。
在作者的方法中，在训练阶段，作者首先提取交叉注意力图和加权累积自注意力图，以微调文本嵌入，使每个文本嵌入都包含来自各个分段区域的语义信息（图6）。
在推断阶段，作者使用这些优化后的嵌入来突出显示分段部分，然后获得分段图（图7）。在接下来的小节中，作者将首先深入讨论文本优化的细节，然后再讨论推断过程。
3.1 优化文本嵌入
接下来，作者将解释如何优化文本嵌入以了解分段区域。给定一个成对的图像（）和 segmentation Mask（），其中 segmentation Mask涉及K类别，作者的初始步骤涉及向文本分词器提供一个空文本，表示为B = “”，以生成P。
重要的是要强调，P中的Token数量，表示为T，必须超过K类别。T的数量通常非常大。例如，SD包括77个Token，使得该方法可以容纳多达76个分段类别。P中的第一个Token表示文本的起始（SOT），由于B没有任何文本内容，所以分词器会在P中的其余T - 1个Token位置插入文本的结束（EOT）Token。（SOT和EOTToken具有特殊功能，它们用于指示文本的启动并填充Token化的输入文本以达到特定长度。）随后，这些Token随后通过SD的文本编码器进行处理，得到表示为P的文本嵌入。然后将输入图像I传递给E以获取图像嵌入I。
然后，作者对标准高斯噪声进行采样，并将其与时间步t相对应地添加到I中，并将其命名为。最后，作者将和馈送给UNet以获取去噪后的图像嵌入。在将这些数据传递给UNet之后，作者提取了前K个文本Token的对应的规范化交叉注意力图，，其中。
此外，SD在不同层中具有几个交叉注意力模块，作者使用表示第层的规范化交叉注意力图。作者对从不同层中提取的不同规范化交叉注意力图进行平均，因为根据经验，作者已经看到平均会增强结果。但是，由于和对于不同层中的不同注意力模块是不同的，因此作者将用于计算的层的调整为相同的大小。最后，作者在作者的损失函数中使用的注意力图计算如下：
其中，Averagel计算在层之间进行平均，Resize是双线性插值。图8直观地展示了这个过程。现在，对于，作者希望优化，使得突出显示分段Mask中第i个类别的区域。
请注意，作者不优化并将其分配给背景类别，因为如果更改其值，SD将无法正常运行。
首先，作者堆叠的，然后与M一起使用它们来计算交叉熵损失，如下所示：
其中CE表示交叉熵。尽管此损失优化，使得突出显示所需区域，但由于用于计算的的大小明显小于输入图像，其中对象边缘在其中是模糊的，因此结果的分段非常不准确。为了提高质量并提高准确性，在与交叉注意力一起，作者还利用了自注意力图（）。
使用自注意力图有三个基本原因。首先，它包含了有关图像语义的丰富信息。实际上，对于自注意力模块输入中的每个像素p，它都有一个激活图，突出显示与p具有相同语义意义的像素。其次，作者使用的自注意力模块的输出大小大于作者使用的交叉注意力模块的输出大小。因此，它提供了更精确的分割。第三，与交叉注意力图相比，它更详细地显示了边缘。
与交叉注意力层类似，出于相同的原因，作者使用了多个自注意力层，并将它们平均如下：
其中。方程式5中没有Resize函数的原因是每层中自注意力图的通道数量取决于该层中自注意力图的大小。因此，为了能够对它们进行平均，作者使用的自注意力图应具有相同数量的通道，因此它们的大小应相同。
现在，要使用自注意力图，作者首先通过线性插值将的大小调整为与相同的大小，并将其称为。通过这样做，对于中的每个像素p，作者都有一个在中突出显示与p相关像素的通道。然后，作者计算的通道的加权和，并将其命名为（WAS-attention图），其中每个通道的权重是中相应通道的像素的值：
这个过程如图8所示，可以改进边缘，因为对语义区域的边缘和边界有很好的理解。
现在作者有了，首先将其调整大小，并显示为。然后，以以下方式计算均方误差（MSE）损失：
其中是来自M的二进制Mask，其中属于第i类的像素为1，所有其他像素为0。
作者使用的最后一个损失函数是SD的损失函数（LLDM），它是添加噪声和预测噪声之间的均方误差损失。作者使用此损失，因为它防止了文本嵌入偏离SD可理解的空间。
最后，作者的优化文本嵌入的目标定义为：
其中α和β是损失函数的系数。
3.2 推断
在推断过程中，作者的目标是以与训练图像相同的细节级别对测试图像进行分割。为实现这一目标，作者从表示为的测试图像开始，将其编码为SD的潜在空间。随后，引入标准高斯噪声到编码图像中，噪声大小由时间参数t确定。重要的是要注意，这个t和高斯噪声值在训练时间的选择保持一致。有关选择特定t的详细全面分析在附录中提供。随后，作者使用优化的文本嵌入（P）以及编码图像来从UNet模型的前向传递中导出相应的交叉注意力和自注意力图。
如图8所示，这些注意力图使作者能够获得每个文本嵌入索引的WAS-attention图。接下来的步骤涉及选择与K个类别相对应的前K个WAS-attention图。然后，使用双线性插值将这些所选图的大小调整为与输入图像的尺寸匹配，并将它们沿通道维度堆叠。随后，作者继续通过在通道上执行argmax计算来生成分割Mask。值得注意的是，此过程可以在推断过程中针对多个测试图像重复进行，无需重新训练。
4、实验
4.1 PASCAL-PART
这个数据集专注于对象解析和语义部分分割。它提供了图像中对象部分的详细注释，涵盖了20个不同的对象类别。在作者的实验中，作者使用了包括车辆类别在内的车类别，其中包括车身、车灯、车牌、车轮和车窗部分。
作者还使用了包括马类别在内的马类别，其中包括头部、颈部+躯干、腿和尾巴的Mask。对于车辆类别，作者在ReGAN的作者提供的测试集上进行测试。但是对于马类别，由于他们没有提供他们的测试集，作者在PASCAL-Part数据集的测试集中测试了所有图像。
表1显示了作者在马类别上的结果。由于ReGAN的作者只报告了10样本设置的结果，所以作者只包括了这些结果。正如表1所示，作者的方法SLiMe在平均值和大多数部分上都优于ReGAN和SegGPT。在10样本设置中，作者的方法在除颈部+躯干以外的所有部分都优于ReGAN，而SLiMe只落后约2%。在1样本设置中，SLiMe在大多数部分上表现优于SegGPT，除头部和腿部外。
表2显示了车辆类别上类似的结果。SLiMe在10样本设置中在平均值和所有部分分割上都优于ReGAN。在1样本设置中，SLiMe在平均值和车身、车牌和车灯分割上优于SegGPT，但在其余3个部分上表现不佳。这背后的主要原因是，尽管对于推断，SegGPT只需要一个分割示例，但它是一种完全监督的方法，需要大量的训练集。相反，SLiMe是一个真正的一次性技术，作者只需要一个示例进行训练。
4.2 CELEBAMASK-HQ
CelebAMask-HQ是一个用于计算机视觉中面部分割任务的数据集。该数据集包含超过30,000张高分辨率的名人图像，带有19种不同面部属性的注释。
在本文中，作者使用了该数据集的两个版本：CelebAMask-HQ10和CelebAMaskHQ19，其中数字表示语义类别的数量。
SegDDPM使用已训练的DDPM中的UNet特征进行少样本分割。它首先在数据集上训练DDPM或使用预训练模型，然后利用少数样本分割进行特征提取。在SegDDPM中，PASCAL-Part没有可用的预训练模型。
因此，作者将作者的方法与SegDDPM在CelebAMaskHQ19上进行了比较，这是他们在论文中使用的数据集。为了进行公平比较，与SegDDPM一样，作者使用20张图像进行训练。表3显示了SLiMe在平均值上优于SegDDPM的结果。
为了将作者的结果与ReGAN进行比较，作者使用CelebAMask-HQ10。表4显示了作者的定量结果，图9显示了一些定性结果。如表所示，作者的方法在10样本和1样本设置中的平均值上均优于ReGAN。在1样本设置中，SLiMe在所有部分分割中要么优于其他方法，要么在所有部分分割中表现相当。
请注意，SegGPT受益于在大型数据集上的训练。在10样本设置中，虽然SLiMe在平均值上优于ReGAN，但它们在多个类别的性能上存在差异。例如，SLiMe在服装和耳朵部分的性能要好得多，而在鼻子部分的性能不如ReGAN。作者认为这种差异的原因是ReGAN利用了StyleGAN，它在人脸上的表现自然比SD要好得多。由于ReGAN具有类别特定的GAN模型，因此可能会为某些部分产生更好的结果，但它并不是一个通用的框架。
与ReGAN相反，作者可以根据SD的语义知识，无需对其类别进行预训练，对其他对象（例如乌龟）进行分割。为面部图像量身定制的SD将有助于消除这种差异，但会降低通用性。
4.3 更多定性结果
4.3.1 视觉结果
在这一部分中，作者展示了作者的方法的多功能性，它可以在训练时针对被遮挡的对象进行训练，然后在没有遮挡的测试图像上进行推理，或者反过来，可以在完全可见的对象上进行训练，并对遮挡对象进行预测。
这展示了作者的方法具备理解部分和对象语义的能力。图10是这一情况的示例。尽管训练图像中由于一个人导致了目标区域的遮挡，但作者的方法能够分割该类别中的其他图像。另外，考虑图11作为另一个示例，其中训练图像在这个实例中完全可见，而测试图像则受到遮挡。
这突显了作者的技术在准确捕捉语义时的鲁棒性，即使在面对遮挡对象或片段化分割时也是如此。
4.3.2 与基线的进一步比较
在这一部分，作者首先呈现了与ReGAN的定性比较，如图14所示。显然，SLiMe在第二行和第三行中展示了更复杂的头发分割，显示出与ReGAN相比，它能够捕捉更精细的细节。此外，在第二行中，ReGAN生成的耳朵分割似乎噪音较多。
接下来，在图15中，作者使用两种伪装动物（螃蟹和蜥蜴）对作者的方法和SegGPT进行了比较分析。值得注意的是，SLiMe在这些动物的精确分割中取得了成功，即使在人眼难以检测到它们的情况下。
4.3.3 SLiMe的通用性
SLiMe的另一个值得注意的特点是其通用性，如图12所示。这张图表明，尽管只在一个包含分割部分的狗图像上进行了训练，SLiMe可以理解头部、身体、腿部和尾巴的概念，并将其有效地应用于来自不同类别的测试图像。
然而，图13说明，当在包含狗和牛的图像上进行训练时，SLiMe擅长于学会专门分割狗类别的测试图像。这两个图表突出了SLiMe的能力，可以获取高级概念或专属对象类别。
4.4 消融研究
本节中的所有实验都是在PASCAL-Part数据集的车类别的10个样本上进行的。在表5中，作者展示了文本提示消融实验的结果，以展示SLiMe对于初始文本提示选择的稳健性。"部分名称"列显示了使用部分名称而不是空白提示的结果。在下一列中，标为"部分"，作者使用术语"部分"而不是部分名称，最后一列中，文本提示为空。通过比较这些列，作者可以看到初始文本提示对结果的影响很小。
接下来，作者进行了实验，以确定作者损失的最佳系数。如表6中α下的列所示，LSD的最合适系数是0.005。此外，通过比较最后两列，可以明显看出LMSE的重要性。
作者还提供了更多的实验结果。首先，作者通过改变的值进行了实验，t对应于输入到SD的噪声调度器的时间步骤，用于控制噪声添加。选择适当的t值非常重要，因为它决定了注入到图像嵌入中的噪声程度。选择较高的时间步骤会导致图像嵌入过于嘈杂。
因此，较小的值更可取。在表8中，作者展示了SLiMe在PASCAL Part数据集的马类的mIoUs。结果突显了SLiMe中的默认值，即，在其他值上表现更好。
第二个受剔除的参数是学习率。选择正确的学习率非常重要，因为较高的值可能导致与SD可以理解的文本嵌入相差甚远。相反，较低的学习率可能不会引入足够的改变到嵌入中。通过在不同学习率下训练SLiMe，作者在表9中展示了作者方法的最佳学习率为0.1，与SLiMe中使用的值一致。
作者还进行了另一项研究，以展示SLiMe中自注意力图的有效性。如表10所示，使用自注意力图与交叉注意力图一起可以提高mIoU达到11.0%。
5、参考
[1].SLIME: SEGMENT LIKE ME
点击上方卡片，关注 「AI视界引擎」 公众号