Balancing Client Participation in Federated Learning Using AoI

TL;DR

This work tackles the challenge of balanced client participation and convergence efficiency in federated learning under limited communication resources. It introduces an AoI-based decentralized Markov scheduling policy that governs age-dependent client selection probabilities, aiming to minimize the variance of inter-selection times and thus balance updates across clients. A rigorous convergence analysis links the algorithm’s performance to the aggregation-weight variance $\Sigma$ and selection skew bounds $\underline{\rho}, \overline{\rho}$, showing that AoI-informed scheduling accelerates convergence. Extensive simulations on MNIST, CIFAR-10, and CIFAR-100 demonstrate that the optimal Markov variant consistently outperforms FedAvg and other baselines in both IID and non-IID settings, with improvements ranging from about 7.5% to 20%, underscoring the practical impact of AoI-based scheduling for scalable, fair, and efficient FL systems.

Abstract

Federated Learning (FL) offers a decentralized framework that preserves data privacy while enabling collaborative model training across distributed clients. However, FL faces significant challenges due to limited communication resources, statistical heterogeneity, and the need for balanced client participation. This paper proposes an Age of Information (AoI)-based client selection policy that addresses these challenges by minimizing load imbalance through controlled selection intervals. Our method employs a decentralized Markov scheduling policy, allowing clients to independently manage participation based on age-dependent selection probabilities, which balances client updates across training rounds with minimal central oversight. We provide a convergence proof for our method, demonstrating that it ensures stable and efficient model convergence. Specifically, we derive optimal parameters for the Markov selection model to achieve balanced and consistent client participation, highlighting the benefits of AoI in enhancing convergence stability. Through extensive simulations, we demonstrate that our AoI-based method, particularly the optimal Markov variant, improves convergence over the FedAvg selection approach across both IID and non-IID data settings by $7.5\%$ and up to $20\%$. Our findings underscore the effectiveness of AoI-based scheduling for scalable, fair, and efficient FL systems across diverse learning environments.

Figures (7)

• Figure 1: Markov chain with $m'+1$ states.
• Figure 2: Total Variance $\Sigma$ over $1000$ rounds. The simulation was performed with $n = 100$ clients, $k = 15$ selected clients per round, and a maximum client age of $m = 10$. The figure compares : Policy 1 (Random Weighted), Policy $2$ (Probabilistic), Policy $3$ (Markov Optimal), and Policy $4$ (Markov Non-Optimal).
• Figure 3: Empirical inter-selection interval distributions for four client selection policies over $1000$ rounds, with $100$ clients and $15$ clients selected per round.
• Figure 4: Normalized variance $\frac{\sqrt{\text{Var}(Y)}}{T}$ versus window size $T$ for four client selection policies: Random selection (Policy 1), Probabilistic (Policy $2$), Markov Optimal (Policy $3$), and Markov Non-Optimal (Policy $4$). The setup involves $100$ clients, $1000$ rounds, and the selection of $15$ clients per round.
• Figure 5: Accuracy on MNIST dataset across different client selection policies: Random selection (Policy $1$), Probabilistic (Policy $2$), Markov Optimal (Policy $3$), and Markov Non-Optimal (Policy $4$). The experiment was conducted with $100$ clients, $15$ clients selected per round, and a maximum client age of $10$. (a) IID dataset. (b) Non-IID dataset (generated using a Dirichlet distribution with $\alpha = 0.3$).

Theorems & Definitions (18)

• Theorem 1
• Definition 1: Local-Global Objective Gap and Selection Skew
• Theorem 2
• Theorem 3
• Theorem 4
• Theorem 5
• Lemma 1
• proof
• Lemma 2
• proof