Exploring Age-of-Information Weighting in Federated Learning Under Data Heterogeneity

This paper investigates wireless federated learning in data heterogeneous scenarios, where device selection usually leads to a degradation in learning performance. This paper is motivated by the fact that while training deep learning networks using federated stochastic gradient descent (FedSGD) on non-independent and identically distributed (non-IID) datasets, device selection can generate gradient errors that accumulate, leading to potential weight divergence, which is further exacerbated with low device participation. To mitigate weight divergence, an age-weighted FedSGD algorithm is designed in this paper to scale local gradients according to the previous device selection results. Furthermore, by revealing the relationship between device participation and latency, an energy consumption minimization problem is formulated accordingly, which consists of resource allocation and sub-channel assignment. By transforming the resource allocation problem into convex and utilizing KKT conditions, we derive the optimal resource allocation solution. Moreover, this paper develops a matching based algorithm to generate the enhanced sub-channel assignment. Simulation results indicate that 1) age-weighted FedSGD is able to outperform conventional FedSGD in terms of convergence rate and achievable accuracy, and 2) the proposed resource allocation and sub-channel assignment strategies can significantly reduce energy consumption and improve learning performance by increasing device participation.