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

TL;DR

This paper introduces an age-weighted FedSGD method for federated learning that mitigates gradient errors caused by non-IID data and device selection, while also optimizing resource allocation to enhance learning efficiency and reduce energy use.

Contribution

It proposes an age-weighted gradient scaling approach and an energy-efficient resource allocation framework for federated learning under data heterogeneity and device constraints.

Findings

Age-weighted FedSGD improves convergence and accuracy.

Optimized resource allocation reduces energy consumption.

Enhanced device participation boosts learning performance.

Abstract

This paper investigates federated learning in a wireless communication system, where random device selection is employed with non-independent and identically distributed (non-IID) data. The analysis indicates that while training deep learning networks using federated stochastic gradient descent (FedSGD) on non-IID datasets, device selection can generate gradient errors that accumulate, leading to potential weight divergence. To mitigate training divergence, we design an age-weighted FedSGD to scale local gradients according to the previous state of devices. To further improve learning performance by increasing device participation under the maximum time consumption constraint, we formulate an energy consumption minimization problem by including resource allocation and sub-channel assignment. By transforming the resource allocation problem into convex and utilizing KKT conditions, we derived the optimal resource allocation solution. Moreover, this paper develops a matching based algorithm to generate the enhanced sub-channel assignment. Simulation results indicate that i) age-weighted FedSGD is able to outperform conventional FedSGD in terms of convergence rate and achievable accuracy, and ii) the proposed resource allocation and sub-channel assignment strategies can significantly reduce energy consumption and improve learning performance by increasing the number of selected devices.