Enhancing Federated Learning Generalization Through Momentum Alignment in 6G Networks

As artificial intelligence (AI) technology continues to advance and the demand for AI applications grows, the development of sixth-generation (6G) wireless networks has become increasingly targeted. These networks are designed to be highly compatible with AI technologies, laying the foundation for revolutionary communication capabilities. In this context, federated learning, a technique that trains models across multiple distributed nodes without sharing raw data, becomes particularly significant. However, deploying federated learning in a 6G wireless network environment faces several challenges. Firstly, data heterogeneity may lead to decreased performance of the global model trained. Secondly, global models might converge to sharp loss landscape during distributed training, which weakens their generalizability. Some studies improve generalization by optimizing local models to converge to flatter loss landscape. However, this does not necessarily ensure that the aggregated global model also converges to a flatter landscape. To address these challenges, this paper introduces Federated Momentum Alignment (FedMA), a novel method designed to optimize the performance of client models in 6G networks by aligning momentum during model training. In the FedMA, the momentum alignment mechanism we propose not only prompts local models to reference future update directions during training to alleviate model drift caused by data heterogeneity, but also facilitates the convergence of global models to a flat loss landscape by aligning global objectives with local ones, thereby enhancing the model’s generalization ability. Experimental results show that FedMA improves performance by 4.62% compared to the state-of-the-art federated learning methods.