[论文解读] FedSoup: Improving Generalization and Personalization in Federated Learning via Selective Model Interpolation

Ming‐Hui Chen, Meirui Jiang+3 more|arXiv (Cornell University)|Jul 20, 2023

AI in cancer detectionComputer Science被引用 3

一句话总结

FedSoup 提出了一种在联邦学习中选择性模型插值的方法，通过维护客户端特定的全局模型池并在本地模型与全局模型之间进行插值，同时提升了泛化能力和个性化性能。该方法通过时间模型选择和联邦模型修补实现更平坦的极小值，显著提升了分布外泛化性能，且无需增加推理成本。

ABSTRACT

Cross-silo federated learning (FL) enables the development of machine learning models on datasets distributed across data centers such as hospitals and clinical research laboratories. However, recent research has found that current FL algorithms face a trade-off between local and global performance when confronted with distribution shifts. Specifically, personalized FL methods have a tendency to overfit to local data, leading to a sharp valley in the local model and inhibiting its ability to generalize to out-of-distribution data. In this paper, we propose a novel federated model soup method (i.e., selective interpolation of model parameters) to optimize the trade-off between local and global performance. Specifically, during the federated training phase, each client maintains its own global model pool by monitoring the performance of the interpolated model between the local and global models. This allows us to alleviate overfitting and seek flat minima, which can significantly improve the model's generalization performance. We evaluate our method on retinal and pathological image classification tasks, and our proposed method achieves significant improvements for out-of-distribution generalization. Our code is available at https://github.com/ubc-tea/FedSoup.

研究动机与目标

• 解决在数据分布偏移下跨孤岛联邦学习中本地个性化与全局泛化之间的权衡问题。
• 克服个性化联邦学习模型中因过拟合导致的尖锐极小值和分布外性能差的问题。
• 开发一种通信高效的单次训练方法，将模型汤技术适配至联邦学习设置，无需重新训练。
• 通过插值历史全局模型和本地更新，寻找更平坦的极小值，提升模型鲁棒性。
• 在不增加推理或内存开销的前提下，保持对分布内和分布外数据的高性能。

提出的方法

• 在联邦训练过程中维护基于本地验证性能选择的历史全局模型池，且为每个客户端独立维护。
• 采用贪心的时间模型选择策略，识别适合插值的高性能全局模型，避免选择位于不同损失景观基域中的模型。
• 通过插值选定的全局模型与本地模型，构建‘汤’（soup），形成个性化且泛化的模型。
• 引入联邦模型修补：通过将本地模型与全局模型汤插值来微调，以平衡个性化与全局性能。
• 使用幂迭代估计中位数主导 Hessian 特征值作为尖锐度度量，以量化极小值的平坦度。
• 借鉴模型汤原理——高性能模型的权重平均——通过在单次训练运行中使用不同训练周期的模型，将其适配至联邦学习场景。

实验结果

研究问题

• RQ1 在数据异质性下，联邦学习中本地个性化与全局泛化之间性能权衡的根本原因是什么？
• RQ2 如何利用模型插值在联邦设置中实现更平坦的极小值并提升泛化能力？
• RQ3 能否通过插值不同训练周期的模型来同时提升个性化与分布外泛化性能？
• RQ4 如何在不需多次重新训练或高通信成本的前提下，将模型汤技术适配至联邦学习？
• RQ5 FedSoup 在小规模本地数据集上在多大程度上缓解了过拟合，同时保持了全局性能？

主要发现

• 在 Camelyon17 病理数据集上，FedSoup 相较于 FedAvg 在未见领域泛化任务中实现了 2.87 个百分点的 AUC 提升。
• 在视网膜眼底图像数据集上，FedSoup 在分布外泛化任务中达到 86.24% 的准确率，优于 FedBABU（85.09%）和 FedProx（85.18%），提升超过 1.05 个百分点。
• 损失景观中的尖锐度降低，表现为中位数主导 Hessian 特征值下降，表明极小值更平坦。
• FedSoup 在保持高本地性能（85.71% 准确率）的同时实现了更优的全局性能（96.00% 准确率），相较于基线方法展现出更优的本地-全局权衡。
• 该方法表现出更强的稳定性，实验中方差更低（例如，Retina 数据集上全局 AUC 的方差为 0.43，而 FedAvg 为 0.83）。
• 性能增益在较小数据集（如视网膜眼底数据集）上最为显著，表明其在缓解因本地数据有限导致的过拟合方面具有显著有效性。