PaperlayerView exact revision on arXiv

Source-linked AI summary

Clustered Sampling: Low-Variance and Improved Representativity for Clients Selection in Federated Learning

Yann Fraboni, Richard Vidal, Laetitia Kameni, Marco Lorenzi

arXiv:2105.05883v2cs.LGcs.AI

TL;DR

FL needs communication-efficient client selection without sacrificing representativity or training stability. The paper introduces clustered sampling with sample-size and model-similarity clustering, proving lower aggregation-weight variance and reporting faster, smoother convergence in heterogeneous data settings. The approach preserves minimal communication and integrates with existing FL privacy and compression technologies.

  • Problem

    Existing FL client-sampling approaches can be biased or exhibit large selection variance, making representativity and convergence stability important unresolved concerns.

  • Method

    Clustered sampling generalizes MD sampling through client clustering based on sample size or representative-gradient similarity.

  • Results

    Clustered sampling provides smaller aggregation-weight variance and better client representativity, with faster and smoother convergence on heterogeneous datasets.

  • Takeaways & Limitations

    Clustered sampling is unbiased, simple to implement, requires no behavior changes on the server or clients, and remains compatible with privacy enhancement and communication reduction methods.

  • Takeaways & Limitations

    The similarity-based algorithm communicates model differences rather than gradients, so its communication cost and privacy properties remain those of standard FL.

Abstract

from arXiv · show

This work addresses the problem of optimizing communications between server and clients in federated learning (FL). Current sampling approaches in FL are either biased, or non optimal in terms of server-clients communications and training stability. To overcome this issue, we introduce \textit{clustered sampling} for clients selection. We prove that clustered sampling leads to better clients representatitivity and to reduced variance of the clients stochastic aggregation weights in FL. Compatibly with our theory, we provide two different clustering approaches enabling clients aggregation based on 1) sample size, and 2) models similarity. Through a series of experiments in non-iid and unbalanced scenarios, we demonstrate that model aggregation through clustered sampling consistently leads to better training convergence and variability when compared to standard sampling approaches. Our approach does not require any additional operation on the clients side, and can be seamlessly integrated in standard FL implementations. Finally, clustered sampling is compatible with existing methods and technologies for privacy enhancement, and for communication reduction through model compression.

1. Introduction

Federated learning must reduce costly server–client communication while selecting clients that represent heterogeneous data distributions. The paper introduces clustered sampling to preserve unbiasedness, reduce selection variability, and improve convergence stability.

  • Communication is a primary FL bottleneck because end-user connections vary in rate and reliability, while server capacity limits simultaneous clients.
  • Client sampling minimizes communication by selecting only a subset of clients, but FedAvg’s uniform scheme is biased relative to aggregating every client.
  • MD sampling is unbiased but can produce large per-iteration selection variance, weakening representativity and increasing convergence variability in non-iid settings.
  • Clustered sampling is introduced as an unbiased scheme that lowers client-selection variability while maintaining minimal server–client communication.
  • The paper develops clustering based on client sample size and model-update similarity, then proves theoretical benefits and evaluates convergence on balanced and unbalanced heterogeneous datasets.

2. Related Work

Prior FL work combines local computation and client sampling to reduce communication, with MD sampling preserving aggregation in expectation under standard optimization assumptions. The related theory establishes convergence guarantees but leaves sampling variability and added-communication trade-offs relevant to subsequent work.

  • FedAvg samples clients uniformly, whereas MD sampling selects clients according to relative data ratios and reproduces full-client aggregation in expectation.
  • MD sampling is unbiased, while FedAvg is prone to client drift because its aggregation does not satisfy that expectation property.
  • MD convergence theory assumes smooth local objectives, unbiased bounded-variance stochastic gradients, and bounded dissimilarity among local objectives.
  • Under these assumptions, prior results show that MD client sampling can converge to a stationary point or local minimum of the federated loss.
  • Other sampling schemes use additional client communication or computation, unlike FedAvg and MD sampling, which minimize per-iteration client–server communication.

3. Clustered sampling

Clustered sampling generalizes MD sampling by using multiple client distributions while preserving unbiased aggregation and convergence guarantees. It reduces aggregation-weight variance and improves client representativity.

  • Definition: Clustered sampling selects clients from m independent distributions that can privilege different client subsets, with MD sampling as a special case.The distributions may vary across iterations and are designed to satisfy unbiasedness conditions.
  • Convergence: Unbiased clustered sampling preserves the same asymptotic convergence behavior and convergence bound as MD sampling under the stated assumptions.The result applies to clustered schemes satisfying Proposition 1.
  • Statistical improvements: Every client has smaller aggregation-weight variance under clustered sampling, with equality only when all distributions equal the MD distribution.The comparison uses the client aggregation weight induced by the sampled subset.
  • Statistical improvements: Clustered sampling gives every client a higher sampling probability and therefore better representativity throughout the federated-learning process.The equality case is precisely the MD configuration, where all clustered distributions are identical.
  • Implications: Reduced aggregation-weight variance and increased representativity may support smoother and faster convergence in heterogeneous federated datasets.The paper identifies clients with unique data distributions as potential beneficiaries of improved sampling representativity.

4. Clustered sampling based on sample size

The sample-size implementation constructs m sampling distributions by allocating client representation according to each client's number of samples. It is unbiased, runs in O(nlog(n)), and can be computed once when sample sizes remain fixed.

  • Construction: The sample-size scheme orders clients by descending sample size and allocates their representation across m distributions.The total sample count is M, and each client's allocation is based on mni.
  • Correctness: Algorithm 1 outputs m distributions satisfying the clustered-sampling conditions of Proposition 1.The construction distributes each client's required representation across sets of cardinality M.
  • Complexity: O(nlog(n)) is the complexity of Algorithm 1.The complexity result is stated in Theorem 3.
  • Implementation: Algorithm 1 may sample a client up to ⌊mpi⌋+ 2 times, while MD sampling may sample clients up to m times.The sample-size clustering algorithm needs to run only once if client sample sizes do not change during learning.

5. Clustered sampling based on similarity

The similarity-based implementation clusters clients using representative gradients, then allocates clustered clients across sampling distributions while preserving the unbiasedness conditions. It has a stated complexity of O(n2d+ X).

  • Similarity construction: Similarity-based clustered sampling uses representative gradients, defined as differences between clients' updated and global models, to identify similar clients.Comparing representative gradients is described as effective for detecting similarity among federated-learning participants.
  • Correctness and complexity: Algorithm 2 outputs m distributions satisfying Proposition 1 when every client has pi ≤1/m.Its complexity is expressed in terms of client count n, model parameters d, and clustering-method complexity X.
  • Statistical behavior: The similarity-based scheme bounds client aggregation-weight variance and explicitly promotes sampling clients according to similarity.The distributions are obtained from a similarity tree built from representative gradients.
  • Initialization: Algorithm 2 can operate without representative gradients for some clients by assigning them a constant zero representative gradient and grouping them together.This applies when clients have not yet been sampled during federated learning.
  • Communication and privacy: The similarity-based method requires only differences between local and global models, so its communication cost matches standard federated learning and its privacy properties remain identical.The paper also permits any valid hierarchical clustering algorithm; Ward clustering is one example with complexity O(n2 log(n)).

6. Experiments

Experiments on controlled MNIST and heterogeneous CIFAR10 settings show that clustered sampling improves convergence, representativity, and variability relative to MD sampling, with benefits increasing under stronger heterogeneity.

  • MNIST: On MNIST, Algorithms 1 and 2 outperform MD sampling in training global loss, testing accuracy, and sampled-class representativity.Algorithm 2 converges to the ideal performance of target sampling.
  • MNIST: MD sampling selects fewer than 10 distinct clients in 37% of FL iterations, whereas clustered sampling always aggregates 10 distinct clients in this balanced setting.The probability of sampling 10 different clients with MD sampling is approximately 63%.
  • MNIST: Clustered sampling yields lower convergence variance and better training-loss and testing-accuracy performance than MD sampling.Algorithm 2 quickly partitions clients into 10 clusters and approaches the ideal target distribution.
  • CIFAR10: For CIFAR10 partitions, clustered sampling improves over MD sampling across Dirichlet heterogeneity levels α ∈ {0.001, 0.01, 0.1, 10}.The experiments use N = 100 and m = 10.
  • CIFAR10: The improvement over MD sampling is larger for more heterogeneous CIFAR10 datasets, with faster and better convergence at α = 0.01 and α = 0.001.At α = 10 and α = 0.1, final performances are close, although clustered sampling converges faster.
  • Additional analyses: Additional experiments report similar improvements with L2 and L1 similarity measures and show that larger local work N produces models fitting clients’ data distributions better.Figure 2 reports rolling-mean training losses over 50 iterations.

7. Discussion and conclusion

The paper concludes that clustered sampling generalizes MD sampling while reducing aggregation-weight variance and improving representativity, convergence, and implementation compatibility.

  • Contributions: Clustered sampling generalizes MD sampling and is implemented using client sample size or model similarity.Both algorithms provide smaller client aggregation-weight variance and better client representativity.
  • Empirical conclusion: Clustered sampling experimentally produces faster and smoother convergence on heterogeneous datasets.
  • Future directions: The framework supports further clustering criteria beyond client sample size and model similarity.The paper identifies model-similarity detection during client sampling as novel relative to approaches using all clients each iteration.
  • Practical scope: Clustered sampling is unbiased, simple to implement, requires no changes to server or client behavior, and remains compatible with privacy and communication-reduction technologies.

A. Proof of Theorem 2

The proof establishes that clustered sampling satisfies the conditions used for FedAvg convergence guarantees, extending the corresponding argument for MD sampling.

  • Proof strategy: The proof reproduces the MD-sampling argument showing that its sampling scheme satisfies Lemma 5 under Assumptions 1 to 3.
  • MD sampling: MD sampling is represented through independently sampled client indices drawn from a multinomial distribution with probabilities p_i.
  • Bounding steps: The derivation bounds the proof terms using the sampling identities, Jensen’s inequality, unbiasedness, and Assumption 3.
  • Conclusion: Substituting the derived bounds into the target equation completes the proof.

A.2. Proof of Lemma 5 for Theorem 2

For clustered sampling, the proof uses independent but non-identically distributed client draws, establishes the required identities and bounds, and compares its convergence term with MD sampling.

  • Clustered sampling: Clustered sampling draws client indices independently from cluster-specific distributions W_k, with client probabilities r_k,i.
  • Sampling identity: Clustered sampling satisfies the sampling identity required for the convergence analysis through Proposition 1.
  • Proof bounds: The proof applies the general decomposition and bounds its terms using the same Jensen-based steps as the MD-sampling proof.
  • Gradient terms: The expected gradient of each cluster distribution W_k is defined with respect to model parameters x for the second-term bound.
  • Comparison: Equations (32) and (40) identify the theoretical convergence improvement of clustered sampling over MD sampling.
  • Comparison: The client-gradient bound is no larger for clustered sampling than for MD sampling, with equality only when cluster gradients match or clustered sampling reduces to MD sampling.

B.2. Probability for a client to be sampled at least once

The paper develops two client-allocation algorithms for clustered sampling, using sample counts or model-similarity groups to construct distributions with controlled client participation.

  • Notation: The algorithms can be formulated using data ratios p_i, but the paper works with integer sample quantities n_i = p_iM.The allocated quantity r′_k,i denotes the number of samples from client i assigned to distribution k.
  • Algorithm 1: Algorithm 1 provides a sampling scheme satisfying Proposition 1 with overall complexity O(n log(n)).Client allocation is linear after sorting the n clients.
  • Algorithm 1: Algorithm 1 allocates clients in decreasing sample-count order across distributions, leaving all but at most one distribution with 0 or M samples after each allocation.A client’s samples are assigned to distributions that are not yet full.
  • Algorithm 2: Algorithm 2 first splits clients into K ordered groups, assigns the first m groups to distributions, and splits remaining groups among available slots.Each resulting distribution has M samples from participating clients.
  • Algorithm 2: Algorithm 2 satisfies Proposition 1 and has time complexity O(n^2d + X), including pairwise model-similarity computation and clustering-tree construction.In practice, estimating similarities only for sampled clients reduces the stated complexity to O(nmd + X).

D. Additional experiments

Additional experiments vary CIFAR10 heterogeneity and similarity measures, using Dirichlet-partitioned federated datasets and global-loss convergence evaluations.

  • Experimental setup: The experiments partition CIFAR10 with a Dirichlet distribution and select learning rates from {0.001, 0.005, 0.01, 0.05, 0.1} to minimize FedAvg with MD sampling loss.All experiments use batch size 50.
  • CIFAR10 partitioning: α = 10 produces an almost iid partition, whereas α = 0.001 yields highly heterogeneous clients often owning nearly one class each.The α = 0.001 partition is also unbalanced across classes.
  • Similarity measures: Algorithm 1 outperforms MD sampling by reducing client aggregation-weight variance, while the three similarity measures produce similar Algorithm 2 performance with Ward clustering.This supports using Arccos similarity in the remaining experiments.

D.4. Influence of m the number of sampled clients, and N the number of SGD run

The experiments vary local computation and the number of sampled clients to assess clustered sampling’s convergence speed and smoothness under heterogeneous CIFAR10 settings.

  • Influence of N: More local work N is associated with faster clustered-sampling convergence.The paper relates greater local work to clients fitting their own data more closely.
  • Unbalanced datasets: Figure 7 evaluates clustered sampling on unbalanced CIFAR10 datasets spanning α ∈ {0.001, 0.01, 0.1, 10}, with N = 100 and m = 10.Learning rates are {0.05, 0.05, 0.05, 0.1} for the four respective datasets.
  • Experimental setup: The experiments vary N, the number of client SGD runs, and m, the number of sampled clients, under a CIFAR10 partition with Dir(α = 0.01).Figures 8 and 9 use multiple learning-rate settings and vary N and m across rows.
  • Influence of m: With fewer sampled clients m, clustered sampling improves more over MD sampling; with more clients, both methods converge faster and their performance gap narrows but remains favorable to clustered sampling.The paper associates the stronger low-m improvement with better data representativity.
  • FedProx: With FedProx regularization set to µ = 0.1, Algorithms 1 and 2 still outperform MD sampling.The regularization parameter was selected from {0.001, 0.01, 0.1, 1.}.
Loading 2105.05883v2…