[论文解读] Real World Federated Learning with a Knowledge Distilled Transformer for Cardiac CT Imaging

Malte Tölle, Philipp Garthe+23 more|arXiv (Cornell University)|Jul 10, 2024

Advanced X-ray and CT Imaging被引用 4

一句话总结

本文提出了一种两阶段半监督联邦学习框架，通过将特定任务的CNN模型知识蒸馏到单一视觉Transformer中，实现对八家医院心脏CT影像的局部标注数据的有效利用。该方法在训练统一Transformer以预测多种心脏结构的同时，提升了模型的准确性和泛化能力，开源的代码与权重为未来研究提供了基础模型。

ABSTRACT

Federated learning is a renowned technique for utilizing decentralized data while preserving privacy. However, real-world applications often face challenges like partially labeled datasets, where only a few locations have certain expert annotations, leaving large portions of unlabeled data unused. Leveraging these could enhance transformer architectures ability in regimes with small and diversely annotated sets. We conduct the largest federated cardiac CT analysis to date (n=8,104) in a real-world setting across eight hospitals. Our two-step semi-supervised strategy distills knowledge from task-specific CNNs into a transformer. First, CNNs predict on unlabeled data per label type and then the transformer learns from these predictions with label-specific heads. This improves predictive accuracy and enables simultaneous learning of all partial labels across the federation, and outperforms UNet-based models in generalizability on downstream tasks. Code and model weights are made openly available for leveraging future cardiac CT analysis.

研究动机与目标

• 解决由于隐私限制和标签异质性，从多家医院的局部标注心脏CT数据集中训练鲁棒深度学习模型的挑战。
• 利用临床环境中通常未被使用的大量未标注数据，以提升模型性能。
• 开发一种联邦学习框架，实现在标签覆盖程度各异的机构之间，同时学习多种心脏结构（如钙化、冠状动脉）的联合训练。
• 创建一个可泛化的、开源的基础模型，用于心脏CT影像，其在下游任务中的表现优于标准CNN。
• 通过在本地训练模型并仅聚合模型权重，避免集中式数据存储，确保隐私合规。

提出的方法

• 采用两阶段半监督学习策略：首先在每个客户端本地训练特定任务的CNN（如关键点检测和钙化分割）模型。
• 通过知识蒸馏，利用所有客户端的标注数据和丰富的未标注数据，将这些CNN模型的知识迁移至单一视觉Transformer模型。
• 使用联邦学习在八家医院站点训练Transformer模型，本地训练后仅将模型权重集中聚合，从而保护数据隐私。
• 通过仅微调Transformer模型的最后1层，将最终模型用于下游任务——冠状动脉分割，展示了更优的特征表示能力。
• FedBioMed框架通过TLS加密、客户端身份验证和IP白名单，确保通信安全，实现单向、安全的联邦学习流程。
• 使用自定义Docker编排和基于PyTorch的流水线，结合MONAI和nnU-Net，实现各站点的预处理、训练和验证。

实验结果

研究问题

• RQ1 单一视觉Transformer模型能否有效从多家医院的局部标注数据中学习多种心脏结构（如钙化、冠状动脉）？
• RQ2 在联邦学习设置中，将特定任务CNN的知识蒸馏到Transformer是否能提升下游分割任务的性能？
• RQ3 在隐私保护的联邦学习设置中，多所医院的未标注数据能在多大程度上被利用以提升模型的泛化能力？
• RQ4 与基于CNN的基线模型相比，联邦Transformer模型在特征质量与分割准确性方面表现如何？
• RQ5 能否在不集中收集数据的前提下，跨异质标签分布训练统一的基础模型？

主要发现

• 所提出的联邦知识蒸馏Transformer模型在心脏CT影像任务中的表现优于基线UNet，尤其在下游冠状动脉分割任务中表现更优。
• 通过知识蒸馏利用未标注数据，即使每家客户端的标注数据有限，模型仍显著提升了预测准确性和泛化能力。
• 该方法实现了在联邦架构中，通过单一Transformer模型同时学习多种部分标签（如钙化、关键点、冠状动脉）的能力。
• 仅微调Transformer的最后1层进行冠状动脉分割，其学习到的特征比CNN基线更具语义意义，表明其表示学习能力更强。
• 最终模型成功在八家医院完成训练，尽管各医院标签覆盖程度不同，证明了在真实联邦环境中的可扩展性与鲁棒性。
• 代码、模型权重和数据处理流水线均以符合FAIR原则的方式开源，支持心脏影像研究中的重复使用与进一步开发。