[Paper Review] 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 Science3 citations

TL;DR

FedSoup proposes a selective model interpolation method in federated learning that improves both generalization and personalization by maintaining client-specific global model pools and interpolating between local and global models. It achieves flatter minima through temporal model selection and federated model patching, significantly boosting out-of-distribution generalization with no added inference cost.

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.

Motivation & Objective

• Address the trade-off between local personalization and global generalization in cross-silo federated learning under data distribution shifts.
• Overcome overfitting in personalized FL models that leads to sharp minima and poor out-of-distribution performance.
• Develop a communication-efficient, single-training-session method to adapt model soups to federated settings without retraining.
• Improve model robustness by seeking flatter minima through interpolation of historical global models and local updates.
• Maintain high performance on both in-distribution and out-of-distribution data without increasing inference or memory overhead.

Proposed method

• Maintain a client-specific pool of historical global models during federated training, selected based on local validation performance.
• Apply a greedy temporal model selection strategy to identify high-performing global models for interpolation, avoiding models in different loss landscape basins.
• Construct a 'soup' by interpolating selected global models and the local model to form a personalized, generalized model.
• Introduce federated model patching: fine-tune the local model by interpolating it with the global model soup to balance personalization and global performance.
• Use power iteration to estimate the median dominating Hessian eigenvalue as a sharpness metric to quantify flatness of minima.
• Leverage model soup principles—weight averaging of high-performing models—adapted to FL by using models from different training epochs within a single run.

Experimental results

Research questions

• RQ1 What causes the performance trade-off between local personalization and global generalization in federated learning under data heterogeneity?
• RQ2 How can model interpolation be used to achieve flatter minima and improve generalization in federated settings?
• RQ3 Can selective interpolation of models from different training epochs improve both personalization and out-of-distribution generalization?
• RQ4 How can model soup techniques be adapted to federated learning without requiring multiple retraining runs or high communication costs?
• RQ5 To what extent does FedSoup mitigate overfitting on small local datasets while preserving global performance?

Key findings

• FedSoup achieved a 2.87-point AUC improvement on unseen domain generalization in the Camelyon17 pathology dataset compared to FedAvg.
• On the retinal fundus dataset, FedSoup achieved 86.24% accuracy on out-of-distribution generalization, outperforming FedBABU (85.09%) and FedProx (85.18%) by over 1.05 points.
• The method reduced sharpness in the loss landscape, as evidenced by lower median dominating Hessian eigenvalues, indicating flatter minima.
• FedSoup maintained high local performance (85.71% accuracy) while achieving superior global performance (96.00% accuracy), demonstrating a better local-global trade-off than baseline methods.
• The method showed greater stability, with lower variance across experiments (e.g., 0.43 for global AUC on Retina vs. 0.83 for FedAvg).
• The performance gain was most pronounced on smaller datasets like retinal fundus, indicating effectiveness in mitigating overfitting due to limited local data.