Multi-task Federated Edge Learning (MtFEEL) in Wireless Networks

TL;DR

This paper introduces a novel federated learning algorithm for wireless networks that personalizes models for edge devices with convergence guarantees, improving performance over traditional methods.

Contribution

It proposes a communication-efficient, personalized federated learning algorithm with theoretical convergence guarantees in wireless settings.

Findings

Outperforms FedAvg and FedSGD in practical SNR regimes.

Converges at a rate of max{1/SNR, 1/√T} in Rayleigh fading channels.

Provides a PAC-based bound on empirical loss incorporating Rademacher complexity and discrepancy.

Abstract

Federated Learning (FL) has evolved as a promising technique to handle distributed machine learning across edge devices. A single neural network (NN) that optimises a global objective is generally learned in most work in FL, which could be suboptimal for edge devices. Although works finding a NN personalised for edge device specific tasks exist, they lack generalisation and/or convergence guarantees. In this paper, a novel communication efficient FL algorithm for personalised learning in a wireless setting with guarantees is presented. The algorithm relies on finding a ``better`` empirical estimate of losses at each device, using a weighted average of the losses across different devices. It is devised from a Probably Approximately Correct (PAC) bound on the true loss in terms of the proposed empirical loss and is bounded by (i) the Rademacher complexity, (ii) the discrepancy, (iii) and a penalty term. Using a signed gradient feedback to find a personalised NN at each device, it is also proven to converge in a Rayleigh flat fading (in the uplink) channel, at a rate of the order max{1/SNR,1/sqrt(T)} Experimental results show that the proposed algorithm outperforms locally trained devices as well as the conventionally used FedAvg and FedSGD algorithms under practical SNR regimes.