MSA【4】：DeSAM

文章目录

前言
1. Abstract & Introduction
2. Methods
总结

前言

基于深度学习的自动医学图像分割模型经常会出现领域偏移的问题，即在源领域训练的模型不能很好地泛化到其他未见领域。为此本文提出 DeSAM，通过解耦 image embedding 和 prompt token 来提高 SAM 的效果

原论文链接：DeSAM: Decoupling Segment Anything Model for Generalizable Medical Image Segmentation

1. Abstract & Introduction

1.1. Abstract

基于深度学习的自动医学图像分割模型经常会出现领域偏移的问题，即在源领域训练的模型不能很好地泛化到其他未见领域

Segment Anything Model（SAM）在提高医学图像分割的跨域鲁棒性方面展现出了潜力。然而，与人工提示相比， SAM 及其微调模型在全自动模式下的表现要差得多。经过进一步调查，本文发现性能下降与提示不佳和掩膜分割的耦合效应有关

在全自动模式下，不可避免的不良提示（如遮罩外的点或明显大于遮罩的框）会严重误导遮罩的生成

1.2. Introduction

1.2.1. Brief Introduction

深度模型对未知域外数据的泛化能力可能较差，这阻碍了模型在临床环境中的应用。为了缓解领域转移造成的性能下降，以往的尝试主要集中在无监督领域适应和多源领域泛化，这些方法却也存在以下局限性：

基于输入空间的增强需要专业知识来设计增强函数
基于特征的增强通常需要复杂的对抗训练

本文关键概念：

单源域泛化
- 仅使用一个源域的训练数据来训练对未见数据具有鲁棒性的深度学习模型
- 主要的解决方案包括基于输入空间的数据增强和基于特征的数据增强
耦合效应
- 图像嵌入和提示标记在 SAM 掩码解码器的交叉注意变换层中相互作用，使得最终输出的掩码高度依赖于提示
- 即使经过微调，模型仍倾向于对错误的提示（即不在掩码中的点或明显大于掩码的方框）更加敏感

1.2.2. Motivation

将基于大型数据集的模型直接移植到医学图像分割中以提高泛化能力是一种很有吸引力的方法

Segment Anything Model（SAM）在超过 10 亿个遮罩上进行了训练，在各种自然图像上实现了前所未有的泛化能力。

将 SAM 应用于医学图像分割有两种主要方法，这两种主要方法的局限性在于

冻结 SAM 的图像编码器和提示编码器
- 仅对掩码解码器进行微调的方法通常会产生次优结果
使用 adaption 或 visual prompt 技术训练图像编码器的方法可提高模型在特定领域的性能
- 然而，由于图像嵌入无法提前计算，训练这样的模型需要消耗大量 GPU 内存
- 此外，这些微调方法仍然需要使用人工给出的方框或点，因此很难实现全自动医学影像分割

直接通过 SAM 实现自动分割有两种方法

一种是使用网格点作为提示，本文称之为网格点模式
另一种是使用与图像大小相同的方框作为提示，本文称之为全方框模式

然而，即使经过完全微调，全自动 SAM 也往往会产生大量假阳性掩膜，性能远不能满足临床要求

1.2.3. Contribution

本文将 SAM 的掩码解码器解耦为两个子任务

prompt-relevant IoU regression (PRIM)
- 根据给定的提示预测 IoU 分数并生成掩码嵌入
prompt-invariant mask learning (PIMM)
- 将图像编码器的图像嵌入与 PRIM 的掩码嵌入融合在一起，生成掩码

2. Methods

2.1. Architecture

在这里插入图片描述

2.1.1. DeSAM Encoder

在训练过程中冻结图像和提示编码器

不对图像编码器进行任何调整，以便在训练阶段之前预先计算图像嵌入
DeSAM 训练过程不需要加载图像编码器，可以更高效地利用 GPU

虽然冻结了提示编码器，但在推理过程中还是用到了提示（分割一切模式）

2.1.2. Prompt-Relevant IoU Module (PRIM)

PRIM 的结构与 SAM 的掩码解码器类似，包括一个交叉注意变换层和一个 IoU 预测头

Attention & Cross Attention
- SAM 中的 Attention
  - $Q$ : image embedding，表示图像中的 features（可以理解为，编码后的分割对象的特征）
  - $K$ : tokens（prompt 编码后得到），引导模型学习期望的 feature（学习后可以分割出对应的物体 / 器官），换句话说，想从包含图像所有 feature 的 $Q$ 中提取需要的 feature，需要计算和包含对应所有 prompt 信息的 $K$ 的匹配程度，和prompt 匹配程度越高，越是期望学习到的特征
  - $V$ : 一个额外的参数，加权 $Q$ 和 $K$ 学到的内容
- Cross Attention
  - 基于跨模态的注意力机制
  - Transformer 架构中混合两种不同嵌入序列的注意机制
  - 两个序列必须具有相同的维度两个序列可以是不同的模式形态（如：文本、声音、图像）
  - 一个序列作为输入的 $Q$ ，定义了输出的序列长度，另一个序列提供输入的 $K$ & $V$
  - Example：如果 $Q$ 是图像， $K$ & $V$ 是文本，注意的应该是文本的图片

为了使提示和输出掩码脱钩，本文只舍弃了掩码预测头，从交叉注意变换层中提取掩码嵌入

2.1.3. Prompt-Invariant Mask Module (PIMM)

采用了经典的编码器-解码器结构

从图像编码器的全局注意力层 $i$ 提取图像嵌入
- 首先，从全局注意力层 $i = (8, 16, 24)$ 中提取形状为 $1024 \times 64 \times 64$ 的图像嵌入
- 其次，将图像嵌入信息传递给不同数量的 squeeze and excitation residual blocks (SEResBlock) 并进行上采样操作
  - SEResBlock
    - squeeze: 包括全局平均池化，将输入特征的空间维度减少到单通道
    - excitation: 学习通道依赖关系并产生一组权重。然后，这些权重将用于对原始特征重新加权，从而使网络专注于最重要的通道
  - 上图是 SEResBlock 的示意图，给定一个输入 $x$ ，其特征通道数为 $c_1$ ，通过一系列卷积等一般变换后得到一个特征通道数为 $c_2$ 的特征。与传统的 CNN 不一样的是，接下来通过三个操作来重标定前面得到的特征。
    - 首先是 Squeeze 操作，顺着空间维度来进行特征压缩，将每个二维的特征通道变成一个实数，这个实数某种程度上具有全局的感受野，并且输出的维度和输入的特征通道数相匹配；它表征着在特征通道上响应的全局分布，而且使得靠近输入的层也可以获得全局的感受野
    - 其次是 Excitation 操作，它是一个类似于循环神经网络中门的机制。通过参数 $w$ 来为每个特征通道生成权重，其中参数 $w$ 被学习用来显式地建模特征通道间的相关性
    - 最后是一个 Reweight 的操作，将 Excitation 的输出的权重看做是进过特征选择后的每个特征通道的重要性，然后通过乘法逐通道加权到先前的特征上，完成在通道维度上的对原始特征的重标定
- 最后，通过跳过连接来融合图像嵌入
此外，我们还合并了 PRIM 的掩码嵌入和 PIMM 的瓶颈嵌入，以利用预训练权重并确保 PIMM 和 PIRM 在训练过程中的梯度流
- 简单来说，梯度流是将我们在用梯度下降法中寻找最小值的过程中的各个点连接起来，形成一条随（虚拟的）时间变化的轨迹，这条轨迹便被称作“梯度流”
- 本文应该是为了实现梯度共享

4.2. Training strategies

在训练过程中，我们会加载 SAM 的预训练权重，冻结图像和提示编码器，并微调 PIMM 和 PRIM 中的层

于自动分割包括网格点模式和全框模式，本文采用了两种不同的策略来训练所提出的模型：

grid points mode
- 训练过程中 $1 : 1$ 随机选择前后景中的点（需要 GT 进行辅助）
  - 一次性只能提供一个点作为 prompt
  - 为了防止随机选取的过程中喂进去的都是背景点，从而导致模型性能下降
- 推理过程中按一定间隔选点即可，此时 SAM 已经基本上具备了分辨前景和背景的能力
- loss 计算公式如下所示，权重为 $1, 1, 10$
whole box mode
- Box 的大小即图像的大小
- 由于 GT 必须在方框内，本文只对掩膜的生成进行监督

总结

本文提出了一种用于医学图像分割中的单源域泛化模型，DeSAM。它将掩码生成与提示解耦，并利用 SAM 的预训练权重，从而降低了对 GPU 的要求。DeSAM 激励解码器从鲁棒的图像嵌入中学习与提示无关的特征，并且由于融合了多个尺度的图像嵌入，具有很强的抵抗看不见的分布变化的能力。本文在多站点数据集上验证了 DeSAM 的性能，表明所提出的方法优于其他最先进的方法