SAMed
是基于大规模图像分割模型 Segment Anything Model (SAM) 构建的,旨在探索将大规模模型定制化应用于医学图像分割的新研究范式
原论文链接:Customized Segment Anything Model for Medical Image Segmentation
本文提出了医学图像分割的通用解决方案 SAMed
SAMed
对 SAM
图像编码器采用基于低秩的微调策略(LoRA
),并在标注的医学图像分割数据集上对其与提示编码器和掩码解码器一起进行微调Warmup
微调策略和 AdamW
优化器使 SAMed
成功收敛并降低了损失与 SAM
不同,SAMed
可以对医学图像进行语义分割
由于缺乏医学图像数据及其相应的语义标签,大规模 CV 模型无法直接用于解决医学图像分割问题
给定一幅空间分辨率为 H × W H \times W H×W、通道数为 C C C 的医学图像 x ∈ R H × W × C x \in \mathbb{R}^{H \times W \times C} x∈RH×W×C,其中每个像素都属于预定义类别列表 Y = { y 0 , y 1 , ? ? , y k } Y = \{y_0, y_1, \cdots, y_k \} Y={y0?,y1?,?,yk?} 中的一个元素( y 0 y_0 y0? 通常被视为背景)
SAMed
冻结了图像编码器中的所有参数,并为每个变压器模块设计了可训练的旁路
LoRA
所示,这些旁路首先将变换器特征压缩到低级空间,然后将压缩后的特征重新投影,使其与冻结变换器块中输出特性的通道保持一致SAMed
没有使用提示编码器,因为 SAMed
在推理过程中无需任何提示即可执行自动分割
SAM
中的所有提示,SAM
将更新默认的默认嵌入,因此 SAMed
还会在训练过程中对该嵌入进行微调LoRA
对变换解码器进行微调是可选的
LoRA
层对其进行微调,但不对其所有参数进行微调,则可以进一步缩小更新参数的模型大小,以方便部署,但性能会略有下降SAM
的原始分割头输出多个分割掩码,以解决分割提示的模糊性
SAMed
也会预测多个分割掩码,但每个掩码代表
Y
Y
Y 中的一个类别,因此 SAMed
会预测 k 个分割掩码给定编码令牌序列 F ∈ R B × N × C i n F \in \mathbb{R}^{B \times N \times C_{in}} F∈RB×N×Cin? 和由投影层 W ∈ R C o u t × C i n W \in \mathbb{R}^{C_{out} \times C_{in}} W∈RCout?×Cin? 操作的输出令牌序列 F ^ ∈ R B × N × C o u t \hat{F} \in \mathbb{R}^{B \times N \times C_{out}} F^∈RB×N×Cout?
LoRA
假定
W
W
W 的更新应该是渐进和稳定的,SAMed
的 LoRA
策略如下所示:
SAMed
首先冻结变压器层以保持
W
W
W 固定不变,然后添加一个旁路来完成低秩近似
因此更新层
W
^
\hat{W}
W^ 的处理过程可描述为
由于多头自关注机制是通过余弦相似度来确定要关注的区域,因此对
Q
Q
Q、
K
K
K 或
V
V
V 的投影层应用 LoRA
来影响关注分数是明智的
观察发现,将 LoRA
应用于
Q
Q
Q 和
V
V
V 投影层时,SAMed
可以获得更好的性能,因此多头自关注的处理策略将变为
参数含义:
SAM
中冻结的投影层,
A
q
A_q
Aq?、
B
q
B_q
Bq?、
A
v
A_v
Av? 和
B
v
B_v
Bv? 是可训练的 LoRA
参数由于 SAM
中的提示编码器在没有提示时使用默认嵌入,因此 SAMed
保留了这一默认嵌入,并在微调过程中对其进行训练
SAM
中的掩码解码器由一个轻量级转换层和一个分割头组成
可以选择将 LoRA
应用于轻量级转换器层并对分割头进行微调
或者直接微调掩码解码器中的所有参数
SAMed
对 SAM
的分割头稍作修改,以定制
Y
Y
Y 中每个语义类别的输出,如图所示:
与 SAM
的模糊预测不同,SAMed
以确定性的方式预测
Y
Y
Y 中的每个语义类别
SAMed
的掩码解码器会同时预测 k 个语义掩码
S
^
l
∈
R
h
×
w
×
k
\hat{S}_l \in \mathbb{R}^{h \times w \times k}
S^l?∈Rh×w×k,与每个语义标签相对应
最后,预测的分割图生成为
Softmax
和 argmax
操作SAMed
采用交叉熵和骰子损失来监督微调过程
SAMed
采用预热来稳定训练过程,并在训练初期熟悉医疗数据SAMed
利用指数学习率衰减使训练过程逐渐收敛问题总结:
贡献点: