[论文解读] FedDAR: Federated Domain-Aware Representation Learning

Aoxiao Zhong, Hao He+3 more|arXiv (Cornell University)|Sep 8, 2022

Machine Learning in Healthcare被引用 4

一句话总结

FedDAR 提出了一种新颖的联邦学习框架，用于处理领域混合的非独立同分布（non-i.i.d.）数据，其中每个客户端的数据是预定义领域的混合。该方法通过在不同领域间共享编码器，同时使用领域特定的分类头，将表示学习与预测解耦，借助交替优化和二阶聚合实现领域感知的个性化，实现了线性收敛，并在合成数据集和真实世界医疗数据集上表现出优越性能。

ABSTRACT

Cross-silo Federated learning (FL) has become a promising tool in machine learning applications for healthcare. It allows hospitals/institutions to train models with sufficient data while the data is kept private. To make sure the FL model is robust when facing heterogeneous data among FL clients, most efforts focus on personalizing models for clients. However, the latent relationships between clients' data are ignored. In this work, we focus on a special non-iid FL problem, called Domain-mixed FL, where each client's data distribution is assumed to be a mixture of several predefined domains. Recognizing the diversity of domains and the similarity within domains, we propose a novel method, FedDAR, which learns a domain shared representation and domain-wise personalized prediction heads in a decoupled manner. For simplified linear regression settings, we have theoretically proved that FedDAR enjoys a linear convergence rate. For general settings, we have performed intensive empirical studies on both synthetic and real-world medical datasets which demonstrate its superiority over prior FL methods.

研究动机与目标

• 通过将客户端数据建模为预定义领域的混合，而非假设独立同分布（i.i.d.）或客户端特定分布，来解决跨孤岛联邦学习中的数据异构性问题。
• 克服现有联邦学习方法在个性化客户端时忽略潜在领域结构的局限性，特别是在医疗领域中，种族或扫描仪协议等领域的差异会导致分布偏移。
• 开发一种方法，在学习跨领域共享表示的同时，支持领域特定的预测头，以提升模型的鲁棒性和公平性。
• 从理论和实证上验证，在领域混合的联邦学习设置下，基于领域的个性化优于基于客户端的个性化。

提出的方法

• FedDAR 将模型分解为共享编码器和领域特定的预测头，将表示学习与预测解耦，以提升泛化能力。
• 通过在多个本地轮次中交替更新本地分类头，并使用加权平均聚合编码器权重，同时通过样本重加权确保领域公平性。
• 在分类头聚合方面，FedDAR 采用一种新颖的二阶优化策略，以提升全局分类头的收敛性和最优性。
• 该方法采用两阶段训练流程：先进行 FedAvg 热身，随后交替执行本地更新和编码器与分类头的全局聚合。
• 在训练过程中应用领域感知的损失重加权，以平衡低资源领域对模型训练的贡献。
• 该框架支持线性和非线性模型，在真实医疗应用中对卷积神经网络（CNNs）和 ResNet-34 进行了评估。

实验结果

研究问题

• RQ1 是否有一种联邦学习方法，通过显式建模领域混合，能在非独立同分布的医疗数据设置下超越标准联邦学习和个性化联邦学习基线？
• RQ2 与联合优化相比，将表示学习与领域特定预测解耦是否能提升收敛速度和性能？
• RQ3 对预测头进行领域感知聚合，如何影响模型在多样化客户端数据组成下的公平性和泛化能力？
• RQ4 在线性设置下，该方法可建立哪些理论保证，例如收敛速率？
• RQ5 FedDAR 在真实世界医疗数据集中在多大程度上减少了偏差，并提升了低资源领域的性能？

主要发现

• 在线性回归设置下，FedDAR 实现了线性收敛，展现出强大的理论收敛特性。
• 在合成数据上，FedDAR 在所有领域中的测试准确率均优于 FedAvg、FedProx、FedPer 和 FedRep。
• 在 FairFace 数据集上，与最先进方法相比，FedDAR 在所有种族群体上均实现了更高的平均准确率，同时提升了公平性并减少了偏差。
• 在真实世界的 EXAM 数据集上，FedDAR 在本地验证集上实现了平均 AUC 0.89，显著优于 FedProx、FedRep 和 FedMinMax。
• 消融研究证实，二阶分类头聚合和领域感知重加权对性能至关重要，尤其在数据不平衡的领域中表现显著。
• 即使某个客户端每个领域的样本仅 100 个，FedDAR 仍能保持优异性能，展现出对低资源领域的强鲁棒性。