点击下方卡片,关注「集智书童」公众号
本文主要解决了什么问题
单一域泛化(SDG)在医学图像分割中的挑战 :现有的方法在跨不同成像协议、扫描仪供应商和临床地点的未见领域上表现不佳,尤其是在形状先验知识的捕捉和泛化方面存在局限性。
字典学习方法的局限性 :传统的字典学习方法在固定的小规模字典下表征能力有限,而当字典规模增大时容易过拟合,且与大型基础模型(如SAM)兼容性不佳。
如何高效利用语义形状先验 :需要一种方法能够动态地选择和组合形状信息,以增强模型对未见领域的泛化能力。
本文的核心创新是什么
提出MoSE框架 :将专家混合(MoE)训练的思想与字典学习结合,每个字典原子被概念化为一个“形状专家”,专门负责编码不同的语义形状信息,通过门控网络动态融合这些专家以生成鲁棒的形状图。
与SAM的双向集成 :将生成的形状图作为Prompt提供给SAM,利用其强大的泛化能力,同时通过Sparse激活防止过拟合,实现了端到端的训练。
解决字典规模与性能的权衡 :通过增加形状专家的数量(从256扩展到1024),持续提升性能,避免了传统方法中因字典规模增大而导致的过拟合问题。
平衡和预热策略 :引入变异系数(CV)正则化和
范数惩罚,确保门控网络在训练初期能够有效激活形状专家,并在后期实现严格的Sparse性。
结果相较于以前的方法有哪些提升
Dice系数显著提高 :在多个未见过的目标域上,MoSE的平均Dice系数达到91.4%,超越了其他最先进的方法(如D-Norm 71.0%、CCSDG 88.4%、SAMed 85.4%、DAPSAM 88.4%)。
Hausdorff距离明显降低 :MoSE的平均HD为17.0 mm,远低于D-Norm(75.8 mm)和CCSDG(24.4 mm),表明其在边界定位上的优越性。
字典规模扩展的优势 :与传统方法相比,MoSE在字典规模从128增加到1024的过程中性能持续提升,而传统方法在规模超过256后出现过拟合。
局限性总结
计算资源限制 :由于GPU内存限制,作者未能尝试更大的字典规模(如超过1024),可能影响进一步性能优化。
依赖SAM的基础能力 :尽管MoSE增强了SAM的泛化能力,但其效果仍受限于SAM本身对特定医学任务的支持程度。
超参数敏感性 :模型性能对超参数(如
、
、Top-K值)较为敏感,需要精细调参以达到最佳效果。 4. 4. 当前仅适用于2D图像 :MoSE目前专注于2D医学图像分割,未来需扩展至多类形状和3D基础分割模型以应对更复杂的任务。
导读
单一域泛化(SDG)在医学图像分割领域近年来备受关注。一种有前景的SDG策略是利用跨不同成像协议、扫描仪供应商和临床地点的语义形状先验知识。然而,现有的编码形状先验的字典学习方法,在少量离线计算的形状元素集合下往往表现出的表征能力有限,或者当字典规模增大时会出现过拟合。此外,它们与大型基础模型(如Segment Anything Model(SAM))的兼容性不佳。在本文中,作者提出了一种新颖的形状专家混合(Mixtureof-Shape-Experts,MoSE)框架,将专家混合(MoE)训练的思想无缝集成到字典学习中,以高效地捕获多样且鲁棒的形状先验知识。MoSE将每个字典原子概念化为一个“形状专家”,专门负责编码不同的语义形状信息。一个门控网络将这些形状专家动态融合成一个鲁棒的形状图,通过SAM编码引导的Sparse激活来防止过拟合。作者进一步将此形状图作为 Prompt 提供给SAM,通过双向集成利用SAM强大的泛化能力。所有模块,包括形状字典,均以端到端的方式进行训练。在多个公共数据集上的大量实验证明了其有效性。
- 引言
准确界定病灶或解剖结构是临床诊断、干预和治疗规划的基本步骤[17, 31-33]。尽管深度学习方法近期在分割任务上展现了卓越性能,但在跨不同领域(如不同临床场所、扫描设备供应商和成像参数)时面临泛化挑战[16, 34, 36, 37]。这些模型在评估未见过的领域数据时往往出现显著性能下降,限制了其实际应用。为解决这一问题,研究者提出了领域泛化(DG)方法[35],旨在使模型具备跨未见领域泛化的能力。这些方法通常基于多源领域泛化协议,依赖于获取多个源领域数据,但由于隐私问题和临床常规中数据共享的复杂性,这种获取方式往往不切实际。
相反,单域泛化(SDG)[9, 11, 14, 29, 55, 64] 越来越受到关注,因为它仅在一个源域上进行训练,但旨在在不同目标域上保持鲁棒性能。近期研究[5, 29, 60]强调语义形状信息在不同医学域间保持一致,使其成为SDG的有价值先验。例如,[29]进一步整合了字典学习来捕获语义形状先验,通过求解迭代优化问题,将每个分割 Mask 表示为形状字典中一组基元素(或原子)的线性组合。然而,这种方法对字典大小高度敏感,难以实现脆弱的平衡(例如,图3(a))。在应用于医学SDG分割时,固定的较小源域字典难以捕捉未见域中形状的多样性。虽然增加字典大小以涵盖更广泛的形状变化会引入过拟合风险,并导致模型记忆特定特征而非有效泛化。此外,[29]中的形状字典在两个阶段独立于分割定义,存在次优问题。
此外,近年来蓬勃发展的基础大模型[12, 13, 39],例如Segment Anything Model (SAM)[24],已展现出强大的泛化能力。作为自然图像分割中的可 Prompt 基础模型,SAM在许多零样本场景中表现出色。尽管由于自然医学数据间存在较大差异[38, 63],它仍难以应对特定的医学分割任务,且需要通过 Mask 或点 Prompt 来指示待分割结构。此外,如何将字典学习整合到SAM流程中尚未得到充分探索。
受这些挑战的启发,作者提出了一种新颖的形状专家混合(Mixture-of-Shape-Experts,MoSE)框架,该框架将专家混合(Mixture-of-Experts,MoE)训练的优势与字典学习相结合。MoE方案能够存储多样化的专家集合,通过端到端网络基于输入Sparse地选择相关子集[4, 45, 52]。在MoSE中,每个字典原子被概念化为一个"形状专家",它编码了从源域中提取的独特语义形状信息。作者的可学习形状专家无需在分割训练之前进行离线计算。门控网络动态地为每个输入图像的SAM嵌入生成Sparse编码,选择并组合最相关的形状专家,形成鲁棒且域不变的形状图。这种Sparse激活使模型能够捕捉更广泛的形状变化,降低过拟合的风险。此外,该形状图可以无缝地重新整合回SAM作为 Prompt ,从而指导分割过程,并显著提高在未见过的域上的泛化能力。
作者的主要贡献是:
- • 作者提出了一种端到端的形状字典框架MoSE,该框架采用MoE策略,通过动态门控机制高效地捕获多样且鲁棒的形状先验。它通过Sparse激活选择性地融合形状专家,以缓解过拟合问题。
- • 一种与SAM流程的双向集成方案,利用其强大的编码器进行Sparse编码,同时将形状图作为 Prompt 。通过利用平衡和预热策略进行正则化。
- • 与传统的字典学习方法不同,通过将字典大小(形状专家的数量)从256增加到1024,作者能够持续获得性能提升。大量实验表明,作者提出的MoSE在性能上显著优于最先进的SDG方法以及适配的SAM。
- 相关工作
2.1. 单域泛化
领域泛化(DG)[7, 23]旨在将一个在多个源域上训练的模型泛化到分布差异较大的目标域。以往在领域泛化方面的进展主要在于使模型能够从多源域数据中学习,以减少模型偏差。
近期,单域泛化(SDG)[9, 11, 14, 29, 55, 64] 越来越受到关注,其特点是在训练过程中仅使用一个源域,而模型需在多个未见过的目标域上进行评估。这尤其具有挑战性,因为不同医疗中心在成像协议、设备和患者群体上存在差异。在图像分类方面[10, 44, 50, 53, 61],已有先驱工作提出在单源域训练时提升模型在未见过的域上的性能。特别是[44, 50, 53]提出采用数据增强方案,通过对抗训练生成多样化的输入图像。而[10, 61]则提出利用归一化机制来适应特征分布以适应未见过的域。该方法也已被应用于目标检测[9, 57]。[49]进一步提出利用自监督视觉语言模型来指导目标检测器的训练。
在分割任务中,[43, 46] 也采用数据增强方法,使分割模型接触合成域偏移的训练数据。类似地,[58] 通过学习对抗域合成器来合成新域。此外,[15] 使用图像及其风格增强对应图像的浅层特征进行对比训练。然而,SDG 的有效性直接受到生成数据真实性的影响,这也是一个长期存在的挑战。近期研究 [5, 29, 60] 强调了不同医学域间语义形状信息的一致性,将其确立为 SDG 的宝贵先验知识。如何高效利用这种先验知识是一项具有挑战性的任务。
2.2. 字典学习
旨在寻找一组基本元素来构建字典,使得给定输入能够通过这些学习到的元素的良好Sparse线性组合来表示[30, 48]。将字典学习融入医学图像分割通常涉及学习一组基元素(即字典),这些元素能够Sparse地表示图像块[47, 62]。它已被应用于通过捕获医学图像的基本结构特征来提升分割性能。例如,将字典学习与聚类算法相结合,导致了无监督自适应分割方法,其中字典作为聚类中心,Sparse表示用于分割[62]。SDG分割的一个显著方法是通过对可用的源域数据进行字典学习来整合形状先验[29]。通过从可用的源域数据中构建形状字典,模型能够捕获跨域不变的语义形状先验。然而,字典大小和性能之间存在权衡[29]。这提出了一个重要问题:如何增加字典大小以涵盖更广泛的形状先验范围,同时同时降低过拟合的风险?
2.3. 专家混合模型
其基础是一个简单而有效的概念:不同的模型组件,称为专家,被训练以专注于不同的任务或数据特征[4]。它最初在[18, 21]中提出,在基础模型时代[40, 42]经历了广泛的探索和进步。在MoE框架中,仅对给定输入中最相关的专家进行激活,优化计算效率的同时利用多样化的专业知识。这种灵活且可扩展的方法符合规模法则,能够在不按比例增加计算成本的情况下实现更大的模型容量。[51]提出将不同的微调SAM模型与MoE相结合。在这项工作中,作者引入了一个新的概念——形状专家混合(MoSE),以探索MoE的效率。它与网络混合本质上不同。
- 方法论
在SDG设置中,作者提供了一个单一源域数据集
,其中包含
对源域输入切片
及其对应的分割 Mask
。这里,
、
和
分别代表高度、宽度和分割类别。在源域数据集上训练后,模型在未见过的目标域
上进行性能评估。作者提出的MoSE框架的概述如图5所示。
3.1. 形状专家混合 (MoSE) 训练
在传统的字典学习方法[29]中,字典是离线构建的,具有固定的原子,这些原子与回归分支生成的系数相结合,以产生形状先验。然而,这种方法在医学SDG分割中存在两个关键局限性:1) 固定的小源域字典难以捕捉未见目标域中的多样形状,阻碍泛化。2) 字典的离线构建与分割任务脱节,限制了其适应性和有效性。
受专家混合(Mixture-of-Experts, MoE)训练系统[45]的启发,该系统通过Sparse激活专家在保持计算成本不变的情况下显著提升模型容量,作者提出将MoE集成到字典学习中。在MoSE中,字典的原子被概念化为"形状专家",每个专家存储从源域数据中学习到的多样化形状先验。具体而言,作者定义
个可学习的形状专家[3, 8]为
,其中每个
。为减少参数开销并匹配SAM Prompt 编码器的大小,作者设置
和
。随后,作者采用轻量级卷积神经网络(CNN)作为门控网络,该网络生成Sparse编码来组合形状专家。具体来说,对于通过SAM图像编码器嵌入到
的每个输入图像
,门控网络生成输出
,该输出为形状专家分配像素级权重。为强制Sparse性并仅选择最相关的专家,作者应用基于绝对值的Top-K选择策略,在每个像素处保留响应最高的
个形状专家:
其中
表示在每个像素位置上与
绝对值最大的前
个值对应的索引集合。最终形状图则计算如下:
其中
表示第
个形状专家的Sparse化门控权重图,
表示在空间维度上进行逐元素乘法的广播。直观上,yshape 中的每个像素
都是由门控网络认为最相关的 top-
专家的线性组合形成的。此外,通过这种逐像素组合生成的形状图
,随后与标签
结合,用于比较分割损失,从而实现端到端的训练,具体细节详见第 3.3.1 节。
3.2. 将Shape Map集成到SAM流程中
作者的MoSE无缝地将SAM与双向方案相结合,以充分利用其泛化能力。除了辅助Sparse编码外,受 Prompt 学习[20, 27, 28, 54]的启发,作者将形状图Whape作为 Prompt 引入,使其能够与图像嵌入交互。这种集成引入了先验形状信息,指导分割过程,并增强模型对未见领域的泛化能力。
具体而言,形状映射9shape首先通过Sigmoid激活函数进行处理,以将其值归一化到区间[0,1]内,确保与 Prompt 表示的兼容性:
其中
表示Sigmoid函数。处理后的形状 Prompt
被集成到SAM框架
中,与图像嵌入交互以指导 Mask 解码器。最终得到分割预测结果:
3.3. 训练目标和策略
3.3.1. 分割损失
最终的分割预测
和中间形状图输出
均使用标签
进行监督。遵循SAMed [59],作者采用复合分割损失,结合交叉熵(CE)损失和Dice损失来训练网络模块和形状字典,该单一源域定义为:
其中
是权衡因子。
表示
或
或
中的
。
3.3.2. MoSE训练中的利用平衡和预热策略
在形状专家混合(MoSE)框架中的门控网络倾向于强化ego不平衡状态,将高权重分配给少数形状专家。为鼓励均衡使用,作者对Sparse编码引入了变异系数(CV)正则化[1, 45]:
最小化这一项鼓励专家被更均衡地调用。尽管在整个数据集上计算这一正则化项是不可行的,但作者选择在批次 Level 进行估计和最小化。
此外,门控网络的硬Sparse化可能会阻止梯度反向传播,导致其在初始状态停滞。为解决这个问题,作者提出延迟硬阈值处理,并使用
范数惩罚来初始训练MoSE层。
在[52]之后进行了几次迭代。随后,作者启用Top-K选择以在
中实施更严格的Sparse性。
3.3.3. 总损失
将所有术语综合起来,作者的总训练损失是
其中
控制着
预热惩罚或 CV 正则化在迭代次数 Twarm-up 时的相对重要性。
3.4. 未知领域的推理
在
上训练后,作者的MoSE with SAM模型可以通过计算,无需对
进行额外适配,即可处理未知的源域图像
。
此处,
是从
生成的目标嵌入。最终的分割 Mask 是
,它受益于字典中封装的形状先验。
- 实验与结果
4.1. 数据集和预处理
作者从五个公开数据集收集了肝脏CT扫描图像:A) BTCV [25](1542个切片)、B) LITS [2](19160个切片)、C) CHAOS [22](2341个切片)、D) AMOs的评估集 [19](4865个切片)和E) WORD的评估集 [41](1110个切片)。
所有切片均重采样为统一的轴向分辨率
像素。作者选择 BTCV 数据集作为训练的源域,并在其他四个未见过的域上评估模型的性能。对于 B-E 的每个域,作者采用 20/80 的分割方式用于验证和测试。
随后,每个卷通过减去其最小强度值并除以其强度范围进行归一化处理。这些数据集的详细信息如表1所示。本研究中的所有数据集都经过预处理,采用一致的风窗宽度(WW)和风窗水平(WL)值。WW设置为600,对于大多数数据集,WL通常设置为100,依据[26]的研究,但CHAOS数据集除外,其使用更高的WL值1100,更适合肝脏。此外,对于LITS数据集,作者按照[56]中采用的相同预处理方法,将肝脏肿瘤类别合并到肝脏类别中。
4.2. 实现细节
作者采用vanilla SAM [24],并按照[6]中提出的设计,在每个编码层引入两个可训练的MLP结构 Adapter ,将其作为作者的 Baseline 模型。基于超参数验证,形状专家的总数为
,Top-K选择参数设置为
。作者进一步在消融研究中检验这些设置的影响。所有实验均使用SAM的'ViT-B'版本进行。超参数
和
分别设置为
和
。作者经验性地将
设置为0.8,以平衡两个损失组件的贡献。
可训练的MLP结构 Adapter 的秩设置为4,这有助于在计算效率和性能之间取得平衡。作者使用学习率
初始化训练,并采用权重衰减为0.1的AdamW优化器。为了进行公平比较,所有方法的训练过程均被限制在最多150个epoch。由于GPU内存限制,作者没有尝试更大的
值,尽管进一步增加
可能会带来更好的结果。
4.3. 单域泛化结果
采用Dice系数[ % ](Dice,越高越好)和Hausdorff距离[mm](HD,越低越好)两种广泛使用的指标定量评估分割性能。在表2中,作者将所提出的MoSE与其他最先进的SDG方法进行比较,包括基于CNN的DNorm [64]和CCSDG [14],以及基于ViT的DeSAMB [11]、DeSAM-P [11]和DAPSAM [55]。
MoSE在所有领域中实现了最高的平均Dice系数91.4%,超越了其他方法,如D-Norm(71.0%)、CCSDG(88.4%)、SAMed(85.4%)和DAPSAM(88.4%)。特别地,MoSE在所有单个目标领域中表现优异,在LITS(B)上取得了最高的Dice分数92.7%,在CHAOS(C)上为90.3%,在AMOS τ (D) A上为91.6%,在WORD(E)上为90.9%。在Hausdorff距离(HD)方面,MoSE也表现出色,其平均HD最低为17.0 mm,显著低于D-Norm(75.8 mm)和CCSDG(24.4 mm)等方法。具体而言,MoSE在AMOS(D)上实现了最低的HD为9.7 mm,在CHAOS(C)上为25.6 mm,同时在其他领域中保持了具有竞争力的性能。另一基于ViT的方法DAPSAM在HD方面表现优于大多数基于CNN的方法,其平均HD为27.6 mm。
在消融研究中,不包含MoE的MoSE模型在平均Dice指标上比 Baseline 模型高出1.2%。更令人瞩目的是,作者的MoSE方法进一步实现了平均Dice指标提升3.2%,并显著降低了平均HD值,从25.3 mm降至17.0 mm。如图2所示,形状图能够精确聚焦于目标器官,为后续分割提供形状先验信息和位置信息作为 Prompt 。这种引导有助于模型更好地定位和分割器官,从而提高分割精度和跨未见领域的泛化能力。
4.4. MoE架构的有效性
在图3(a)中,随着字典规模从128增加到256,传统方法的性能得到提升,在
时达到峰值,而对于更大的
则出现过拟合现象。相比之下,MoSE架构随着形状专家数量(记为
)的增加,性能持续提升。MoSE取得的最佳结果是在
时平均Dice系数达到
,明显优于传统方法。这突显了作者MoSE方法的优势,该方法有效克服了字典规模与性能之间的权衡,在不产生传统字典学习方法中观察到的过拟合风险的情况下,提供了更优的结果。
4.5. 超参数敏感性分析
如图3(b)所示,当
小于
时,门控网络上的惩罚力度过弱,难以防止将较高的权重分配给少数形状专家。随着模型过度依赖一小部分专家,其泛化性能会下降。相反,较大的
可能导致形状专家的选择过于均匀。在这种情况下,门控网络无法适当优先考虑与特定任务最相关的专家,导致模型无法有效利用形状专家的专业知识。在图3(c)中,作者发现将
设置为
能够在训练中取得平衡。当
小于
时,形状专家在训练初期的激活不足。Top-K选择会阻断梯度,导致某些形状专家无法学习到有用的形状信息。反之,如果
过大,形状专家Sparse化的训练过程会变得缓慢,最终阻碍泛化性能。
4.6. 形状专家的数量和Top-K
这些是作者MoSE的关键因素。如图6所示,
的网格搜索范围设置为128、256、512和1024,而
的范围从64到1024。增加
(即存储更多样化的形状专家)通常能带来更好的性能,因为它允许模型存储更丰富的形状先验信息,从而增强其在未见过的领域中的泛化能力。关于
的选择,当
或
时,较小的
值能获得更好的性能,这可能是由于模型能从关注更少、更相关的形状专家中受益,从而提升其泛化能力。然而,当
时,较大的
值更有效,因为形状专家数量的增加需要更广泛的选取来有效地捕获最相关的形状先验。
4.7. 在未见领域的可视化结果
如图4所示,MoSE在感兴趣区域(ROI)的分割上实现了更高的准确性,其分割 Mask 与其他方法相比更精确地与解剖结构对齐,展示了其在未见过的领域中的优越泛化性能。
4.8. MoSE中形状专家的可视化
如图5所示,作者的MosE框架中的形状专家以 Heatmap 形式可视化,每个 Heatmap 代表1000oSE学习到的不同形状先验。每个专家以 Heatmap 表示,捕捉不同的形状先验。颜色强度反映了每个专家的激活程度。这些专家通过在每个像素处动态选择最相关的先验来为分割任务做出贡献。每个专家的激活强度各不相同,并且为每个像素的分割动态选择最相关的形状先验。
- 结论
作者提出了MoSE,一种利用MoE原理增强医学图像分割中SDG的新型端到端形状字典学习框架。与难以避免过拟合且使用次优独立计算的形状字典的传统字典学习方法不同,MoSE将每个字典原子概念化为一个独立的"形状专家",并采用门控网络动态选择最相关的专家。通过将学习到的形状图作为 Prompt 集成到SAM中,作者进一步利用大型基础模型的强表示能力,在多个未见过的目标域上取得了显著改进。
未来方向包括将MoSE扩展到多类形状和3D基础分割模型,以及探索其他基础模型和 Prompt 工程技术以超越器官级任务。作者相信,MoSE框架通过成功融合形状先验学习与基于 Prompt 的大型分割模型,为领域泛化的医学成像开辟了新的途径。
参考
[1]. Mixture-of-Shape-Experts (MoSE): End-to-End Shape Dictionary Framework to Prompt SAM for Generalizable Medical Segmentation
扫码加入👉「集智书童」交流群
(备注:方向+学校/公司+昵称)