上海大学马丽艳团队Patterns：部分标注多器官分割联邦学习框架

2024-02-05 16:01

上海

原创 Cell Press CellPress细胞科学

交叉学科

Interdisciplinary

部分类别监督技术能够有效减少医生的标注代价。随着深度学习的发展，数据隐私的问题得到了公众的普遍关注，人们对于训练模型的数据隐私提出了更高的要求，联邦学习可以在原始数据不离开本地的基础上，训练通用模型。对于部分类别标注的医学图像分割任务，中心化训练时通常合并未标记类别为背景，让模型自适应地调整前景与背景的关系，但是这种做法在没有前景类别的交集时会产生严重的类别混淆问题。以往的研究通常不考虑数据集间的特征差异，而该问题在医学图像中普遍存在，在联邦聚合时会导致模型权重的冲突。

近日，上海大学计算机工程与科学学院马丽艳团队在Cell Press细胞出版社旗下期刊Patterns上发表了题为“UFPS: A unified framework for partially annotated federated segmentation in heterogeneous data distribution”的新研究。UFPS针对无类别交集的CT数据，设计了统一类别与特征空间的训练框架，促进了联邦全局模型对任意客户数据集的分割效果。该框架与模型结构无关，可适配于通用医疗大模型的学习。

研究亮点

提出了联邦学习下部分类别标注挑战的解决方案；

基于多个无类别交集的特征异质数据集设计了全局模型训练策略；

呈现了全局模型在多类多域的有效性。

研究简介

近年来，部分类别监督技术的相关研究工作促进了通用医疗模型的发展，这些模型仅需以少量类别标注的多个数据集进行训练，大幅减少了医学领域的标注成本；然而，该技术主要依赖于中心化的数据集群，不符合现实医疗应用中的隐私法规。作为一种分布式学习框架，联邦学习通过聚合模型权重或梯度，使得多个协作方共同训练一个全局模型而不共享原始数据成为了可能。

目前，学术界还未充分探索基于多个部分标注数据集的联邦分割任务。不同于数据中心化的训练，将部分类别监督方法直接运用到无类别交集的联邦分割场景中，会导致全局模型的收敛效果不佳，如直接将无标注类别合并到背景中会产生知识混淆，不对背景类别进行监督会导致深度网络的任意优化（图1）。一种可行方案是用这些无类别交集但所处同一区域的图像训练统一的特征提取器，同时对各个数据集的特定任务单独训练分割器；该策略能够获得有效的训练结果，但是相比于通用分割模型，其推理时间成倍增加。此外，由于医学领域中的成像仪器、图像采集协议、人口密度等差异，会导致图像特征不同，进而提升模型训练的难度，但是已有的部分类别监督方法没有考虑到该问题。

图1 部分类别监督方案在联邦分割任务中的问题说明

在该研究中，马丽艳团队针对联邦学习场景下部分类别标注且特征异质的图像分割任务，提出了统一训练框架（UFPS）。在该框架中，分别设计了统一标签学习（ULL）和稀疏统一锐度感知最小化（sUSAM）策略，同时处理类别与特征的双重异质性（图2）。

图2 统一训练框架的完整流程图

在ULL中，为了以无类别冲突的方式训练联邦通用模型，通过本地预训练模型来提供背景类去重的伪标签，实现了部分类别标注到噪声标签学习的问题转换。在联邦训练的过程，干扰模型性能的噪声来源主要分为三类。第一类源于给出伪标签的教师模型，考虑到全局模型可以集成多域的知识，泛化能力强于预训练教师模型，在其训练达到稳定状态后可以作为主教师模型；然而联邦聚合的过程会弱化对图像边缘的区分能力，因此采用基于预测共识的方法使得伪标签更可靠（图3）。二是联邦聚合时的噪声，预训练教师模型为高质量数据集给出的伪标签置信度更高，因此调整了聚合权重与数据质量的关系，让以优质伪标签训练的模型主导对全局模型的影响。最后一类是本地训练时的噪声，通常和难以学习的类别混淆，可以解耦类别学习的过程让难类训练更稳定；同时考虑到本地模型的可靠程度逐步提高，逐步提升了其预测相比于伪标签对模型的影响力。

图3 基于预测共识的伪标签生成示意图

进一步的，sUSAM提供了优化方向解耦的CT图像通用学习方法，以此在保证对本地数据集性能的前提下，对于未知域也具有较强的分割能力。为了进一步加快训练速度，同时挖掘少数训练者能够发现的数据特征，提供了基于动量梯度的全局非交集掩码，使得全局模型的泛化效果进一步地提升（图4）。

图4 基础方法（上）与sUSAM（下）的损失平面对比图

作为一种具有隐私保护效果的类别通用分割框架，UFPS能够对参与训练的客户数据集取得较好的分割性能；相较于已有研究，对潜在的联邦客户也能有很好的泛化效果，以此鼓励更多的客户加入联邦合作；对于易产生混淆的脾脏类别，分割效果甚至超过了全类别标注的中心化训练策略（图5）。

图5 客户级性能（左）与类别级性能（右）

总的来说，该研究调研了部分类别监督方法在联邦学习中的训练效果，设计了能够保护数据隐私的UFPS，一种可以处理无标注类别交集的特征异质医学图像分割框架。该框架包含了ULL和sUSAM，分类实现类别和特征空间的双重对齐。此外，该方法与模型架构的选择无关，为通用医疗大模型的训练提供了可行的处理方案。

UFPS相关网址如下：

Python源码：

https://github.com/tekap404/unified_federated_partially-labeled_segmentation

作者介绍

马丽艳，上海大学计算机工程与科学学院副研究员。2013年于北京交通大学获得博士学位，2013-2018年就职于中国科学院微电子研究所，先后任助理研究员和副研究员。2018年至今在上海大学任教。主要研究方向为医学影像分析、计算机视觉和机器学习。相关工作发表在Patterns、SIAM、IEEE Transaction汇刊等期刊和会议上。