MSA【3】：SAMed

---

前言

SAMed 是基于大规模图像分割模型 Segment Anything Model (SAM) 构建的，旨在探索将大规模模型定制化应用于医学图像分割的新研究范式

原论文链接：Customized Segment Anything Model for Medical Image Segmentation

---

1. Abstract & Introduction

1.1. Abstract

本文提出了医学图像分割的通用解决方案 SAMed

• SAMed 对 SAM 图像编码器采用基于低秩的微调策略（LoRA），并在标注的医学图像分割数据集上对其与提示编码器和掩码解码器一起进行微调
• Warmup 微调策略和 AdamW 优化器使 SAMed 成功收敛并降低了损失

与 SAM 不同，SAMed 可以对医学图像进行语义分割

1.2. Introduction

由于缺乏医学图像数据及其相应的语义标签，大规模 CV 模型无法直接用于解决医学图像分割问题

• 大规模 CV 模型根据强度方差来决定不同分割区域的边界，这在自然图像中是合理的，但在医学图像中却不合理
  • 基础的分割大模型可以通过识别强度值发生显著变化的区域，以判断不同区域之间可能存在的边界
  • 但是，医学图像包含复杂的结构和纹理，需要专门的算法和模型才能精确分割
    • 在核磁共振成像图像中，不同的组织具有不同的信号强度，它们之间的界限并不总是明确定义的
    • 医学图像分割需要能够考虑图像的解剖学和病理特征（例如形状、纹理和背景）的模型
• 大规模 CV 模型无法将分割区域与有意义的语义类别联系起来

---

2. Methods

2.1. Overview

给定一幅空间分辨率为 $H\times W$、通道数为 $C$ 的医学图像 $x\in {\mathbb{R}}^{H\times W\times C}$，其中每个像素都属于预定义类别列表 Y = { y 0 , y 1 , ⋯   , y k } Y = \{y_0, y_1, \cdots, y_k \} Y={y0​,y1​,⋯,yk​} 中的一个元素（$y_0$ 通常被视为背景）

• SAMed 冻结了图像编码器中的所有参数，并为每个变压器模块设计了可训练的旁路
  • 如 LoRA 所示，这些旁路首先将变换器特征压缩到低级空间，然后将压缩后的特征重新投影，使其与冻结变换器块中输出特性的通道保持一致
• SAMed 没有使用提示编码器，因为 SAMed 在推理过程中无需任何提示即可执行自动分割
  • 如果剥离 SAM 中的所有提示，SAM 将更新默认的默认嵌入，因此 SAMed 还会在训练过程中对该嵌入进行微调
• 使用 LoRA 对变换解码器进行微调是可选的
  • 如果冻结变换解码器，并使用 LoRA 层对其进行微调，但不对其所有参数进行微调，则可以进一步缩小更新参数的模型大小，以方便部署，但性能会略有下降
• SAM 的原始分割头输出多个分割掩码，以解决分割提示的模糊性
  • SAMed 也会预测多个分割掩码，但每个掩码代表 $Y$ 中的一个类别，因此 SAMed 会预测 k 个分割掩码

2.2. LoRA in image encoder

给定编码令牌序列 $F\in {\mathbb{R}}^{B\times N\times C_{in}}$ 和由投影层 $W\in {\mathbb{R}}^{C_{out}\times C_{in}}$ 操作的输出令牌序列 $\hat{F}\in {\mathbb{R}}^{B\times N\times C_{out}}$

• 投影层
  • 将输入数据从原始空间映射到新空间的层，新空间通常具有不同的维度
  • 投影层通常由一组可学习的参数组成，如权重和偏置，这些参数在训练过程中得到优化

LoRA 假定 $W$ 的更新应该是渐进和稳定的，SAMed 的 LoRA 策略如下所示：

• SAMed 首先冻结变压器层以保持 $W$ 固定不变，然后添加一个旁路来完成低秩近似
  • 该旁路包含两个线性层 $A\in {\mathbb{R}}^{r\times C_{in}}$ 和 $B\in {\mathbb{R}}^{C_{out}\times r}$
  • 其中 r ≪ { C i n , C o u t } r \ll \{ C_{in}, C_{out} \} r≪{Cin​,Cout​}

因此更新层 $\hat{W}$ 的处理过程可描述为

由于多头自关注机制是通过余弦相似度来确定要关注的区域，因此对 $Q$、$K$ 或 $V$ 的投影层应用 LoRA 来影响关注分数是明智的

观察发现，将 LoRA 应用于 $Q$ 和 $V$ 投影层时，SAMed 可以获得更好的性能，因此多头自关注的处理策略将变为

参数含义：

• $W_q$、$W_k$ 和 $W_v$ 是 SAM 中冻结的投影层，$A_q$、$B_q$、$A_v$ 和 $B_v$ 是可训练的 LoRA 参数

2.3. Prompt encoder and mask decoder

由于 SAM 中的提示编码器在没有提示时使用默认嵌入，因此 SAMed 保留了这一默认嵌入，并在微调过程中对其进行训练

SAM 中的掩码解码器由一个轻量级转换层和一个分割头组成

• 可以选择将 LoRA 应用于轻量级转换器层并对分割头进行微调

• 或者直接微调掩码解码器中的所有参数
  
  SAMed 对 SAM 的分割头稍作修改，以定制 $Y$ 中每个语义类别的输出，如图所示：

• 与 SAM 的模糊预测不同，SAMed 以确定性的方式预测 $Y$ 中的每个语义类别

• SAMed 的掩码解码器会同时预测 k 个语义掩码 ${\hat{S}}_l\in {\mathbb{R}}^{h\times w\times k}$，与每个语义标签相对应

• 最后，预测的分割图生成为

  • 其中，$d=-1$ 表示在最后一个维度（通道维度）上执行 Softmax 和 argmax 操作
    

2.4. Training strategies

2.4.1. Loss function

SAMed 采用交叉熵和骰子损失来监督微调过程

2.4.2. Warmup

• SAMed 采用预热来稳定训练过程，并在训练初期熟悉医疗数据
• SAMed 利用指数学习率衰减使训练过程逐渐收敛
  

2.4.3. AdamW optimizer

---

总结

问题总结：

1. 自然图像分割更关注于边界，但解剖或病理结构的分析在医学图像分割中同样起着至关重要的作用
2. SAM无法用医学图像做到语义分割

贡献点：

1. 首先对SAM进行扩展，探索其在医学图像语义标签分割上的能力
2. 考虑到性能、部署和存储开销，提出了图像编码器的自适应和一系列调优策略