相似度感知蒸馏的统一弱监督个性化联邦图像分割

2024-03-20 10:31潘建珊林立吴洁伟刘翼翔陈孝华林其友黄建业唐晓颖
中国图象图形学报 2024年3期
关键词:联邦站点像素

潘建珊,林立,吴洁伟,刘翼翔,陈孝华,林其友,黄建业,唐晓颖*

1.深圳市公共信用中心,深圳 518000;2.南方科技大学电子与电气工程系,深圳 518000;3.香港大学电机电子工程系,香港 999077;4.南方科技大学嘉兴研究院,嘉兴 314000

0 引言

近年来,机器学习尤其是深度学习模型在多个计算机视觉和医学影像分析任务中取得了先进性能并因此在许多相关领域得到广泛应用(周传鑫 等,2021)。然而作为一种数据驱动的方法,深度学习模型的性能在很大程度上取决于可用数据的数量及标注的质量。受限于患者数量和标注成本,单个医疗机构往往面临着数据量不足的问题,因而联合多家医院的数据进行协作训练对于获取一个性能强大的智能分析模型意义重大,特别是对于医学影像分割任务而言。然而,随着隐私和安全意识的增强,许多政府和组织通过不同的法规或指南规定了对于个人或患者数据及影像的保护,限制了数据的共享(Voigt 和von dem Bussche,2017)。在这种情况下,联邦学习(federated learning,FL)应运而生并受到研究者的广泛关注,因其允许不同中心在无需共享或集中数据的情况下联合训练性能强大的全局模型(Yang 等,2019)。在常见的联邦学习范式中,每个本地客户端(如医疗站点)使用本地的训练数据训练模型,并将各个本地/局部模型参数汇聚到中心服务器进行某种方式的聚合并向各个站点广播更新后的参数来实现协调。举例说明,最常用的平均式联邦学习(federated averaging,FedAvg)依据各站点样本数量对各局部模型参数进行加权平均,在服务器端获得单个全局模型(McMahan等,2017)。

尽管联邦学习已经在医学图像分割领域取得了一定的进展,但现有工作仍大多属于传统的中心化联邦范式,也即通过图像风格迁移(Liu等,2021)、图像模拟(Li 等,2020a)或调整聚合权重(Li 等,2019)等手段来提升单个全局模型在每个站点的鲁棒性和泛化性。然而由于各站点之间数据存在着不可避免的分布差异,使用单一全局模型以适应联邦中的所有站点是极具挑战的,尤其是在医学影像场景中,如图1(a)所示。医生的拍摄经验、成像设备的型号、成像参数及患者的群体差异等都是导致客户端图像存在领域偏移的重要因素。

图1 本文使用的数据及标注的样例Fig.1 Representative examples of data and labels utilized((a)the two datasets constructed and used in this paper,with domain gaps across sites;(b)examples of original images with various sparse labels and full masks)

在这种背景下,研究者提出个性化联邦学习,它为提升联邦中每个客户端模型的性能提供了一个有潜力的解决思路(Sun等,2023)。目前已经提出的个性化联邦方法主要分为两大类,第1类延续了中心化联邦获取中心模型的思路,并在各个站点进行局部微调以获取更适合自身分布的个性化模型,如Wang等人(2019)提出的FT(FedAVG with fine-tuning),这类方法易受限于深度学习模型的固有弱点:灾难性遗忘,即遗忘在第1阶段学到的公共知识而在局部数据上过拟合。另一类更先进的解决方案通过将模型划分为全局共享部分和个性化部分,将模型的共享部分汇聚到中心服务器进行聚合,个性化部分则保留本地训练所得参数,从而实现模型性能的提升,例如,Li 等人(2021)将批标准化(batch normalization,BN)层作为模型个性化部分,而Collins等人(2021)则将个性化层设置在模型的预测头部分。这些方法仅关注各站点自身的信息而忽略了站点间的相似性,不可避免地丢失了其他站点的有用知识,从而弱化了模型的表征能力。另一方面,上述方法仍需要额外的中央服务器以聚合共享部分的参数,探索一种无需中央服务器的联邦方法对降低计算消耗、提升方法实用性起到重要意义。此外,由于个性化联邦学习发展时间尚短,多数方法仅在简单的分类任务上进行设计及验证,而在相对复杂的分割任务上鲜有研究。

深度学习方法在医学影像分割任务上受到极大关注并取得重大成功,特别是在Ronneberger 等人(2015)提出U-Net 之后。此后大多数的分割方法主要关注于提出精心设计的更为强大的网络架构(Lin等,2021),或者各种融合医学、拓扑、体积等先验的损失函数(Lin等,2020)以提升分割性能。然而数据可用性和标注质量是全监督分割范式的主要限制,这对于需要专家知识和临床经验进行标注的医学影像而言尤甚。弱监督学习提供了一种数据高效的解决方案(田萱 等,2019),通过使用稀疏粒度(如点、涂鸦、边界框、多边形块等)的标注进行模型的监督,已经有部分工作分别基于上述各种稀疏标签的弱监督训练方法进行研究,并取得了良好的性能(Obukhov 等,2019;Liang 等,2022)。在联邦中,一个更符合实际需求的设定是允许不同的中心站点使用异构形式的稀疏标签,如图1(b)所示,并以统一、兼容的方式进行训练,使各站点都能从联邦中获得性能收益。据调研,目前针对弱监督下的个性化联邦方法仍未有相关研究报道。

针对上述现状,本文立足于医学影像分割这一任务,提出了一种统一的弱监督个性化联邦学习范式,称为pFedWSD(unified weakly suppervised personalized federated image segmentation via similarityaware distillation)。所提出的pFedWSD 通过循环知识蒸馏为每个中心训练个性化模型,主要包含两个阶段:不确定度感知的动态循环公共知识积累阶段和分布相似度感知的个性化阶段。在弱监督方面,采用了部分交叉熵、门控条件随机场和树能量损失函数的综合优化目标,有效统一了各种不同稀疏标注的数据来进行训练,可以在线得到边缘更准确的伪标签,从而获得优异的分割性能,而无需额外的监督数据、交替的优化过程或耗时的后处理。

本文的主要贡献归纳如下:1)据调研,这是第1个在各客户端采用异构标签设定下的弱监督医学图像分割个性化联邦学习方法;2)所提出的两阶段个性化联邦,通过不确定度感知的方式动态地排序每一轮训练中各客户端模型性能,并以循环知识蒸馏的形式积累公共知识,在第2 阶段通过批标准化层的统计信息度量各站点之间的相似性并聚合得到各站点教师模型后进行知识蒸馏,可以有效地在无需中心站点的条件下进行所有中心的公共知识积累及自适应个性化,并且有效利用其他站点模型的有效信息;3)经过在眼底视杯(optic cup,OC)、视盘(optic disc,OD)和视网膜中心凹无血管区(foveal avascular zone,FAZ)分割任务上的定量和定性实验,所提方法相比其他联邦学习方法取得了更好的性能,所有中心都能从联邦中获益,并且取得了接近于使用全监督标签进行集中式训练的性能表现。

1 本文方法

1.1 问题表述

式中,L为损失函数。

1.2 个性化联邦流程

受Chen 等人(2023)所提出的MetaFed(federated learning among federations with cyclic knowledge distillation for personalized healthcare)的启发,pFed-WSD 同样采用循环知识蒸馏的基本框架,而无需额外的中心服务器,在不损害数据隐私和安全的前提下实现多中心通用知识的积累并保留个性化信息,其主要框架展示在图2(a)中。不失一般性,本文假设有5 个中心参与联邦学习,整个训练过程主要包含两个阶段,不确定度感知的动态知识积累阶段(黄色箭头表示)和分布相似度感知的知识蒸馏个性化阶段(蓝色箭头表示)。

图2 本文弱监督个性化联邦学习框架图Fig.2 Overview of the proposed pFedWSD pipeline((a)the personalized FL framework;(b)the weakly-supervised learning approach)

在通用知识积累阶段,所提出框架采用循环知识蒸馏的训练范式以完成模型对于公共知识的积累,即按一定的顺序将多个中心排序并将上一个中心作为下一个中心的教师模型。不同于MetaFed 的固定顺序蒸馏,为了更有效地完成知识的蒸馏并加速训练,各中心的模型应按照当前训练轮次的性能进行排序,即性能最高的模型作为性能次优模型的教师,以此类推,并在每个联邦轮次中进行动态调整,以达到各中心均衡提升的目的。pFedWSD 采用在各终端验证集上的分割性能(Dice系数)及模型不确定度进行模型性能排序,该综合评价指标既考虑了模型的绝对分割精度,又将鲁棒性/方差考虑在内,可以更全面地排序当前轮次各站点的模型性能。具体为

式中,Dci和Ui分别表示第i个站点当前训练轮次的Dice 系数和不确定度,λu为权衡参数。受贝叶斯网络中不确定度估计的启发,本文方法使用蒙特卡洛丢弃法来估计不确定度(Kendall 和Gal,2017)。具体而言,在不同的加性高斯噪声和随机丢弃下,每幅验证图像通过模型随机地向前传递T次。对于输入样本的每个像素,可以获取一组softmax 概率向量,选择预测的熵作为度量来近似不确定度,从而可以表述为

式中,Nva为验证样本的个数,H和W分别表示图像的高和宽。经过观察发现,不同站点模型的性能与样本的平均不确定度之间存在一定的负向关联。为了直观呈现这一观察结果,图3 展示了不同模型对同一样本的预测结果,并提供了相应的不确定度图样例。

图3 不同模型对同一样本的预测结果及不确定度图样例Fig.3 Examples of segmentation predictions and uncertainty maps from different models for the same samples((a)original images;(b)predictions;(c)uncertainty maps;(d)ground truth)

在确定循环蒸馏排序之后,本文框架使用KL(Kullback-Leibler)散度损失来拉近教师与学生模型预测级别之间的距离,实现公共知识的传递。具体为

式中,ϕi为当前站点(学生)模型,而ϕi-1为上一站点(教师)模型,x为当前站点的样本。因此,训练每个站点模型的总目标函数为

式中,λd为权衡知识蒸馏损失和当前数据监督损失之间的权重系数,Lseg为弱监督分割损失,其具体形式将在1.3节展开。需要注意的是,在第1阶段的每一轮训练中,对每个站点模型的验证性能与教师模型的验证性能进行对比,如教师模型在验证集上的Dice 系数高于学生模型的Dice 系数,则模型使用式(7)进行训练,否则仅利用本地数据进行训练,即式(7)中λd设置为0。在这个阶段中,为了保持各站点模型对其数据分布的个性化,本文方法沿用了FedBN(Li 等,2021)保留每个客户端模型批标准化(BN)层的操作。经过一定的训练轮数,可以获得一个带有各站点公共知识的通用模型ϕ。

在第2 阶段,即分布相似度感知的个性化阶段,首先将ϕ分发给各站点,并结合各站点BN层初始化各站点模型。第2 阶段仍然采用循环蒸馏的方式进行本地模型的迭代训练。为了更好地实现个性化,并从相似站点模型学习获取有用知识,第2 阶段各站点的教师模型为每个训练轮次中各站点模型的相似度加权聚合。具体而言,每个客户端基于本地数据计算BN层的统计量并广播,而后各终端计算得到相似度权重矩阵M,以指导每轮训练中各客户端教师模型的加权聚合。在每轮训练中,各个客户端将模型进行广播,并使用M聚合知识蒸馏教师模型。在实际设置中,为了稳定和提升计算效率,仅使用第1 阶段中获得的各站点BN 层参数计算M,并在第2阶段的训练中固定M。M为一个N×N矩阵,矩阵元素mij∈[0,1]表示站点i与站点j之间的相似度,其值越大,相似度越高。需要注意的是,各站点教师模型的BN 层使用本地上一轮更新模型的BN层参数,而其他层参数为

式中,L代表网络的层数。受Lu等人(2022)的启发,使用推土机距离(Wasserstein distance,WD)的近似值来度量两个高斯分布之间的距离,具体为

假设每个通道都是相互独立的,则σj,l为对角矩阵,σj,l=diag(ri,l)。因此,两个客户端i,j之间的距离可以表示为

di,j值越大代表分布差异越大,则相似度权重mij应越小。因此将设置为di,j的倒数,即=1/di,j,j≠i,将进行归一化,可得

使用滑动平均方式更新ψt+1,并将mi,i设置为常数α,则

由此,可得到各个终端的教师模型的相似度聚合权重。在第2 阶段,每个站点模型延续式(7)进行训练,教师模型将在当前中心的验证数据上得到的Dice 系数作为性能指标,当其值低于本地模型的性能指标,则λd设置为0,反之,教师模型性能越高,则λd越大,式中λ0为超参数,Dcvate和Dcvast分别表示教师模型和学生模型在验证集上的Dice系数。经过若干轮次训练,直至各客户端模型收敛。

1.3 弱监督训练范式

图像的语义分割本质为像素级别的分类任务。最经典的弱监督分割方法将传统交叉熵损失函数修改为部分交叉熵函数(partial cross entropy,pCE)(Tang 等,2018),也即仅在标注区域的像素点进行交叉熵损失计算。然而仅仅在稀疏标注区域上应用pCE 通常无法在靠近语义边界的地方提供足够的监督信号,因而在没有监督的地方会导致差的预测结果。一种较为可行的思路是通过特定方式从已标注的监督区域向外扩展,形成未标注区域的伪标注,从而为模型训练提供更充足的监督信号(Obukhov 等,2019;任冬伟 等,2022)。

属于同一物体或对象的像素在不同的特征语义空间上应共享相似的模式,受树滤波器的启发(Song等,2019;Liang 等,2022),本文根据其结构保持的特性来建立这种成对相似性,而成对相似性与模型预测则用于为未标记的像素生成软伪标签,通过不断在线训练与修正,实现模型预测与伪标注的逐渐改进。图2(b)中展示了pFedWSD 中采用的弱监督训练范式,主要由一个对标注像素的分割分支和一个对未标注像素的辅助分支构成。分割分支将稀疏标注Y作为监督信号并沿用pCE 损失进行监督,辅助分支则从原始图像I与分割模型中选定层的特征F得到一对亲和性矩阵Alo和Ahi,用于精细化网络预测P并生成软伪标注,生成的伪标注用于监督未标注的像素区域。此外,为了进一步修正模型对于目标边缘的预测精度,本框架采用了Obukhov 等人(2019)提出的门控条件随机场损失(gated conditional random field loss,LgCRF)作为正则化项,用于抑制单纯使用树能量损失Ltree可能会出现的目标区域伪标注过度膨胀或坍缩的情况,总体目标函数可以表示为

式中,λt和λg为平衡各弱监督分割损失函数的权重参数。

以下给出树能量损失Ltree的具体计算过程,首先是树亲和度计算。一幅图像可以用无向图G=(V,E)来表示,V表示所有像素形成的顶点集,而相邻点之间的边组成边集E,在本文中采用四方位联通建图,即仅将每个顶点的上下左右相邻像素连接成边。给定相邻像素i和j,它们之间的低阶权重和高阶权重分别定义为

式中,I(i) ∈R3×h×w和F(i) ∈R256×h×w分别为原始图像像素值和特征图值,h和w代表图像的高和宽。F(i)由分割网络选定层的特征经过1 × 1 卷积改变通道数得到。得到权重之后,在保证图连通性前提下依次从E中删除权重最大的边以构建最小生成树(minimum spanning tree,MST)。使 用Borůvka(Gallager 等,1983)算法生成高阶和低阶MST,与Song 等人(2019)方法相似,MST 的两个顶点之间的距离可以通过它们相连边的权重求和来计算。顶点之间最短路径的距离,记为超边S,构成了MST 的距离图,具体为

式中,i,j,k和m为顶点索引,* ∈{lo,hi}。为了捕捉顶点间的长程关系,将距离图投射到正亲和度矩阵,即

式中,σ为用于调节强度信息的超参数。给定一幅图像,低阶亲和度矩阵是静态的而高阶亲和度矩阵是动态的,两个矩阵提取特征层面的成对关系,可以训练网络获取互补知识。由于低阶亲和矩阵Alo包含边界信息,而高阶亲和矩阵Ahi保持语义一致性,通过级联滤波器模块,将模型预测图精细化后的结果作为软伪标签

式中,P为softmax层输出的概率值图。通过与低阶、高阶亲和度矩阵相乘并经过滤波器F,其运算式为

式中,Ω为所有像素集合,而zi=为归一化项。

获得伪标注之后,树能量损失可以表示为

式中,δ为标签分配函数,用于测量预测概率P和伪标签之间的距离,此处本文框架使用L1 范数损失。因此,最终树能量损失可以表示为

式中,ΩU代表图像未标注区域像素,带标注区域通过监督信号由部分交叉熵损失进行优化,而无标注区域由树能量损失进行优化,从而使每个站点模型在弱监督设定下实现在线自训练。

此外,值得注意的是,点、涂鸦和块标注均可以作为像素级别的稀疏类别监督信号,然而边界框作为描述目标边界的矩形框,需要经过一定的预处理转化才可作为类别监督信号。根据所分割目标的形状、尺寸等先验,可以将其转换为点、涂鸦、块或者三者的混合。以本文使用数据集为例,如图4 所示,由于已知OD、OC 为两种类椭圆结构,因而首先计算获得两种目标区域边界框的最大内接椭圆,而后对内椭圆中间挖去一个小圆区域并对两椭圆环形区域应用骨架化操作,从而获得两个目标区域的骨架/涂鸦稀疏标注,同时对OD 边界框外的背景部分进行骨架化操作得到背景类的稀疏标注,从而将边界框标注转换为类涂鸦标注;对于FAZ,由于其形状较不规则,将边界框四边往外扩张若干像素(本文具体设置为10 像素),扩张后的矩形框的外部像素设置为背景类,同时,将目标区域边界框的长与宽分别缩小为原本的1/3得到缩小框,将缩小框内的像素设置为目标类,由此可将边界框转换为块标注。对于其他复杂结构,同样可以根据形状、尺寸等先验,进行一定的预处理将框标注转换为稀疏像素标注。

图4 对采用边界框标注的数据的标签预处理流程Fig.4 Label preprocessing process for data annotated with bounding boxes

2 实验结果与讨论

2.1 实验数据集及预处理

依据实验和验证需要,本文利用多个公开医学影像分割数据集,使用自动化方法生成其对应稀疏标注(同一站点使用同一种标注形式,包含点、涂鸦、边界框和块4 种形式),构建两个分别用于眼底OD/OC 分割及视网膜FAZ 分割的多领域分布弱监督分割数据集。其中,眼底OD/OC 分割数据集分别使用Drishti-GS(retinal image dataset for optic nerve head segmentatio)(Sivaswamy 等,2014)、RIM-ONE-r(open retinal image database for optic nerve evaluation)(Fumero 等,2011)、REFUGE-train(retinal fundus glaucoma challenge)、REFUGE-val(Orlando 等,2020)及Gamma(Wu 等,2023)数据集充当站点A—站点E的数据,各站点标注形式分别为涂鸦、涂鸦(另一种风格)、边界框、点、块,以下统称为Fed-ODOC 数据集;对于FAZ 分割任务,本文使用了FAZID(foveal avascular zone image database)(Agarwal 等,2020)、OCTA500-3M、OCTA500-6M(Li 等,2020b)、OCTA-25K-IQA-SEG(Wang 等,2021)和ROSE(retinal OCTangiography vessel segmentation dataset)(Ma 等,2021)眼底相干光层析血管图像(optical coherence tomography angiography,OCTA)来充当各站点数据,其标注形式分别为涂鸦、点、块、边界框、涂鸦(另一种风格),以下称为Fed-FAZ 数据集。各站点数据集均遵循原有数据集的训练集与测试集划分,并在训练集中按一定比例进行训练集和验证集的划分。各站点图像样例见图1(a),图中同时列出了训练及测试样本数。对于Fed-ODOC 数据集,依据原始眼底照的掩码标注截取视杯视盘区域并统一调整大小为384 × 384 像素;而Fed-FAZ 数据集则维持原图像范围并统一调整大小为256 × 256 像素。数据预处理包括将所有图像像素值归一化到0至1之间(数据精度使用float32),数据增强包括随机水平与垂直翻转图像,随机旋转一定角度(范围为-45°~45°)。

接下来说明各稀疏标注形式的自动化生成方法。给定一个全监督掩膜标注,如图1(b)最后一列所示,可以通过形态学腐蚀变换对每一类区域进行处理得到块标注,通过scikit-image 库中的骨架化操作得到涂鸦标注,通过对取得的骨架做局部形变变换并随机擦除得到另一种风格的涂鸦标注,通过计算目标区域掩膜的外接矩形得到边界框标注,通过计算目标区域的内接矩形四边中点作为目标类点标注,通过扩展边界框并取四边中点作为背景类点标注。

2.2 实验环境与参数设置

实验环境为一台装配8 块NVIDIA GeForce RTX 3090 显卡,系统为CentOS 7,CPU 型号为Intel Xeon Gold 6242R 的服务器。编程语言为Python,深度学习框架为PyTorch,本文框架在联邦学习框架Flower 的基础上进行开发。在每个站点使用原版U-Net作为分割模型架构,网络的通道数从上到下分别为16、32、64、128 和256,并且在树能量损失的计算中选用解码器第2 层特征进行级联滤波。模型的初始化方式为随机初始化,优化器选用默认参数下的AdamW 优化器,初始学习率为10-2,并随着迭代步数进行更新。

式中,lr表示当前学习率,lr0表示初始学习率,e为当前迭代步数,Ne为总迭代步数,其数值为3万。在超参数选择方面,根据内部探究性实验经验将λu(式(2)),λd(式(7)),α(式(13)),λ0(式(14))均设置为0.5,将式(15)中的λt和λg均设置为0.1。所设计的联邦方法中第1 阶段的训练轮数为50,第2 阶段轮数为1 000。所有对比方法均在相同的实验环境中部署并采用相同的学习率、优化器、优化策略及模型架构,以达到尽可能公平的对比。

2.3 评判指标

本文着眼于医学影像分割任务,因而选取了Dice 相似性系数(Dice similarity coefficient,DSC)与95%豪斯多夫距离(Hausdorff distance,HD95),以及精确率(precision,Pre)、召回率(recall,Rec)作为主要的衡量指标,其中Dice 系数表示预测结果与真实标签图的重叠程度,其值在[0,1]之间。而HD 表明了预测结果与真实标签图的最大不匹配程度。Dice的数学表达式为

HD的数学表达式为

式中,P与Y分别表示预测结果与真实标签,p、y为预测结果与真实标签上的像素点,HD95 与HD 类似,是基于计算P与Y中边界点之间距离的第95 百分位数,以此消除小异常值子集的影响。较大的Dice 系数和较小的 HD 95 代表更好的分割结果。

2.4 与其他联邦学习算法对比

将本文方法与几个代表性联邦学习框架进行性能比较,包括常用的传统中心化联邦方法(如FedAvg 与FedProx)以及较先进的个性化联邦方法(如FT、FedBN、FedAP(federated learning with adaptive batchnorm for personalized healthcare)、FedRep、FedALA(adaptive local aggregation for personalized federated learning)等)。在实现方面,由于这些方法最初是为图像分类任务设计,本文实验尽量保持它们原本的设计原则并使其适应弱监督图像分割任务。FedAvg 依据各站点样本数量对各局部模型参数进行加权平均,在服务器端获得单个全局模型(McMahan 等,2017);FedProx 通过在每次迭代中强制执行模型参数之间的平衡,以确保各个站点训练的模型具有相似的参数值,从而减少数据偏差,是FedAvg 的泛化形式(Li 等,2020c)。个性化联邦方法中,FT 为带有 站点微调的FedAvg;FedBN 和FedRep 分别将所有的BN 层和分割头(最后一个卷积层)作为模型的个性化部分;FedAP 通过BN 层统计各站点相似度并聚合,同时保留各站点BN 层的个性化参数;MetaFed 基于静态(固定顺序)环形知识蒸馏实现各站点的个性化;FedALA(adaptive local aggregation for personalized federated learning)(Zhang 等,2023)通过可学习的自适应局部聚合(adaptive local aggregation,ALA)模块,针对每个客户端上的局部目标自适应地聚合全局模型和局部模型,在每次迭代中训练之前初始化局部模型。在对比联邦方法中,各站点均使用原始稀疏标注进行直接监督(边界框标注站点与本文方法同样经过预处理为可直接监督的稀疏标注形式),而未针对弱监督进行额外的范式设计。同时本文还与基线设定(弱监督下的本地训练)进行比较,在该设定下各站点使用各自的数据及对应的稀疏标注训练本地分割模型。类似地,本文进行了全监督下的本地训练、弱监督下的集中训练及全监督下的集中训练,其中,集中训练指不考虑数据的隐私性及安全性,将各站点数据集中于一台中心服务器上合并,而后进行直接训练。

表1 和表2 显示了Fed-ODOC 数据集的量化结果。首先,通过表1和表2的最后一列总计指标可以看到,相对于本地训练而言,所有联邦学习方法都可以有效提升各站点模型在各自测试数据中的整体性能。其中站点D 通过加入联邦,性能获得了显著提升,其主要原因是站点D 使用了点标注这一种监督信号最弱的标注形式,加入联邦使其可以从其他站点获取有效知识。站点F 同样在加入联邦之后获得了较大的性能提升,其原因在于该站点数据/患者的分布差异较大,仅通过单一站点的数据难以训练一个强大、泛化性能强的深度分割模型。从方法类别上来看,个性化联邦方法由于可以个性化模型以应对各个站点分布各异的数据,因而性能总体优于传统的中心化联邦框架。值得注意的是,本文个性化联邦框架取得了最好的性能,每个站点均能在本文联邦框架中获益,大多数站点取得了所有联邦方法中最优或者次优的性能,并且是唯一在视杯视盘分割的综合性能中取得Dice 系数高于90%的方法,显著优于集中式训练(弱监督)的性能,并取得与本地训练(全监督)、集中式训练(全监督)最接近的性能。这意味着,每个站点在保护隐私的前提下,通过多站点的数据合作可以训练得到一个性能强大的模型,并且各站点可以采用不同的稀疏标注方式,极大程度上节约了标注成本。

表1 不同联邦学习算法及不同本地训练、集中训练设定下的算法在视杯视盘分割任务上的Dice系数对比Table 1 Dice scores of OD/OC segmentation from different federated learning approaches,as well as under various localized and centralized training conditions /%

在Fed-FAZ 上的结果进一步验证了所提方法带来的提升,如表3所示,Fed-FAZ数据更具挑战性,存在部分站点数据较少的问题(如站点E)。此外,由于各站点存在不同的图像质量问题和分布差异,导致在集中训练以及部分联邦学习方法中只有小幅度提升甚至存在性能下降,各站点间存在负向干扰。相较于本地训练,并非所有的联邦方法都能为各站点带来性能提升,例如FedBN、FedAP 与MetaFed。可以看出,简单地将部分参数个性化并不能为各站点模型带来显著的正向效果,模型的共享部分仍容易受到其他站点由于分布差异、质量差异等因素带来的负面干扰,同时静态的环形知识蒸馏也容易受到不相似站点模型充当教师模型所带来的无效知识影响。本文方法由于经过公共知识积累阶段以及相似度聚合知识蒸馏阶段,可以凝练积累各站点数据的相似且有效的知识,并在第2 阶段为相似站点赋予更高权重,从而蒸馏得到有益知识,为差异站点赋予小权重从而减轻干扰。同样地,本文方法可以提供与本地训练(全监督)及集中式训练(全监督)相竞争的性能。

表3 不同联邦学习算法及不同本地训练、集中训练设定下的算法在中心凹无血管区分割任务上的性能对比Table 3 Performance comparison of different federated learning approaches,as well as under various localized and centralized training conditions on FAZ segmentation

图5 展示了本文方法与其他联邦方法的分割结果可视化对比。其中,CT 代表汇总集中数据训练,weak 指数据使用稀疏标签,full 指数据使用全监督标签,图中数值表示Dice 系数(%)。得益于所提出的个性化方法及更优的弱监督损失,本文方法在目标区域的边缘精确度、拓扑结构保持以及离群值抑制等方面都取得了更好的性能表现。

图5 分割结果对比Fig.5 Visualization of representative segmentation results from pFedWSD and other federated learning approach((a)original images;(b)FedAvg;(c)FT;(d)FedProx;(e)FedBN;(f)FedAP;(g)FedRep;(h)ours;(i)CT(weak);(j)CT(full);(k)ground truth)

2.5 消融实验

为了证明所提出的弱监督损失以及联邦范式的有效性,本文在Fed-ODOC 数据集上进行了5项关键组件的消融实验,分别是弱监督损失中的树能量损失及门控条件随机场损失,即分别对两项弱监督损失、pFedWSD 中的两个阶段、pFedWSD 第2 阶段中的相似度感知权重(替换为FedAvg 中的样本权重系数)进行消融,结果如表4 所示。可以看出,任何一个损失项的消除均导致一定程度的分割性能下降;导致模型在预测中边缘精确度下降或离散预测区域增加。随后进行pFedWSD 框架的两个阶段——不确定度感知的动态循环公共知识积累阶段(无公共知识积累阶段)和分布相似度感知个性化阶段(无个性化阶段)的消融。可以看出,随着公共知识积累阶段的去除,各站点综合性能大幅下降,表明对各站点知识进行动态循环蒸馏可以使得各站点模型保留各中心训练数据所得到的有效知识,并摒弃冗余知识,对最终模型性能有着重要意义。个性化阶段的消除造成了较小程度的性能下降,表明即使仅使用公共知识积累阶段的模型,也可以取得较好的分割结果。即便如此,个性化阶段仍然为各站点模型的综合性能在Dice方面提升了约2.2%。此外,将个性化阶段的教师模型采用的相似度感知聚合权重消除,替换为FedAvg 中基于样本占比的权重,各站点模型无法高效地从教师模型中学习相似站点的有效知识,并且容易受到不相似站点带来的负面影响,因此性能存在较大下降。

本文还通过对各站点依次消融来分析在pFed-WSD 框架下各站点数据对整个联邦的贡献程度。同时,随着各站点的依次消融,可以分析各站点之间数据的相似程度及相互之间的影响。在pFedWSD框架下,无站点A、无站点B、无站点C、无站点D、无站点E与全站点设置下的各站点综合平均Dice系数依次为88.17%、88.63%、88.19%、87.98%、89.14%和90.38%,豪斯多夫距离依次为13.18、12.59、11.562、12.42、11.74和9.63。从图6的详细结果展示中可以看出,各站点均在全站点设定中取得最优性能,表明了尽管各中心存在数据分布差异,但所有站点均能从尽可能多的联邦站点中学习到更好的个性化模型。对于站点A 而言,站点D 退出时取得了最低的性能,表明两中心间数据相似度较高,能有效相互促进性能;同样地,对于站点D,站点A 的退出导致了最大的性能下降。类似地,可以看出站点E与其他站点数据分布差距较大,随着站点E的退出,其余多数站点取得了在4 个站点联邦条件下最优或次优的性能。实际中,站点E可能要提供更多的样本或更高质量的标注形式以提升该站点对联邦的贡献程度。与之相反,站点D的退出导致其余站点的综合性能较大跌幅,这可能是因为站点D贡献着所有中心中最大的样本量,其他的中心应允许其使用更为稀疏(宽松)、高效的标注方式。

3 结论

针对医学影像全监督分割标签标注成本高、联邦学习中可能出现的数据分布差异以及标签异构等问题,本文提出了一种弱监督个性化联邦分割方法,通过循环知识蒸馏为每个中心训练个性化模型,主要包含两个阶段:不确定度感知的动态循环公共知识积累阶段和分布相似度感知的个性化阶段。

通过在视网膜FAZ 及眼底OD 和OC 两个分割任务上的大量定性和定量对比实验,所提出的pFed-WSD 在与其他中心式联邦学习和先进的个性化联邦学习方法的对比中取得了更优的性能,可以提取各站点的有效公共信息并进行自适应个性化,在各站点均是不同形式弱监督稀疏标签的设定下,取得了接近于使用全监督标签进行本地训练甚至集中训练的性能。但由于在两个阶段中均需要进行多个教师模型的计算与通讯,本文方法存在通讯和计算开销上的短板。如何降低上述开销以及引入例如提示学习等通用大模型训练范式,将数据分布与标注稀疏程度作为模型的额外条件输入,同时促进模型的泛化性能和针对具体数据分布的自适应个性化性能是今后的主要研究方向。

猜你喜欢
联邦站点像素
赵运哲作品
像素前线之“幻影”2000
一“炮”而红 音联邦SVSound 2000 Pro品鉴会完满举行
“像素”仙人掌
303A深圳市音联邦电气有限公司
基于Web站点的SQL注入分析与防范
2017~2018年冬季西北地区某站点流感流行特征分析
首届欧洲自行车共享站点协商会召开
怕被人认出
高像素不是全部