[论文解读] Federated Split Vision Transformer for COVID-19 CXR Diagnosis using Task-Agnostic Training

Sang Joon Park, Gwanghyun Kim+3 more|arXiv (Cornell University)|Nov 2, 2021

COVID-19 diagnosis using AI参考文献 67被引用 19

一句话总结

该论文提出了一种联邦分割视觉变换器（FeSTA）框架，结合联邦学习与分割学习，实现在医院间隐私保护的协作式新冠肺炎胸部X光（CXR）诊断。通过利用视觉变换器的模块化、可分解架构，FeSTA在客户端与服务器之间拆分模型，减少了带宽和计算负载，同时在非独立同分布（non-IID）数据分布下以及多任务学习设置中，仍能保持与集中式训练相当的性能表现。

ABSTRACT

Federated learning, which shares the weights of the neural network across clients, is gaining attention in the healthcare sector as it enables training on a large corpus of decentralized data while maintaining data privacy. For example, this enables neural network training for COVID-19 diagnosis on chest X-ray (CXR) images without collecting patient CXR data across multiple hospitals. Unfortunately, the exchange of the weights quickly consumes the network bandwidth if highly expressive network architecture is employed. So-called split learning partially solves this problem by dividing a neural network into a client and a server part, so that the client part of the network takes up less extensive computation resources and bandwidth. However, it is not clear how to find the optimal split without sacrificing the overall network performance. To amalgamate these methods and thereby maximize their distinct strengths, here we show that the Vision Transformer, a recently developed deep learning architecture with straightforward decomposable configuration, is ideally suitable for split learning without sacrificing performance. Even under the non-independent and identically distributed data distribution which emulates a real collaboration between hospitals using CXR datasets from multiple sources, the proposed framework was able to attain performance comparable to data-centralized training. In addition, the proposed framework along with heterogeneous multi-task clients also improves individual task performances including the diagnosis of COVID-19, eliminating the need for sharing large weights with innumerable parameters. Our results affirm the suitability of Transformer for collaborative learning in medical imaging and pave the way forward for future real-world implementations.

研究动机与目标

• 解决在去中心化医院之间训练鲁棒医疗AI模型的同时保护患者数据隐私的挑战。
• 克服在使用如视觉变换器等大型表达性模型时，联邦学习面临的带宽和计算限制。
• 整合联邦学习（全局模型聚合）与分割学习（减少数据传输和客户端资源消耗）的优势，用于医学影像任务。
• 在异构客户端之间实现有效的多任务学习，提升各任务性能，同时不共享原始数据或大模型权重。
• 证明与任务无关的视觉变换器主干网络可泛化于多种CXR任务，包括新冠肺炎诊断、分割与检测。

提出的方法

• 将视觉变换器分解为三部分：客户端的头部（特征提取）、共享服务器端的主干（自注意力建模）以及客户端特定的尾部（任务特定输出）。
• 通过仅传输中间特征（而非完整梯度或权重）实施分割学习，最大限度减少带宽并增强隐私保护。
• 利用联邦学习在服务器端聚合客户端更新，实现在不交换原始数据的前提下进行全局模型优化。
• 采用两步训练策略：首先在多个任务上预训练共享的与任务无关的主干，然后在各客户端上微调特定任务的头部和尾部。
• 采用多任务学习设置，使同一共享主干在诊断、分割和检测任务中学习共享表征。
• 将框架与Flower联邦学习框架集成，实现在真实异构客户端环境中无缝部署与测试。

实验结果

研究问题

• RQ1 视觉变换器的模块化架构是否能有效支持医学影像中的分割学习，同时保持模型性能？
• RQ2 在典型多中心医院协作中常见的非独立同分布（non-i.i.i.d.）数据分布下，所提出的FeSTA框架表现如何？
• RQ3 通过共享的、与任务无关的视觉变换器主干进行多任务学习，在如新冠肺炎诊断等单个下游任务上的性能提升程度如何？
• RQ4 两步训练策略在收敛性和泛化能力方面是否优于一步训练或交替训练？
• RQ5 与标准联邦学习相比，该框架是否能在显著降低客户端计算和通信开销的同时保持高性能？

主要发现

• 在新冠肺炎诊断任务中，FeSTA框架在非独立同分布数据分布下仍实现了0.931 ± 0.004的平均AUC，与集中式训练性能相当。
• 更大的视觉变换器主干（H=12, L=12, D_hidden=768）优于较小配置，表明更高容量仍有进一步提升性能的潜力。
• 两步训练策略表现最佳（AUC 0.931），优于一步和交替训练方法，表明其具有更优的收敛性与泛化能力。
• 采用共享的、与任务无关主干的多任务学习显著提升了所有任务的性能：诊断（AUC 0.931）、分割（Dice 0.821）、检测（mAP 0.204）。
• 该框架在异构客户端中保持了稳定性能，并通过避免原始数据和大模型权重的共享，有效实现了隐私保护。
• 消融实验证实，所提出的训练方案与模型容量对于在联邦分割学习设置中实现最优性能至关重要。