Federated Learning over Connected Modes

2024-03-05

Neural Information Processing Systems

AI Summary

Statistical heterogeneity in federated learning poses two key challenges: slow global convergence and difficulty in achieving effective personalization. This paper introduces Linear Mode Connectivity (LMC) to federated learning for the first time, proposing a collaborative training framework based on a *solution simplex*: it constructs a shared, low-loss linearly connected region in parameter space and dynamically assigns each client a local subregion within it, thereby unifying rapid global convergence with local distribution adaptation. The method integrates gradient-signal-driven subregion allocation, solution-simplex parameterization, and joint optimization. Experiments under realistic cross-institutional federated settings demonstrate significant improvements in both global convergence speed and local accuracy, with negligible computational overhead. The core contribution lies in the novel application of LMC theory to model the structural geometry of the federated solution space—enabling simultaneous gains in efficiency and personalization.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Learning to Search
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

Abstract

Statistical heterogeneity in federated learning poses two major challenges: slow global training due to conflicting gradient signals, and the need of personalization for local distributions. In this work, we tackle both challenges by leveraging recent advances in emph{linear mode connectivity} -- identifying a linearly connected low-loss region in the parameter space of neural networks, which we call solution simplex. We propose federated learning over connected modes ( extsc{Floco}), where clients are assigned local subregions in this simplex based on their gradient signals, and together learn the shared global solution simplex. This allows personalization of the client models to fit their local distributions within the degrees of freedom in the solution simplex and homogenizes the update signals for the global simplex training. Our experiments show that extsc{Floco} accelerates the global training process, and significantly improves the local accuracy with minimal computational overhead in cross-silo federated learning settings.

Problem

• Federated learning faces slow training due to conflicting gradients.
• Personalization of local models in federated learning is challenging.
• Statistical heterogeneity complicates global and local model optimization.

Innovation

• Federated Learning
• Linear Mode Connectivity
• Solution Simplex