On Large-Cohort Training for Federated Learning

1 Introduction

• Larger cohorts are expected to improve convergence by sampling clients more representatively, and may improve privacy by reducing communication rounds.Distributed differential privacy can also benefit when noise is divided among more clients.Source: §1 · pp. 1–2
• The study finds that larger cohorts may yield little practical convergence improvement and introduce optimization and generalization issues, including challenges unique to federated learning.It provides partial solutions informed by federated training dynamics, intended as simple benchmarks for future work.Source: §1 · p. 2

1.1 Related Work

• Large-batch training also exhibits diminishing speedup returns and generalization gaps, while this work examines cohort-size effects and identifies additional challenges from the federated setting.The cited large-batch findings include diminishing returns and generalization issues in deep neural networks.Source: §1.1 · p. 2
• Prior federated convergence analyses often predict cohort-size gains asymptotically, but rely on learning-rate and heterogeneity assumptions that may not hold in practice.The paper therefore examines cohort-size effects in practical, communication-limited settings.Source: §1.1 · pp. 2–3
• Unlike cohort-selection methods that require server control, this work studies the size of randomly sampled cohorts, reflecting client availability constraints in cross-device FL.Prior selection methods can use client quantities such as loss to choose participants.Source: §1.1 · p. 3

2 Preliminaries

• FedSGD communicates client gradients for every update, whereas FedAvg uses multiple local-training epochs to potentially reduce communication rounds.FedSGD is the mini-batch SGD approach described for solving the weighted client-loss objective.Source: §2 · p. 3
• FedOpt sends the server model to a cohort, averages client updates into a pseudo-gradient, and applies a first-order server optimizer that need not use an unbiased gradient estimate.FedAvg, FedAvgM, FedAdagrad, and FedAdam are special cases determined by the server optimizer.Source: §2 · p. 3

2.1 Experimental Setup

• The evaluation varies cohort size across four datasets with natural or specified client partitions, using models for image, character, and language tasks.Table 1 reports client counts and training and test example counts; models include a ResNet-18, convolutional network, and two RNNs.Source: §2.1 · pp. 3–4, Table 1 · p. 4
• The experiments include FedSGD, FedAvg, FedAvgM, FedAdagrad, and FedAdam, plus newly developed FedLARS and FedLamb variants.These methods are special cases of the FedOpt framework.Source: §2.1 · p. 4
• Results show representative experiments from five random trials, with mean curves and one-standard-deviation bands; learning rates are tuned on held-out validation data.The paper says the full experimental results are in Appendix B.Source: §2.1 · pp. 4–5

3.1 Catastrophic Training Failures

• On EMNIST, FedAvg failures rose from 0% at M = 10 to 80% at M = 800; adaptive clipping eliminated failures across tested cohort sizes.Source: §3.1 · p. 5, Figure 1 · p. 5

3.2 Diminishing Returns

• Convergence gains from larger cohorts saturate across tested optimizers, although larger cohorts can still help in communication-limited settings.The typical saturation threshold is M = 10–50; M = 50 represents 10% of CIFAR-10 clients but approximately 0.015% of Stack Overflow clients.Source: §3.2 · pp. 5–6, Figure 2 · p. 6, Figure 3 · p. 6

3.3 Generalization Failures

• Generalization effects depend on task and optimizer: larger FedAdam cohorts worsen CIFAR-100 generalization despite similar training accuracy, while EMNIST and Stack Overflow show no such issue.Source: §3.3 · p. 7, Figure 5 · p. 7

3.4 Fairness Concerns

• Cohort-size effects span client-accuracy percentiles: CIFAR-100 favors M = 50 over larger cohorts, while Stack Overflow sees slight gains as M grows; percentile spread changes little beyond a point.Source: §3.4 · pp. 7–8, Figure 6 · p. 8

3.5 Decreased Data Efficiency

• Larger cohorts generally reduce FedAdam rounds to accuracy thresholds but require significantly more examples per unit-accuracy, and stragglers can further increase compute time.Source: §3.5 · p. 8, Figure 7 · p. 8

• For FedAvg, pseudo-gradient norms decrease approximately with the inverse square root of cohort size, a pattern that predicts norms after a few rounds.The comparison is shown for Stack Overflow; similar results are reported for other tasks and optimizers.Source: §4 · p. 9, Figure 8 · p. 9
• FedAvg client updates are nearly orthogonal on average, so averaging them produces smaller pseudo-gradients as cohort size increases.FedSGD's pseudo-gradient norm instead decreases only slightly with cohort size and has high variance.Source: §4 · p. 9
• Smaller updates help explain diminishing returns and why each client's examples contribute little in large cohorts; adaptive server optimizers partly compensate through scale-invariant normalization.Source: §4 · p. 9

5 Designing Better Methods

• The section evaluates simple baselines for improving large-cohort training, intended as a reference for future work rather than a complete solution.Source: §5 · p. 9

5.1 Learning Rate Scaling

• Linear server-rate scaling fails catastrophically beyond M = 100; square-root scaling avoids failure but yields task-dependent gains and can hurt the largest cohorts.For EMNIST with FedAvg, linear scaling beyond M = 100 caused failures, while M = 100 brought no real accuracy improvement.Source: §5.1 · p. 10, Figure 9 · p. 10

5.2 Layer-wise Adaptivity

• FedLamb often matches FedAdam at large cohorts but eventually loses test accuracy for M > 100; FedLARS performs well on CIFAR-10 but poorly on EMNIST and Shakespeare.Layer-wise adaptive methods are not universally better than coordinate-wise adaptive methods.Source: §5.2 · pp. 10–11, Figure 10 · p. 11

5.3 Dynamic Cohort Sizes

• Dynamic cohorts approach M = 50 data efficiency while reaching final accuracy closer to a large fixed cohort, but do not outperform M = 800 and retain generalization issues.Source: §5.3 · p. 11, Figure 11 · p. 11

5.4 Normalized FedAvg

• The proposed normalized FedAvg applies SGD to the normalized server pseudo-gradient and adds no hyperparameters relative to Algorithm 1.The experiments reuse learning rates tuned for unnormalized FedAvg.Source: §5.4 · p. 12
• Normalized FedAvg improves training accuracy for larger cohorts on most tasks, with exceptions on EMNIST and the largest Stack Overflow cohort.The authors report significant improvements on CIFAR-100 and all but the largest Stack Overflow cohort sizes.Source: §5.4 · p. 12, Figure 12 · p. 12

5.5 Changing the Number of Local Steps

• With cohort size fixed, more local steps can speed convergence only up to a point before failures, and can reduce generalization more than increasing cohort size.Source: §5.5 · p. 13, Figure 13 · p. 13

• The paper links reduced data efficiency to potentially higher energy consumption and carbon emissions, but says these impacts warrant more specialized analysis.Source: §6 · p. 13
• Future work includes studying differential-privacy noise and personalization for generalization and fairness, and whether findings extend to cross-silo settings with client subsampling.The authors also call for broader connections to other federated-learning topics and large-batch training.Source: §6 · p. 13

• CIFAR-100 has 32 × 32 × 3 images with 100 labels, partitioned among 500 training and 100 test clients with heterogeneous label distributions.The experiments use a ResNet-18 with group normalization.Source: §A.1 · p. 19
• EMNIST contains handwritten alphanumeric images with 62 characters, partitioned by author among 3,400 clients with natural writing-style heterogeneity.The experiments use a convolutional network.Source: §A.1 · p. 19
• Shakespeare assigns clients to play roles with at least two dialogue lines; 715 clients contribute sequences for next-character prediction.The model uses an RNN with two LSTM layers.Source: §A.1 · p. 19
• Stack Overflow is used for next-word prediction, with train clients restricted to pre-2018 examples and test clients to post-2018 examples.The split has 342,477 train clients, 38,758 validation clients, and 204,088 test clients.Source: §A.1 · p. 20

A.2 Implementation and Hyperparameters

• FedAvg, FedSGD, FedAvgM, FedAdam, and FedAdagrad are implemented alongside the novel FedLARS and FedLamb methods.All are special cases of Algorithm 1.Source: §A.2 · p. 20
• For methods other than FedSGD, algorithms differ only in their server optimizer; FedSGD instead uses full-batch client SGD with ηc = 1.Other optimizers use task-specific mini-batch sizes, and E = 1 except in Section 5.5.Source: §A.2 · p. 20

A.3 Adaptive Clipping

• Client updates are clipped before averaging, while the server adapts the clipping level to estimate a target norm percentile.The adaptive rule uses the fraction of client updates that remain below the clipping threshold; experiments use q = 0.8, ρ1 = 1, and ηa = 0.2.Source: §A.3 · p. 21

A.4 Learning Rates and Tuning

• Learning rates are selected using held-out validation performance after 1500 rounds with M = 10 across five trials, then held fixed across cohort sizes.The search tests 25 client/server learning-rate configurations per algorithm and dataset; for three datasets, training clients are split 80/20 for tuning.Source: §A.4 · pp. 22–23

B.1 Test Accuracy Versus Communication Round

• Test-accuracy curves across communication rounds are reported for FedSGD, FedAvg, FedAvgM, FedAdagrad, FedAdam, FedLARS, and FedLamb across tasks and cohort sizes.The appendix presents one figure per algorithm.Source: §B.1 · p. 23, Figure 14 · p. 23, Figure 15 · p. 23

B.2 Accuracy Versus Cohort Size

• Figures report train and test accuracy for FedAvg, FedAdam, FedAdagrad, FedLARS, and FedLamb after 1500 rounds as participation rate varies.The x-axis gives the percentage of training clients in each cohort.Source: §B.2 · p. 25, Figure 21 · p. 25, Figure 22 · p. 25

B.3 Cohort Size Speedups

• Increasing cohort size generally reduces rounds to reach a test-accuracy threshold, but speedups are usually not linear.In communication-limited settings, larger cohorts can reach thresholds that small-cohort training cannot; the figures compare five algorithms across cohort sizes.Source: §B.3 · p. 25, Figure 23 · p. 25

B.4 Measures of Accuracy Across Clients

• Cohort size affects accuracy percentiles similarly in nearly all cases, and larger cohorts do not consistently yield fairer models.On CIFAR-100 and Shakespeare, larger cohorts can bring generalization failures across percentiles; this pattern is not observed on EMNIST and Stack Overflow, which have more train and test clients.Source: §B.4 · pp. 26–27, Table 6 · p. 27, Table 8 · p. 27

B.5 Simulating Straggler Effects

• When the straggler parameter λ is large, large-cohort training may require 5–10 times more client compute time than M = 10; at small λ, larger cohorts can reach higher accuracy in comparable time.The simulated round duration is the maximum of M shifted-exponential client runtimes and excludes server and communication time.Source: §B.5 · p. 28, Figure 28 · p. 29

B.6 Pseudo-Gradient Norms

• With local training, pseudo-gradient norms are generally stratified by cohort size, with larger cohorts producing smaller norms; FedSGD usually does not show this pattern.After enough rounds, norms follow inverse square root scaling; FedAdagrad on EMNIST is a notable exception to the local-training pattern.Source: §B.6 · p. 30, Figure 32 · p. 30, Figure 33 · p. 30

B.6.1 Cosine Similarity of Client Updates

• FedAvg client-update cosine similarity approaches zero after a few rounds, while FedSGD similarity remains higher, especially in intermediate rounds.On EMNIST, FedSGD similarity approaches FedAvg's as T → 1500.Source: §B.6.1 · p. 32, Figure 39 · p. 32

B.7 Server Learning Rate Scaling

• The scaling rules increase the server learning rate with cohort size and use a 100-round linear warmup.The tested rates scale with either the square root of the cohort-size ratio or the ratio itself.Source: §B.7 · p. 32
• With linear scaling, FedAvg on EMNIST had one catastrophic failure across five trials at M = 200, and all trials failed at M ≥ 400.The authors suggest the rule may be too aggressive in federated settings, potentially due to client heterogeneity.Source: §B.7 · p. 32, Figure 40 · p. 33
• Square-root scaling avoided the reported training failures, but its accuracy effects varied by task and metric.It improved training accuracy on CIFAR-100 and Shakespeare, yet Shakespeare's test generalization worsened; EMNIST and Stack Overflow saw minor training gains or degradation.Source: §B.7 · p. 32, Figure 41 · p. 33

B.8 Dynamic Cohort Sizes

• Training starts at M = 50 and doubles cohort size every 300 communication rounds, using all training clients when another doubling would exceed their number.Source: §B.8 · p. 33
• For most accuracy thresholds, the strategy matches M = 50 data efficiency before transitioning toward that of larger cohorts.Data efficiency is plotted against the total number of examples processed.Source: §B.8 · p. 33, Figure 42 · p. 34
• Across communication rounds, FedAvg test accuracy differs little between M = 50 and M = 400 or M = 800, while dynamic cohorts generally lie between them.The larger cohort sizes are M = 400 for CIFAR-100 and Shakespeare, or M = 800 for other tasks.Source: §B.8 · p. 34, Figure 43 · p. 34
• The dynamic strategy improves data efficiency for large-cohort training, but does not remedy diminishing returns or generalization failures.It can interpolate data efficiency between small and large cohorts.Source: §B.8 · pp. 33–34, Figure 42 · p. 34, Figure 43 · p. 34
• FedAdam has more returns from increasing cohort size; dynamic cohorts typically start near M = 50 accuracy and later match larger cohorts, helping Stack Overflow but hurting CIFAR-100 generalization.The CIFAR-100 outcome reflects the generalization issues described in Section 3.3.Source: §B.8 · p. 34, Figure 43 · p. 34