Deep Unfolding-based Weighted Averaging for Federated Learning in Heterogeneous Environments

Ayano Nakai-Kasai,Tadashi Wadayama

2023-08-28

Abstract:Federated learning is a collaborative model training method that iterates model updates by multiple clients and aggregation of the updates by a central server. Device and statistical heterogeneity of participating clients cause significant performance degradation so that an appropriate aggregation weight should be assigned to each client in the aggregation phase of the server. To adjust the aggregation weights, this paper employs deep unfolding, which is known as the parameter tuning method that leverages both learning capability using training data like deep learning and domain knowledge. This enables us to directly incorporate the heterogeneity of the environment of interest into the tuning of the aggregation weights. The proposed approach can be combined with various federated learning algorithms. The results of numerical experiments indicate that a higher test accuracy for unknown class-balanced data can be obtained with the proposed method than that with conventional heuristic weighting methods. The proposed method can handle large-scale learning models with the aid of pretrained models such that it can perform practical real-world tasks. Convergence rate of federated learning algorithms with the proposed method is also provided in this paper.

Machine Learning,Signal Processing

What problem does this paper attempt to address?

The paper attempts to address the problem of how to improve model performance in federated learning within heterogeneous environments by optimizing aggregation weights. Specifically, due to the heterogeneity of devices and data involved in federated learning (such as device configurations, network conditions, inconsistent data distributions, etc.), the effectiveness of model training can be reduced. Therefore, the paper proposes a method based on deep unfolding to adjust aggregation weights, allowing these weights to directly reflect the heterogeneity of the environment. This enables the assignment of appropriate weights to each client during the aggregation phase of federated learning, thereby improving the model's test accuracy and convergence speed.

The main contributions of the paper include:
1. Proposing a new method that utilizes deep unfolding technology to optimize aggregation weights in federated learning. This method can directly incorporate environmental heterogeneity into the weight adjustment process.
2. Experimentally validating the effectiveness of the proposed method, particularly in handling data with unknown class balance. Compared to traditional heuristic weighting methods, this method achieves higher test accuracy.
3. Discussing the potential application of the proposed method in large-scale learning models, especially when combined with pre-trained models, to perform practical real-world tasks.

In summary, the paper aims to address issues caused by device and data heterogeneity by improving the aggregation weight adjustment mechanism in federated learning, thereby enhancing the overall performance of federated learning.