点击下方卡片,关注「AI视界引擎」公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
Image
医学图像分割在临床应用中具有重要意义,然而数据隐私问题以及专家标注的成本限制了标注数据的获取。
联邦半监督学习(Federated Semi-Supervised Learning, FSSL)提供了一种解决方案,但面临两大挑战:伪标签的可靠性依赖于局部模型的性能,而客户端设备由于计算资源有限,通常需要轻量级或异构的模型架构。这些约束导致伪标签的质量和稳定性下降;尽管大型模型具有更高的准确性,但无法在客户端设备上进行训练或用于常规推理。
为此,作者提出 SAM-Fed,一种基于高容量分割基础模型(segmentation foundation model)指导轻量级客户端训练的联邦半监督学习框架。SAM-Fed 结合双知识蒸馏(dual knowledge distillation)与自适应一致性机制(adaptive agreement mechanism),以优化像素级监督信号。
在皮肤病变和息肉分割任务中,针对同质与异构设置的实验表明,SAM-Fed 始终优于现有的先进 FSSL 方法。
- 引言
深度学习在医学图像分析中取得了优异的性能,这得益于大规模标注数据集的支持 [1]。然而,由于隐私法规和伦理问题,跨机构共享医学数据受到限制 [2]。联邦学习(Federated Learning, FL)能够在不交换原始数据的情况下实现协作模型训练:每个客户端在局部进行训练,仅将模型更新上传至中央服务器(例如,使用 FedAvg [2] 进行聚合)。
大多数联邦学习(Federated Learning, FL)方法假设局部数据集为完全标注,但在临床场景中由于专家标注成本高昂,这一假设往往难以实现 [3]。由于未标注的医学图像数量远超标注数据,联邦半监督学习(Federated Semi-Supervised Learning, FSSL)逐渐成为更具实际意义的替代方案 [4]。现有的FSSL方法通过促进客户端间的预测一致性 [5]、利用全局与局部的知识蒸馏(knowledge distillation)[6],或采用无数据蒸馏(data-free distillation)与动态聚合策略来应对客户端数据分布的异构性 [7, 8]。最近,Ma等人 [9] 提出了HSSF,一种支持客户端特定架构的模型异构FSSL框架,通过正则性凝练与融合(Regularity Condensation and Fusion)实现知识迁移。
尽管在FSSL(联邦自监督学习)方面取得了最新进展,但在半监督语义分割中确保伪标签的可靠性仍然是一个核心挑战[10, 11],因为其质量高度依赖于各客户端局部模型的准确性与容量。
原则上,容量更高的模型能够生成更精确的分割预测,从而产生更可靠的伪标签。然而,在联邦医疗环境中,客户端设备通常受到严格的计算资源限制,难以训练或维持大型模型以进行局部学习与推理[12]。这种对强伪标签监督的需求与客户端模型能力受限之间的矛盾,进一步加剧了联邦设置下半监督学习的复杂性。为应对这一问题,作者提出了一种双知识蒸馏策略,包括:(I) 联邦知识蒸馏(federated knowledge distillation),促进全局服务器模型与局部客户端模型之间的双向知识迁移;以及 (II) SAM引导的知识蒸馏(SAM-guided knowledge distillation),利用Segment Anything Model(SAM)为轻量级局部模型提供细粒度的像素级监督。据作者所知,此前尚无研究在完全无标签客户端场景下针对分割任务实现服务器-客户端联合知识蒸馏。
本文的主要贡献总结如下:
• 作者提出了一种新颖的框架,利用服务器端丰富的计算资源来增强客户端无标签数据下的伪监督效果。
• 作者提出了一种基于服务器端训练的教师模型与客户端模型之间像素级预测一致性的半监督语义分割策略,从而实现可靠的伪标签生成。
• adas-SAM-Fed_2511支持通过 FedAvg 实现的同质聚合以及异质聚合,能够在计算资源各异的临床环境中灵活部署。
- adas-SAM-Fed_2511
问题定义。考虑一个联邦学习设置,包含一个中心服务器和
个客户端。服务器托管一个公开的带标签数据集
,其中包含
个样本对;而每个客户端
拥有一个私有的无标签数据集
,其大小为
,且各客户端之间的数据分布存在异构性(heterogeneous)。目标是通过一种半监督学习过程,协作训练局部分割模型,该过程利用公开共享的带标签数据,以及每个客户端局部私有的无标签数据。
Image
概述。所提出的异构客户端环境下框架的整体架构如图1所示。作者提出的框架包含三个主要组成部分:(1) 服务器端的监督训练,(2) 客户端侧的自监督学习(SSL)与基于自适应一致性的伪标签生成,以及(3) 联邦聚合。adas-SAM-Fed_2511具有模型无关性(model-agnostic),可集成到多种联邦半监督学习框架中。在此,作者以HSSF模型作为 Baseline ,并解决伪标签不可靠的问题,该问题通常会导致监督信号不稳定并降低分割性能。 Baseline 框架的进一步实现细节可参见文献[9]。
服务器端。在初始化阶段,服务器在公开的标注数据集
上训练一个监督全局模型,并在同一数据集上训练一个高容量的教师模型。尽管在adas-SAM-Fed_2511中任何高性能模型均可作为教师模型,但作者采用 SAM(Segment Anything Model),因其具备强大的泛化能力,记为
。作者在服务器端使用标注数据集对 SAM 进行微调,并采用低秩适应(Low-Rank Adaptation, LoRA)技术,随后分发得到的教师模型。
服务器端
输入:公开数据集
,SAM 教师模型
,全局模型
在
上微调
,在
上训练
将
和
的权重广播到所有客户端
客户端
输入:公开数据集
,局部 未标注 数据集
,客户端
模型训练局部模块
,其中
,
与
相关联。
对于
执行以下操作:
- 使用
和
生成伪标签
- 使用自适应一致性机制(Adaptive Agreement Mechanism)生成像素级一致 Mask ▷ 公式 (1)–(2)
- 为每个像素计算置信度权重
▷ 公式 (3)
- 将
加入伪标签数据集
,并附带对应的伪标签
和权重
在
上使用置信度加权的一致性损失训练
。
应用联邦聚合。
结合公开的标注数据集
,将其提供给所有客户端以实现伪监督。在每个客户端上,该教师模型通过一种自适应的、基于一致性的机制,将知识蒸馏到轻量级的半监督模型中。
客户端侧。每个客户端
使用自监督学习(SSL)结合基于自适应一致性的伪标签生成方法进行局部训练。局部模型
最初在公开的标注数据集
上进行训练。这确保了模型能够通过学习标注数据建立坚实的基础,从而在利用客户端私有的未标注数据之前做好充分准备。需要注意的是,SAM 仅通过一次前向推理过程生成初始伪标签,相较于训练过程,这一操作在计算上非常轻量,随后这些预测结果将在半监督学习过程中与客户端模型的输出进行动态集成。
自适应一致性机制:在训练初期,轻量级客户端模型对未标注数据的预测通常置信度较低,而教师模型则能提供更可靠的指导。随着训练的进行,客户端模型逐渐与局部数据分布对齐,其预测结果的可靠性也逐步超过冻结的教师模型。为应对这一变化,作者提出一种自适应伪标签策略,该策略在每次迭代中根据教师模型与客户端模型预测的相对置信度动态调整二者贡献权重。
在每次训练迭代中,客户端模型和冻结的教师模型分别对无标签批次中的每个样本
进行标签预测:
其中,
和
分别表示教师模型和客户端模型的 softmax 输出,
和
为对应的预测类别标签。
每个像素位置
的最终伪标签由一种共识感知的选择策略确定。当教师模型与客户端模型的预测一致时,作者接受它们的共识结果。
否则,作者选择置信度更高的预测:
其中第一种情况显式地捕捉了模型一致性,通过共识强化可靠预测。
为了加权每个伪标签的置信度,作者定义了一个自适应置信度权重
。设
和
分别表示教师模型和客户端模型在位置
处的最大置信度得分。则:
联邦聚合。作者针对同质和异质客户端设置采用不同的聚合策略。在同质情况下,所有客户端共享相同的模型架构,作者使用 FedAvg 在每轮迭代后计算局部参数的数据加权平均值,并广播更新后的全局模型。在异质情况下,作者遵循 HSSF Baseline [9] 的方法,采用规则性凝练(Regularity Condensation, RC)和规则性融合(Regularity Fusion, RF):服务器从客户端在公共数据集
上的预测中提取可靠知识(RC),随后将这些优化后的知识返回给客户端,客户端通过 KL 散度损失(KL-divergence loss)将其融入自身模型(RF),从而实现不同架构间的协作。
该加权方案在两个模型预测一致时赋予完全置信度,否则采用置信度更高的模型的最大置信度分数。
在生成带有自适应一致性 Mask 的伪标签数据集
后,局部模型将在 未标注 数据上进一步训练,使用生成的伪标签及其对应的权重。 未标注 数据的无监督损失计算为加权交叉熵损失:
其中
表示 未标注 样本的批次,
为交叉熵损失,
根据预测置信度调节每个像素的贡献权重。该阶段使模型能够利用大量 未标注 数据,同时优先关注可靠的伪标签区域。通过迭代的联邦更新,局部模型与全局模型均得以优化,伪标签质量在通信轮次中逐步提升。算法 1 提供了该过程的详细分步说明。
- 实验设置
数据集。作者在两个任务上评估了adas-SAM-Fed_2511:皮肤病变分割和息肉分割。对于皮肤病变任务,作者采用 ISIC2018 数据集(共 2,694 张图像),保留 100 张图像作为公开数据,并将剩余样本分配给四个客户端,分别包含 200、400、800 和 1,194 张图像(记为
–
)。对于息肉分割任务,作者采用 Non-IID 设置,使用五个数据集:CVC-ColonDB(公开,380 张图像)以及四个私有客户端数据集,分别为 CVC-ClinicDB(300 张图像)、EndoTect-ETIS(612 张)、CVC-300(396 张)和 Kvasir(1,000 张)。每个私有数据集均按 0.8:0.1:0.1 的比例划分为训练集、验证集和测试集。
网络与训练设置。作者在同质和异质两种设置下模拟四个客户端站点。在异质配置中,
和
使用 ResNet18 作为主干网络(Backbone),而
和
使用 ResNet34;服务器端采用 ResNet101 进行特征提取。在同质设置中,所有客户端均使用 ResNet34。所有图像均被缩放至
,并采用与文献 [9] 相同的数据增强和超参数设置。训练过程中使用 AdamW 优化器,初始学习率为 0.0001。在服务器端进行 SAM 微调时,为提高效率,视觉编码器(vision encoder)和 Prompt 编码器(prompt encoder)均被冻结。通过 LoRA 对 SAM 进行适配,设置秩
,缩放因子
,Dropout 为 0.1,最终得到 6.65M 可训练参数。
- 实验结果
Image
图2展示了异构设置下伪标签生成的过程,展示了SAM与客户端模型的预测结果及其一致性的 Mask ,其中绿色表示一致,红色表示不一致。每一行对应不同的训练轮次,突出了客户端模型的逐步优化过程。在早期阶段,客户端模型显著受益于SAM更优的预测结果,将其作为可靠的监督信号。随着训练的进行,客户端模型逐渐变得更加准确,并与SAM的预测结果日益趋同,最终在具有挑战性的区域也能生成高质量的预测。这种视觉上的演变凸显了自适应像素级一致性机制在整个训练过程中提供稳定且富有信息量监督的关键作用。结合前述的定量结果,这些观察结果证实,SAM-Fed能够有效利用SAM的泛化能力,同时使具有异构架构的客户端模型收敛至稳健且一致的分割性能。
- 结论
作者提出了 SAM-Fed,一种由 Segment Anything Model(SAM)引导的联邦半监督学习框架,旨在解决临床机构中计算资源受限的问题,这些机构通常难以在局部训练大型模型。
SAM-Fed 引入了一种双知识蒸馏策略,该策略结合了联邦知识蒸馏,实现全局模型与客户端模型之间的双向知识交互,以及 SAM 引导的知识蒸馏,利用 Segment Anything Model 为轻量级客户端提供细粒度的像素级监督。此外,一种自适应像素级一致性机制进一步动态优化伪标签,确保训练过程中监督信号的可靠性和一致性。在皮肤病变和息肉分割任务上的实验表明,SAM-Fed 在同质与异构配置下均持续优于现有的先进联邦半监督学习(FSSL) Baseline 方法,验证了其在真实世界联邦医学图像分割场景中的有效性与实用性。
参考
[1]. SAM-FED: SAM-GUIDED FEDERATED SEMI-SUPERVISED LEARNING FOR MEDICALIMAGE SEGMENTATION
点击上方卡片,关注「AI视界引擎」公众号
