Joint Probability Selection and Power Allocation for Federated Learning

TL;DR

This work tackles energy- and time-constrained federated learning over wireless networks by relaxing binary client selection to probabilistic participation, modeled with $a_{ik}\in[0,1]$, and jointly optimizing with transmit power $P_{ik}$. An alternating optimization framework solves two subproblems: (i) power allocation for fixed participation probabilities via a fractional program solved with Dinkelbach's method, yielding $P_{ik}^* = \frac{\lambda B_i}{a_{ik} S \ln(2)} - d_i^2 \sigma^2$ (projected to $[0, P^{\max}]$); (ii) participation probability updates with a closed-form expression ensuring per-round time and energy constraints. The method is validated on a non-iid MNIST CNN task with 100 devices, showing notable improvements in convergence time, energy consumption, and accuracy compared with benchmarks, particularly under highly biased data distributions. The proposed probabilistic scheduling fosters broader participation and fairness while respecting device budgets, offering a practical framework for energy-aware FL in wireless settings.

Abstract

In this paper, we study the performance of federated learning over wireless networks, where devices with a limited energy budget train a machine learning model. The federated learning performance depends on the selection of the clients participating in the learning at each round. Most existing studies suggest deterministic approaches for the client selection, resulting in challenging optimization problems that are usually solved using heuristics, and therefore without guarantees on the quality of the final solution. We formulate a new probabilistic approach to jointly select clients and allocate power optimally so that the expected number of participating clients is maximized. To solve the problem, a new alternating algorithm is proposed, where at each step, the closed-form solutions for user selection probabilities and power allocations are obtained. Our numerical results show that the proposed approach achieves a significant performance in terms of energy consumption, completion time and accuracy as compared to the studied benchmarks.

Figures (2)

• Figure 1: Test accuracy for the highly-biased data scenario.
• Figure 2: Test accuracy for the mildly-biased data scenario.