A Contribution-based Device Selection Scheme in Federated Learning

TL;DR

This paper introduces a device selection scheme for federated learning that balances exploration and exploitation to improve model generalization, personalization, and reduce communication costs.

Contribution

It proposes a novel min-max optimization framework with a primal-dual approach and a modified TMC method for efficient device contribution estimation.

Findings

Lower communication overhead compared to baseline schemes

Improved model generalization and personalization

Competitive performance in experimental evaluations

Abstract

In a Federated Learning (FL) setup, a number of devices contribute to the training of a common model. We present a method for selecting the devices that provide updates in order to achieve improved generalization, fast convergence, and better device-level performance. We formulate a min-max optimization problem and decompose it into a primal-dual setup, where the duality gap is used to quantify the device-level performance. Our strategy combines \emph{exploration} of data freshness through a random device selection with \emph{exploitation} through simplified estimates of device contributions. This improves the performance of the trained model both in terms of generalization and personalization. A modified Truncated Monte-Carlo (TMC) method is applied during the exploitation phase to estimate the device's contribution and lower the communication overhead. The experimental results show that the proposed approach has a competitive performance, with lower communication overhead and competitive personalization performance against the baseline schemes.