ICML 2024 | MH-pFLID：通过注入和蒸馏的模型异构个性化联邦学习用于医疗数据分析

DRUGAI

今天为大家介绍的是来自Tianyu Luan团队的一篇论文。联邦学习广泛应用于医疗领域，用于在不需要访问本地数据的情况下训练全局模型。然而，由于客户间不同的计算能力和网络架构（系统异构性），从非独立同分布（non-IID）数据中有效聚合信息面临着重大挑战。当前使用知识蒸馏的联邦学习方法需要公共数据集，这会引发隐私和数据收集问题。此外，这些数据集需要额外的本地计算和存储资源，这对于硬件条件有限的医疗机构来说是一个负担。在本文中，作者引入了一种新颖的联邦学习范式，称为基于注入和蒸馏的模型异构个性化联邦学习（MH-pFLID）。作者的框架利用一个轻量级的信使模型来收集每个客户的信息。作者还开发了一套接收器和发射器模块，用于接收和发送来自信使模型的信息，以便高效地注入和蒸馏信息。作者的框架消除了对公共数据集的需求，并在客户之间高效地共享信息。作者在各种医学任务上的实验表明，MH-pFLID 在所有这些领域均优于现有的最先进方法，并具有良好的泛化能力。

联邦学习广泛应用于医疗领域，因为它能够在不需要直接访问每个医疗机构的本地数据的情况下，协作训练全局模型。考虑到不同医疗机构的本地数据大多是非独立同分布的，之前的工作如主要集中在解决统计异质性问题，即数据分布是客户间的主要差异。然而，系统异质性，例如客户之间的模型架构差异，也是现实应用中常见的情况。不同的需求、硬件资源考虑以及监管风险导致不同的医院选择和维护不同的模型架构，而不同的模型架构会导致不同的拟合能力、性能表现和模型聚合困难。广泛的应用和问题的复杂性使得模型异质性问题成为一个关键且具有挑战性的任务。

图 1

如图1(a)所示，现有方法解决统计和系统异质性问题基于知识蒸馏。这些方法涉及在不同客户之间交换公共数据集上的软预测，从而实现知识的转移。虽然这些方法在解决系统和统计异质性问题方面取得了进展，但它们仍然依赖于生成软预测的前提公共数据集。医疗公共数据集的高隐私问题和复杂的数据收集与发布过程限制了这些方法的实际应用。此外，大规模公共数据集使得本地训练具有挑战性。许多医疗机构难以获得用于训练大规模公共数据集的计算资源。考虑到这些额外的计算资源在推理时不需要，这种负担更为沉重。

为了消除对公共数据集的依赖，作者提出了一种新的注入-蒸馏范式，以解决非IID数据分布下的异构模型问题。与传统方法依赖公共数据集生成的软预测不同，作者的方法利用一个极其轻量级的信使模型进行信息传递。范式包括三个步骤：知识注入、知识蒸馏和聚合，模型将信使模型插入每个本地客户。在知识注入阶段，知识从信使注入到每个本地模型。在知识蒸馏阶段，客户的知识在本地数据训练时被蒸馏到信使模型中。在知识聚合阶段，通过结合信使模型的参数来聚合知识。此外，信使模型的小参数量确保其本地训练相比于本地数据训练增加的负担极小。这种额外负担远小于在公共数据集上的本地训练所需的负担。

模型方法

问题公式化

作者的目标是为每个客户i训练一个个性化模型 f i(θ i; x)，其中θ i是f i的参数，x是模型输入。每个模型f i只能访问其自己的私有数据集 Di={xij, yij}，其中xij是Di中的第j个输入数据，yij是第j个标签。作者协同训练所有的fi，使每个模型都能利用其他客户的数据集中的信息，而不直接访问这些数据。在实际的医疗场景中，每个客户的数据分布通常是非独立同分布的。MH-pFLID的范式可以表示为：

其中其中G表示fi的集合，N是参与客户的总数。由于每个客户采用的模型架构是定制的，每个客户的模型结构fi是不同的（系统异质性）。因此，MH-pFLID同时面临两个主要挑战：客户间的数据分布差异（统计异质性）和不同的模型架构（系统异质性）。

流程

图 2

MH-pFLID的流程如图2所示。本地和信使模型都分为一个提取特征的主体模型和一个使用这些特征生成网络输出的头部模型。训练过程包括五个步骤：知识注入、知识蒸馏、上传信使模型到服务器、信使聚合和下载聚合后的信使信息。以下是每个步骤的详细说明。

知识注入：知识注入阶段旨在将来自信使模型的信息注入到本地模型中。具体而言，作者冻结信使模型，并利用信使模型来指导本地模型的训练。对于客户i，其知识注入阶段的训练损失函数Linj,i表示为：

其中，M代表本地数据的总量，yijl和yijm分别是本地模型和信使模型对第j个数据的预测。Linjl和Linjm分别表示本地模型和信使模型的损失函数，λinjl 和 λinjm是它们对应的权重，yij是xij的标签。

知识蒸馏：知识蒸馏阶段旨在将本地模型的信息蒸馏到信使模型中。具体而言，作者冻结本地模型，并对信使模型进行知识蒸馏，其损失函数Ldis,i表示为：

其中，Ldism和Ldiscon分别表示训练信使模型的损失函数和知识蒸馏损失函数。对于知识蒸馏损失，作者使用KL散度来约束信使头部和本地头部模型的输出，使得知识能够从本地模型蒸馏到信使模型中。λdism和λdiscon是它们对应的权重。其他变量定义与注入阶段相同。在知识蒸馏阶段的训练中，作者尝试同时使用真实标签yij和本地模型输出yijl来监督信使模型的输出yijm，以便知识能够被蒸馏到信使模型中。

信使上传、聚合和下载：训练后，将信使模型上传到服务器。然后分别对信使的主体和头部模型参数进行聚合，使用的聚合操作是来自McMahan等人的加权平均，将所有参数相加再一起除以总数。最后，下载聚合后的模型并分发到每个客户，进行下一轮训练。

在推理阶段，作者直接使用经过良好训练的本地异构模型进行推理。与其他现有方法相比，MH-pFLID消除了对公共数据集的需求，仅在本地插入轻量级的信使模型。这增强了在医疗场景中异构模型联邦学习的应用。

信息接收器和信息发射器

信息接收器和发射器设计用于在本地模型和信使模型之间进行有效通信。为了实现轻量级的信使模型，接收器和发射器设计使得信使只需要携带本地模型蒸馏的信息代码本。接收器模块结合信使代码本中的信息并将其注入到本地模型中，而发射器模块则从本地模型中提取信息并将其蒸馏到信使模型中。

图 3

信息接收器：信息接收器可以定义为Iloc,R=R(Iloc,Imes)，其中Iloc和Imes分别是本地和信使主体的输出特征。Iloc,R是接收器的输出，是本地主体特征的加权组合。在知识注入阶段，作者设计信息接收器以更好地匹配本地特征与全局特征，使得本地模型能够更好地接收全局知识。信息接收器，如图3(a)所示，首先生成本地客户特征Iloc，然后通过线性层Wd进行上采样或下采样，生成与信使特征Imes具有相同维度的特征I′loc。

信息发射器：信息发射器可以定义为Imes,T=T(Imes,Iloc)，其中Imes,T是发射器的输出，是信使主体特征的加权组合，如图3(b)所示。与知识注入阶段类似，在知识蒸馏阶段，作者让全局特征Imes学习处理后的本地特征I′loc的知识。Imes和I′loc用于生成查询特征Q、键特征K和值特征V，通过Wk、Wq和Wv计算得到。

实验结果

任务与数据集

作者在四个非独立同分布任务上验证了MH-pFLID的有效性。

A. 医学图像分类（不同分辨率）：作者使用了乳腺癌组织病理图像数据库（BreaKHis）。作者对高分辨率图像进行了2×、4×和8×的降采样处理。每种分辨率的医学图像被视为一个客户端，总共四个客户端。每个客户端的数据集随机分为训练集和测试集，比例为7:3，遵循之前的工作。对于不同分辨率的相同图像，它们将在训练集或测试集中使用。在这个任务中，作者使用ResNet{17, 11, 8, 5}。

B. 医学图像分类（不同标签分布）：此任务包括乳腺癌分类任务和眼科疾病（OCT）分类任务。作者设计了八个客户端，每个客户端对应一个异构模型。这些模型包括ResNet、ShuffleNetV2、ResNeXt、SqueezeNet、SENet、MobileNetV2、DenseNet和VGG。类似于FedAvg，作者在BreaKHis（RGB图像）和OCT2017（灰度图像）数据集上的8个客户端应用了非IID标签分布方法。具体来说，在不同客户端中，每个客户端的标签是不同的。此外，数据分布在客户端之间也是不同的。

C. 医学时间序列分类：作者使用了 Sleep-EDF 数据集来进行时间序列分类任务。在非IID分布下，作者设计了三个客户端，分别使用TCN、Transformer和RNN。

D. 医学图像分割：在这里，作者专注于息肉分割。数据集由四个中心收集和标注的内窥镜图像组成，每个中心的数据集被视为一个独立的客户端。每个客户端使用特定的模型，包括 Unet++、FCN、Unet和Res-Unet。

这些任务和数据集的多样性帮助作者全面验证了 MH-pFLID 在处理不同类型的医学数据时的性能和泛化能力。

医学图像分类（不同分辨率）结果

表 1

在表1中，与其他联邦学习框架相比，MH-pFLID在所有不同分辨率的客户端（包括原始高分辨率、降采样到一半（“×2↓”）、四分之一（“×4 ↓”）和八分之一（“×8 ↓”）的图像）上均取得了最佳性能。这表明，基于注入和蒸馏范式的MH-pFLID能够有效地让同一家族内的本地异构模型学习全局知识，从而提高本地模型的性能。此外，MH-pFLID在MF1指标上表现出更显著的优势，突显了其在提高本地异构模型鲁棒性方面的能力。

医学图像分类（不同标签分布）结果

表 2

在表2中，不同标签分布的医学图像分类任务的实验结果显示，每个客户端使用异构模型时，MH-pFLID取得了最佳结果。这表明，与基于软预测的异构联邦学习方法相比，MH-pFLID的注入和蒸馏方法具有优势。它能够更有效地利用来自其他客户端的知识来指导本地客户的学习。相比于单独的本地训练，MH-pFLID提升了每个异构模型的本地性能。这表明，作者提出的特征适配方法通过对齐全局和本地特征，在指导每个异构模型的客户训练时，有效缓解了客户漂移问题。

时间序列分类结果

表 3

实验结果如表3所示，在不同类型的神经网络下，MH-pFLID取得了最佳结果。这进一步证明了MH-pFLID在异构模型联邦学习中的优越性。具体来说，作者在Sleep-EDF数据集上进行了时间序列分类任务的实验，设计了三个客户端，分别使用TCN、Transformer和RNN模型。实验结果表明，MH-pFLID在所有模型上的性能都优于其他现有的联邦学习方法。此结果强调了MH-pFLID在处理异构模型和非独立同分布数据时的有效性和鲁棒性。

医学图像分割结果

表 4

作者再次验证了MH-pFLID在医学图像分割任务中的有效性。表4展示了在分割任务中使用联邦学习的结果，显示MH-pFLID取得了最佳实验结果。这表明MH-pFLID在多种任务中的应用潜力。实验结果进一步强调，MH-pFLID能够有效提升各客户端本地模型的性能，优于现有的同类个性化方法。在不同任务中，MH-pFLID能够提高本地异构模型的表现，展示了其在医学应用中的广泛适用性。

图 4

此外，图4中的可视化结果显示，MH-pFLID 的分割结果更接近于真实情况（ground truth）。这些结果验证了 MH-pFLID 的分割性能，使得其在医学图像分析任务中展现出明显的优势和应用前景。

局限性与结论

作者的方法在医学分类和分割任务中展示了其有效性，但作者尚未在医学对象检测、图像配准、医学3D重建等领域验证和完善作者的方法。在未来的工作中，作者的方法在这些领域的潜力还有待进一步验证和增强。

MH-pFLID有效解决了现有异构模型个性化联邦学习方法面临的挑战，包括公共数据集的收集和标注、本地客户端和服务器的计算负担。基于注入和蒸馏范式的MH-pFLID提供了解决这些问题的方案。MH-pFLID在每个客户端引入了一个轻量级信使模型，并设计了信息接收器和发射器。这些模块使本地异构模型能够在非独立同分布的数据下，良好地转移来自其他客户端的信息。广泛的实验表明，MH-pFLID在异构模型的联邦学习框架中，相较于现有方法具有显著的优越性。

编译 | 于洲

审稿 | 曾全晨

参考资料

Xie L, Lin M, Luan T, et al. MH-pFLID: Model Heterogeneous personalized Federated Learning via Injection and Distillation for Medical Data Analysis[J]. arXiv preprint arXiv:2405.06822, 2024.