Source-linked AI summary
Node Selection Toward Faster Convergence for Federated Learning on Non-IID Data
Hongda Wu, Ping Wang
TL;DR
Non-i.i.d. data creates discrepancies between global and local objectives, slowing federated-learning convergence. The paper combines Optimal Aggregation, which excludes adverse local updates, with FedPNS, which adjusts node-selection probabilities; experiments show faster convergence than randomly selected FedAvg.
Problem
Non-i.i.d. data creates discrepancies between global and local objectives, causing federated-learning models to converge slowly.
Method
Optimal Aggregation selects a subset of participating local updates using local–global gradient relationships, while FedPNS dynamically adjusts node-selection probabilities from that result.
Results
FedPNS converges faster and achieves higher test accuracy than FedAvg across synthetic and real-data settings, with stronger gains under higher data heterogeneity.
Takeaways & Limitations
Selecting nodes according to their contribution can accelerate federated-learning convergence on heterogeneous data.
Abstract
from arXiv · showhide
Federated Learning (FL) is a distributed learning paradigm that enables a large number of resource-limited nodes to collaboratively train a model without data sharing. The non-independent-and-identically-distributed (non-i.i.d.) data samples invoke discrepancies between the global and local objectives, making the FL model slow to converge. In this paper, we proposed Optimal Aggregation algorithm for better aggregation, which finds out the optimal subset of local updates of participating nodes in each global round, by identifying and excluding the adverse local updates via checking the relationship between the local gradient and the global gradient. Then, we proposed a Probabilistic Node Selection framework (FedPNS) to dynamically change the probability for each node to be selected based on the output of Optimal Aggregation. FedPNS can preferentially select nodes that propel faster model convergence. The unbiasedness of the proposed FedPNS design is illustrated and the convergence rate improvement of FedPNS over the commonly adopted Federated Averaging (FedAvg) algorithm is analyzed theoretically. Experimental results demonstrate the effectiveness of FedPNS in accelerating the FL convergence rate, as compared to FedAvg with random node selection.
I. INTRODUCTION
Federated learning keeps training data local but converges slowly when randomly selected nodes hold non-i.i.d. data. The paper proposes Optimal Aggregation and FedPNS to identify adverse updates and preferentially select nodes that improve convergence.
- Non-i.i.d. data produces biased local updates, stagnates convergence, reduces accuracy, and increases communication rounds for resource-constrained edge nodes.
- The paper theoretically analyzes FedAvg’s convergence bound and the expected decrease of global loss under data heterogeneity and local-update aggregation.
- Optimal Aggregation identifies and excludes adverse local updates instead of necessarily aggregating updates from all participating nodes.
- FedPNS dynamically adjusts each node’s selection probability using Optimal Aggregation and preferentially selects nodes that boost model convergence.
- Experiments on synthetic and real datasets with different learning objectives show that FedPNS improves convergence rate compared with FedAvg.
- Prior node-selection methods are contrasted because they had not analyzed data heterogeneity’s impact or the alignment between local and global gradients.
III. PRELIMINARIES
The paper formulates federated learning as minimizing a global empirical risk assembled from private local datasets. FedAvg uses partial participation, local SGD, and server aggregation across communication rounds.
- A central server coordinates multiple nodes, each holding a private dataset, to collaboratively minimize a global model objective.
- The global objective is represented through local objectives that may differ because nodes can have different data distributions.
- FedAvg selects a fraction of nodes in each communication round, and participating nodes perform local stochastic gradient descent.
- The local update consists of mini-batch SGD performed for τ local updates, determined by dataset size, batch size, and local epochs.
- Each selected node’s local training result is communicated back to the server as a model update for aggregation.
C. The Challenges of Non-i.i.d. Data Distribution
Non-i.i.d. data and partial node participation slow FL convergence because local objectives and updates diverge from the global objective. The analysis relates expected global-loss decrease to data heterogeneity and aggregation, motivating contribution-aware node selection.
- C. The Challenges of Non-i.i.d. Data Distribution: Partial participation and non-i.i.d. training data slow FedAvg convergence despite its simple random-selection policy.Increasing local computation can reduce communication rounds, but does not resolve unsatisfactory performance on non-i.i.d. data.
- C. The Challenges of Non-i.i.d. Data Distribution: Randomly selected nodes optimize user-specific local objectives, whose updates can deviate from the global objective when data distributions differ.The discrepancy arises because each local loss is tied to its node’s data distribution.
- C. The Challenges of Non-i.i.d. Data Distribution: The paper uses local-update and global-update inner products to identify nodes whose updates adversely contribute to global convergence.This criterion implicitly profiles differences between node data distributions and the population distribution.
- C. The Challenges of Non-i.i.d. Data Distribution: Under convexity, Lipschitz smoothness, local dissimilarity, and bounded-gradient assumptions, heterogeneity is represented by δ and node gradients by γ_i.Greater data heterogeneity increases δ, while γ_i reflects the relationship between a node’s distribution and the population distribution.
- C. The Challenges of Non-i.i.d. Data Distribution: Lemma 1 bounds expected FedAvg global-loss decrease as a function of local-distribution variance and the aggregation strategy.The bound applies when the current model is not stationary and uses the stochastic gradient convention stated in the paper.
B. Aggregation with Gradient Information
The paper challenges averaging every participating update and proposes Optimal Aggregation to remove updates that harm expected loss reduction. It iteratively evaluates candidate exclusions using gradient information and a loss check.
- B. Aggregation with Gradient Information: The server checks the expected inner product between the global gradient and local gradients to identify potentially adverse updates.Updates are considered for exclusion when the expectation improves after removing them.
- B. Aggregation with Gradient Information: Each candidate exclusion is validated by comparing test losses before the update is removed from the aggregation set.This second step accounts for the effect of exclusion on the global update and data heterogeneity.
- B. Aggregation with Gradient Information: Optimal Aggregation seeks a subset of participating local updates whose exclusion produces a larger expected decrement of global loss.The method challenges the necessity of aggregating all participating updates.
- B. Aggregation with Gradient Information: The iterative procedure removes labeled updates only when the loss criterion is satisfied, stopping when no adverse update remains or the set reaches threshold v.Otherwise, the candidate update remains retained.
- B. Aggregation with Gradient Information: Optimal Aggregation both enlarges expected per-round loss reduction and identifies potential adverse nodes for subsequent probabilistic selection.The identified nodes are typically associated with non-i.i.d. datasets.
C. FL with Probabilistic Node Selection (FedPNS)
FedPNS dynamically adjusts each node’s selection probability using Optimal Aggregation’s identification of adverse updates. It decreases probabilities for repeatedly adverse nodes and increases probabilities for the remaining nodes.
- C. FL with Probabilistic Node Selection (FedPNS): FedPNS dynamically changes node-selection probabilities from the contribution information produced by Optimal Aggregation.The design preferentially selects nodes associated with faster model convergence.
- C. FL with Probabilistic Node Selection (FedPNS): The ratio x measures how often a node has been labeled by CHECK EXPECTATION relative to how often it has been selected.The coefficients α and β control the probability-decrement function.
- C. FL with Probabilistic Node Selection (FedPNS): The algorithm initializes equal node-selection probabilities, samples participating nodes according to the current probabilities, and updates them after aggregation.The procedure repeats across T global rounds.
- C. FL with Probabilistic Node Selection (FedPNS): The empirically selected parameters α=2 and β=0.7 balance aggressive probability changes against weak differentiation between nodes.Large α increases sensitivity to identification mistakes, whereas small α can make probability changes marginal.
- C. FL with Probabilistic Node Selection (FedPNS): Nodes labeled by CHECK EXPECTATION receive probability decreases based on x, while all remaining nodes receive equal probability increases.The adjustment is applied for subsequent global rounds.
D. Convergence Rate of FedPNS
The convergence analysis compares FedPNS with FedAvg through the weight divergence between the federated sequence and a centralized virtual sequence. Under the stated assumptions, smaller FedPNS divergence yields a tighter final global-loss bound.
- D. Convergence Rate of FedPNS: The analysis introduces a virtual centralized sequence v_t optimized for the global loss and synchronized with the full-participation model.This sequence provides the reference for measuring federated weight divergence.
- D. Convergence Rate of FedPNS: Theorem 1 bounds FedPNS weight divergence after synchronization for equal-sized nodes, fixed local updates, and specified data-distribution assumptions.The theorem considers K local nodes whose samples follow node-specific distributions q^(i).
- D. Convergence Rate of FedPNS: Weight divergence arises from bounded participating-node gradients and differences between node distributions and the population distribution.These are the two components identified in the analysis.
- D. Convergence Rate of FedPNS: FedPNS preferentially selects nodes with smaller bounded gradients, producing smaller weight divergence than equal-probability node selection in FedAvg.The analysis connects this smaller divergence to the convergence comparison.
- D. Convergence Rate of FedPNS: A smaller FedPNS weight divergence gives a tighter upper bound on F(w_T) − F(w*) after T global rounds than FedAvg.The proof uses the proportionality between the final global-loss gap and weight divergence.
E. Complexity Analysis
The proposed server-side aggregation procedures have quadratic complexity in the model width, while local backpropagation is cubic, making the server overhead marginal under the stated assumptions.
- The FCNN forward-propagation complexity is O(n^2) because n≫n_layer.The model has n·n_layer parameters, with n≫n_layer typically.
- CHECK LOSS has O(n^2) complexity, while local-gradient addition and multiplication require O(n) and O(n^2), respectively.
- CHECK EXPECTATION also has O(n^2) complexity because participating updates and iterations are much smaller than n.
- Algorithm 2 has Θ(2n^2 + n) complexity, whereas local backpropagation has O(n^3) complexity.The authors characterize the proposed server-side overhead as marginal and ignorable relative to local training.
V. EVALUATION AND ANALYSIS
The evaluation tests FedPNS across synthetic and real datasets, multiple learning objectives and models, using FedAvg as the benchmark and testing accuracy as the usual evaluation measure.
- The experiments implement FedPNS across different tasks, models, and datasets, comparing it with commonly used FedAvg.The evaluation first examines Optimal Aggregation, then FedPNS under different data heterogeneity.
- Synthetic data is generated to characterize data heterogeneity and study its impact on model convergence.
- Real-data experiments cover convex MNIST classification with MLR and non-convex MNIST and CIFAR-10 tasks with CNN models.
- Unless otherwise specified, trained-model accuracy is evaluated on each dataset's testing set.The standard configuration selects 20% of nodes, or |S_t| = 10, and runs for T = 200 rounds.
A. Performance of Optimal Aggregation
Optimal Aggregation lowers training loss relative to FedAvg by excluding adverse local updates, and its node-level behavior is examined under i.i.d. and non-i.i.d. participation.
- Optimal Aggregation achieves lower training loss than FedAvg in the MNIST experiment.The experiment uses CNN-M on MNIST with σ = 0.5, ρ = 1 and keeps the same participating nodes for both methods.
- During early rounds, greater local-update diversity makes excluding adverse local updates more effective.
- i.i.d. nodes are never excluded, while some non-i.i.d. nodes are excluded repeatedly by CHECK LOSS.
- Nearly all non-i.i.d. nodes are selected at least once, indicating that Optimal Aggregation identifies potentially adverse updates among participating nodes.
- As data heterogeneity increases, FedAvg converges more slowly with higher training loss and decreasing test accuracy, while FedPNS has lower loss and higher accuracy in all tested settings.
2) MLR, CNN-M Model for MNIST:
Across MNIST and CIFAR-10 experiments, FedPNS generally converges faster and reaches higher test accuracy than FedAvg, with gains shaped by heterogeneity and parameter choices.
- 2) MLR, CNN-M Model for MNIST:: FedPNS converges faster and achieves higher test accuracy than FedAvg for both MLR and CNN models across data heterogeneity settings.
- 2) MLR, CNN-M Model for MNIST:: Performance gains become larger as data heterogeneity increases, while the gap is not obvious at σ = 0.5, ρ = 2.The smaller gap is attributed to the reduced impact of non-i.i.d. nodes when many i.i.d. nodes can be selected.
- 3) CNN-C Model for CIFAR-10:: On CIFAR-10 with CNN-C, FedPNS converges faster and reaches higher test accuracy than FedAvg, especially at σ = 0.2 and 0.3, ρ = 1.Improvement is not obvious at σ = 0.2, ρ = 2 because FedPNS has difficulty distinguishing node contributions.
- C. Choosing α and β: Increasing α or β improves performance only within a range; α = 2 and β = 0.7 are selected empirically.Larger α increases sensitivity to identification mistakes, while β = 0.8 degrades performance in the reported setting.
- D. Other Comparison: Selecting nodes by higher gradient norm slows convergence because i.i.d. nodes have smaller averaged gradient norms than non-i.i.d. nodes.This comparison is made between FedPNS and BN2 on CNN-M for MNIST with σ = 0.5, ρ = 1.
- VI. CONCLUSION: Across different tasks, models, and datasets, FedPNS accelerates convergence and produces higher test accuracy than FedAvg.
APPENDIX
The appendix develops bounds for the global objective change by controlling the model-step norm and its alignment with the global gradient, then combines the resulting inequalities.
- APPENDIX: Taylor expansion of the L-smooth global objective provides the starting point for the convergence bound.The proof invokes smoothness of F(w) before applying the subsequent bounds.
- APPENDIX: The proof separately bounds the update norm ∥w_t+1 − w_t∥2 using global aggregation and local-update definitions.The stated inequalities use Cauchy-Schwarz and the bounded dissimilarity assumption.
- APPENDIX: The proof separately bounds the inner product ⟨∇F(w_t), w_t+1 − w_t⟩ using global aggregation and SGD optimization.This bound is paired with the update-norm bound in the smoothness argument.
- APPENDIX: Substituting the two auxiliary bounds into the initial smoothness inequality yields the appendix result.The substitution is explicitly described after equations (A2) and (A3).
B. Proof of Theorem 1
Theorem 1 is proved by bounding divergence between partial-participation, synchronized, and centralized models, then substituting these bounds into the theorem’s initial inequality.
- B. Proof of Theorem 1: The proof bounds the divergence between the partial-participation model w_t and centralized model v_t at any global round.This establishes the main divergence quantity used in the theorem analysis.
- B. Proof of Theorem 1: The synchronized and centralized models are indexed by local-update count τ because each global round applies τ local SGD steps.This re-indexing facilitates the divergence analysis.
- B. Proof of Theorem 1: The bound on ∥w̃_t − v_t∥ uses equal-data-size updates, triangle inequalities, class-distribution relations, and the L-smoothness assumption.The proof introduces local-update indices and bounds intermediate terms using the stated distributional and smoothness assumptions.
- B. Proof of Theorem 1: Mathematical induction with a = 1 + ηL supplies the recursive bound, after which substitutions complete the proof of Theorem 1.The proof substitutes (B9) and (B6) into (B1), with Assumption 3 used in the intermediate bound.
- B. Proof of Theorem 1: The synchronized-model divergence is represented between consecutive global rounds after omitting the previous-round accumulation because v_t is synchronized with w̃_t at round start.The resulting expression is then combined with the preceding divergence bounds.
- B. Proof of Theorem 1: The proof bounds ∥w_t − w̃_t∥ by taking expectation over nodes sampled from the identical distribution {p_1, p_2, · · ·, p_|K|}.The sampled-node set is denoted S_t, and the global model is expressed through the selected nodes’ updates.