On Large-Cohort Training for Federated Learning
Authors: Zachary Charles
Zachary Garrett
Zhouyuan Huo
Sergei Shmulyian
Virginia Smith
Venue: NeurIPS 2021 Poster (2021)
Keywords: federated learning
distributed optimization
large batch training
large cohort training

Abstract: Federated learning methods typically learn a model by iteratively sampling updates from a population of clients. In this work, we explore how the number of clients sampled at each round (the cohort size) impacts the quality of the learned model and the training dynamics of federated learning algorithms. Our work poses three fundamental questions. First, what challenges arise when trying to scale federated learning to larger cohorts? Second, what parallels exist between cohort sizes in federated learning, and batch sizes in centralized learning? Last, how can we design federated learning methods that effectively utilize larger cohort sizes? We give partial answers to these questions based on extensive empirical evaluation. Our work highlights a number of challenges stemming from the use of larger cohorts. While some of these (such as generalization issues and diminishing returns) are analogs of large-batch training challenges, others (including catastrophic training failures and fairness concerns) are unique to federated learning.

On Large-Cohort Training for Federated Learning
Zachary Charles Google zachcharlessgoogle..oom com
Zachary Garrett Google zachgarrett&google.com
Zhouyuan Huo Google zhhuo@google.com
Sergei Shmulyian
Virginia Smith
Google sshmulyian@google.c
Carnegie Mellon University smi tertthv@cm..co edu
Abstract
Federated learning methods typically learn model by iteratively sampling updates from a population of clients. this work, we explore how the number of clients sampled at each round (the cohort size) impacts the quality of the learned model and the training dynamics of federated learning algorithms. Our work poses three fundamental questions. First, what challenges arise when trying scale federated learning to larger cohorts? Second, what parallels exist between cohort sizes in federated learning, and batch sizes in centralized learning? Last, how can we design federated learning methods that effectively utilize larger cohort sizes? We give partial answers to these questions based on extensive empirical evaluation. Our work highlights a number of challenges stemming from the use of larger cohorts. While some of these (such as generalization issues and diminishing returns) are analogs of large-batch training challenges, others (including catastrophic training failures and fairness concerns) are unique federated learning.
Introduction
Federated learning (FL) [52] considers learning model from multiple clients without directly sharing training data, often under the orchestration of a central server. In this work we focus on cross-device FL, in which the aim is to learn across large population of edge devices [27, Table 1]. distinguishing characteristic of cross-device FL is partial participation of the client population: Due systems constraints such network size, the server typically only communicates with subset of the clients at time!. For example, in the popular FedAvg algorithm [52], at each communication round the server broadcasts its current model to subset of available clients (referred to cohort), who use the model to initialize local optimization and send their model updates back to the server.
Intuitively, larger cohort sizes have the potential to improve the convergence of FL algorithms. By sampling more clients per round, we can observe a more representative sample of the underlying iopulation-possiblee reducing the number of communication rounds needed to achieve given accuracy. This intuition is reflected in many convergence analyses of FL methods [29, 31, 32, 62, 69], which generally show that asymptotic convergence rates improve the cohort size increases.
Larger cohorts can also provide privacy benefits. For example, when using the distributed differential privacy model [6, 11, 16, 65] in federated learning, noise is typically added to the updates sent from the clients to the server [54]. This helps preserve privacy but can also mar the utility of the learned model. By dividing the noise among more clients, larger cohorts may mitigate detrimental effects of noise. Moreover, since privacy tends to decrease as a function of the number of communication
In contrast, cross-silo settings often have small set of clients, most of which participate in each round [27].
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

rounds [1, 17], larger cohorts also have the potential to improve privacy in FL by reducing the number of rounds needed for convergence.
Motivated by the potential benefits of large-cohort training, we systematically explore the impact of cohort size in realistic cross-device settings. Our results show that increasing the cohort size may not lead to significant convergence improvements in practice, despite their theoretical benefit [69]. Moreover, large-cohort training can introduce fundamental optimization and generalization issues. Our results are reminiscent of work on large-batch training in centralized settings, where larger batches can stagnate convergence improvements [14, 19, 51, 70, 71], and even lead to generalization issues with deep neural networks [23, 30, 46-48, 50, 64]. While some of the challenges we identify with large-cohort training are parallel to issues that arise in large-batch centralized learning, others are unique to federated learning and have not been previously identified in the literature.
Contributions. In this work, we provide a novel examination of cohort sizes in federated learning. We give a wide ranging empirical analysis spanning many popular federated algorithms and datasets (Section 2). Despite the many possible benefits of large-cohort training, we find that challenges exist in realizing these benefits (Section 3). We show that these issues are caused in part by distinctive characteristics of federated training dynamics (Section 4). Using these insights, we provide partial solutions to the challenges we identify (Section 5), focusing on how to adapt techniques from large- batch training, and the limitations of such approaches. Our solutions are designed to serve as simple benchmarks for future work. We conclude by discussing limitations and open problems (Section 6). Throughout, we attempt to uncover interesting theoretical questions, but remain firmly grounded in the practical realities of federated learning.
1.1 Related Work
Large-batch training. In non-federated settings, mini-batch stochastic gradient descent (SGD) and its variants are common choices for training machine learning models, particularly deep neural networks. While larger mini-batch sizes ostensibly allow for improved convergence (in terms of the number of steps required to reach desired accuracy), in practice speedups may quickly saturate when increasing the mini-batch size. This property of diminishing returns has been explored both empirically [14, 19, 51, 64] and theoretically [48, 70]. Beyond the issue of speedup saturation, numerous works have also observed a generalization gap when training deep neural networks with large batches [23, 30, 46, 47, 50, 71]. Our work differs from these areas by specifically exploring how the cohort size (the number of selected clients) affects federated optimization methods. While some of the issues with sarge-batch training appear in large-cohort training, we also identify number of new challenges introduced by the federated setting.
Optimization for federated learning. Significant attention has been paid towards developing feder- ated optimization techniques. Such work has focused on various aspects. including communication- efficiency [5, 34, 37, 52], data and systems heterogeneity [25, 28, 29, 39-42, 67]. and fairness [26, 43]. We provide description of some relevant methods in Section 2, and defer readers to recent surveys such as [27] and [41] for additional background. One area pertinent to our work is that of variance reduction for federated learning, which can mitigate negative effects of data heterogeneity [28, 29, 73]. However, such methods often require clients to maintain state across rounds [29, 731, which may be infeasible in grosse settings [27]. Moreover, such methods may not perform well in settings with limited client participation [62]. Many convergence analyses of federated optimization methods show that larger cohort sizes can lead to improved convergence rates, even without explicit variance reduction [31, 32, 69]. These analyses typically focus on asymptotic convergence, and require assumptions on learning rates and heterogeneity that may not hold in practice [9, 27]. In this work, we attempt to whether increasing the cohort size leads to improved convergence in practical, communication -imited settings.
Client sampling. A number of works have explored how to select cohorts of a fixed size in cross- device FL [10. 12, 18, 59, 63]. Such methods can yield faster convergence than random sampling by carefully selecting the clients that participate at each round, based on quantities such as the client loss. However, such approaches typically require the server to able to choose which clients participate in cohort. In practice, cohort selection in cross-device federated learning is often governed by client availability, and is not controlled by the server [7, 61]. In this work we instead focus on the impact of size of the cohort, assuming the cohort is sampled at random.

---

2
Preliminaries
Federated optimization methods often aim to minimize a weighted average of client loss functions:
K
min f(x) : Σ Pkfx(x),
(I)
where K is total number of clients, the Pa are client weights satisfying Pk 2 0, and fú. is the loss function of client k. For practical reasons, Pk: is often set to the number of examples in client k's local dataset [42, 52].
To solve (1), each client in sampled cohort could send V (xx() to the server, and server could then apply (mini-batch) SGD. This approach is referred to as FedSGD [52]. This requires communication for every model update, which may not desirable
settings. To address this, McMahan et al. [52] propose FedAvg. in which clients perform multiple epochs of local training, potentially reducing the number of communication rounds needed for convergence.
We focus on more general framework, FedUpt, introduced by Reddi et al. [62] that uses both client and server optimization. At each round, the server sends its model M. Each client C₁ C performs E epochs of training using mini-batch SGD with client learning rate Ne producing local model 2g. Each client & E C then communicates their client update Ag the server, where A == x is the difference between the client's 's local model and the server model. The server computes weighted average A of the client updates, and updates its own model via
cohort of clients C' of size
SERVEROPT 17, A)
where SERVEROPT(I, 7/s is some first-order optimizer, 7/s the server learning rate, and g is gradient estimate. For example, if SERVEROPT is SGD, then SERVEROPT n). 7799. The in (2) referred to as a pseudo-gradient [62]. While A may not be unbiased estimate of f. it can serve a somewhat comparable role (though show in Section 4, there are important distinctions). Full pseudo-code of FedOpt is given Algorithm
Algorithm FedOpt framework
Input: M, T E, x' Te, Ns. SERVEROPT, {Pk 1 for ( = 1, do
The server selects cohort C, of M clients uniformly at random, without replacement. The server sends x' to all clients in C..
Each client & E C. performs F epochs of mini-batch SGD on Su with step-size Tle- After training, each C. has local model o and sends A = - x to the server.
The server computes yseudo-gradiens A' and updates its model via
ΣkEc, P
r111 SERVEROPT (2e, Mss At).
Σxec, Pk
Algorithm generalizes a number of federated learning algorithms, including FedAvg [52], FedAvgM [25]. FedAdagrad [62], and FedAdam [62]. These are the cases where SERVEROPT is SGD, SGD with momentum, Adagrad [15, 53], and Adam [33], respectively. FedSGD is realized when SERVEROPT SGD, 17c 1, 1. and each client performs full-batch gradient descent.
2.1 Experimental Setup
We aim to understand how the cohort size M impacts the performance of Algorithm 1. In order to study this, we perform a wide-ranging empirical evaluation using various special cases of Algorithm across multiple datasets, models, and tasks. We discuss the key facets of our experiments below.
Datasets, models, and tasks. We use four datasets: CIFAR- [35], EMNIST 33, Shakespeare [8], and Stack Overflow [3]. For CIFAR-100, we use the client partitioning proposed by Reddi et al. [62]. The other three datasets have natural client partitions that we use. For EMNIST, the handwritten characters partitioned by their author. For Shakespeare, speaking lines in Shakespeare plays

---

Table 1: Dataset statisties.
DATASET
TRAIN CLIENTS
TRAIN EXAMPLES
TEST CLIENTS
TEST EXAMPLES
CIFAR-100
500
50,000
100
10,000
EMNIST
3,400
671,585
3,400
77,483
SHAKESPEARE
715
16,068
715
2,356
STACK OVERFL.OV
342,477
135,818,720
204,088
16,586,035
partitioned by their speaker. For Stack Overflow, posts on the forum are partitioned by their author. The number of clients and examples in the training and test sets are given in Table 1.
For CIFAR-I00 100, we train ResNet-18, replacing batch normalization layers with group normalization (as proposed and empirically validated in federated settings by Hsieh et al. [24]). For EMNIST, we train a convolutional network with two convolutional layers, max-pooling max-pooling, dropout, and two dense layers. For Shakespeare, prediction. For Stack Overflow, we perform next-word- -oodd-pedeiction using an RNN with a single LSTM layer. For full details on the models and datasets, Appendix A.1.
train RNN with two LSTM layers to perform next-character-
Algorithms. We implement many special cases of Algorithm 1, including FedSGD, FedAvg, FedAvgM, FedAdagrad, and FedAdam. We also develop two novel methods: FedLARS and FedLamb, which are the special cases of Algorithm where SERVERÖPT is LARS [71] and Lamb [72], respectively. See Section 5 for the motivation and full details of these algorithms.
Implementation and tuning. Unless otherwise specified. in Algorithm clients perform E epochs of training with mini-batch SGD. Their batch size is fixed per-task. We set pı to be the number of examples client kts dataset. We tune learning rates for all algorithms and models using held-out validation set: We perform T 1500 rounds of training with M 50, E for each algorithm and model, varying T/e, Y)s over {10* - i I } and select the values that maximize the average validation performance over 5 random trials. All other hyperparameters (such as momentum) are fixed. For more details, see Appendix A. We provide spencsource implementations of all simulations in TensorFlow Federated [4] All experiments
conducted using clusters of multi-core CPUs, though our results are independent of wall-clock time and amount compute resources.
Presentation of results. We apply the algorithms above to the tasks listed above with varying cohort sizes. For brevity, we present only a fraction of our results, selecting representative experiments to illustrate large-cohort training phenomena. The full set of experimental results can be found in Appendix B. We run 5 random trials for each experiment, varying the model initialization and which clients are sampled per round. In all subsequent figures, dark lines indicate the mean across the 5 trials, and shaded regions indicate one standard deviation above and below the mean.
Large-Cohort Training Challenges
In this section explore challenges that exist when using large cohorts in federated learning. While some of these challenges mirror issues in large-batch training, others are unique to federated settings. While provide concrete recommendations for mitigating some of these challenges, discussion is generally centered around introducing and exploring these challenges in the context of federated learning.
3.1 Catastrophic Training Failures
We first discuss practical issue unique to large-cohort training. Due to data heterogeneity, the server model x may be misaligned with some client's is loss Ss, in which case V (z(x)) can blow up and lead to optimization problems. This issue is exacerbated by large cohorts, as we are more likely to sample misaligned clients. To demonstrate this, applied FedAvg with varying cohort sizes M, using learning rates tuned for 10. For each M, we performed random trials and recorded whether a catastrophic training failure occurred, in which the training accuracy decreased by factor of at least 1/2 in a single round.
'https://giie www.coo.oooooccccccc http:////hhhhhhooooo///////////a http:///w.....ooo//cc/////////t. cohor_

---

CUNIS Folley
2985: Folleg
ENNSIT Foelhe
EHH ST. hoppey
0.4
0.4
0.c
60.2
550
1550
3500
520 2003
1500
507
1092
1350
500
1250
1500
Figure 1: Applying FedAvg to EMNIST with cohort size 200. We plot the train accuracy and norm of the pseudo- pseudo-gradient for trial that ran successfully (left), and one that experienced catastrophic training failure (right). The trials differed only which clients were randomly sampled each round.
On EMNIST, the failure rate increased from 0% for M 10 to 80% for M = 800. When failures occurred, we consistently saw spike in the norm of the trends occurred all datasets. In order to prevent this spike, we apply clipping to the client updates. We use the adaptive clipping method of [2]. While this technique was originally designed for training with differential privacy, we found that it greatly improved the stability of large-cohort training. Applying FedAvg to EMNIST with adaptive clipping, no catastrophic training failures occurred for any cohort size. We use adaptive clipping in all subsequent experiments. For more details, see Appendix A.3.
A (see Figure 1). These
3.2 Diminishing Returns
In this section, we show that increasing M in Algorithm can lead to improved convergence, but that such improvements diminish with M. To demonstrate this, we plot the test accuracy of FedAvg and FedSGD across multiple tasks, for varying cohort sizes M. Results for CIFAR- 100 and Stack Overflow are given in Figure 2, though we observe similar trends for all tasks (Appendix B.
CIPA8-103, Factug
Stace Overioa Tedtug
OPAR: 100, TadSCD
Stack Overfow Teep2e
0.3
6,125
20.103
0.10
6.053
300
430
0.05
0.023
400
0.0
1200
3500
500
1020
1500
0.00
510
1002
1500
5:0
2000
1520
Figure 2: Test accuracy of FedAvg (top) and FedSGD (bottom), for various cohort sizes M, over the course of 1500 communication rounds.
We see that convergence benefits do not scale linearly with cohort size. While increasing M from to 10 can significantly improve convergence, there is generally a threshold after which point increasing change in convergence. This threshold is typically between M 10 and M = 50. Interestingly, this seems to be true for both tasks, even though M = 50 represents 10% of the training clients for CIFAR-10. 10, but only approximately 0.015% of the training clients for Stack Overflow. We see comparable results for EMNIST and Shakespeare, as well as for other optimizers, including FedAdam and FedAdagrad. See Appendix B.1 for the full results. In short, increasing M alone can lead to diminishing returns, or even no returns in terms of convergence. This mirrors issues of diminishing returns in large-batch training [14, 19, 51, 64].
M incurs little
that
3.3 Generalization Failures
Large-batch centralized optimization methods have repeatedly been shown to converge to models
with worse generalization ability than models found by small-batcl methods [23, 30, 46, 47, 50, 71]. Given the parallels between batch size in centralized learning and cohort size in FL, this raises obvious questions about whether similar issues occur in FL. In order to test this, we applied FedAvg, FedAdam, and FedAdagrad with different cohort sizes to various models. In Figure 3 plot the train and test accuracy of our models after T 1500 communication rounds of FedAvg, FedAdam, and FedAdagrad.
We find that generalization issues do oceur in FL. For example, consider FedAdam on the CIFAR-100 task. While it attains roughly the same training accuracy for M e (50, 100, 200, 400), see that
5

---

CTAY 107
01-40- 300
6.54
20.3
Feridarre!
26.25
2.50
654
Toin
6575
1.3
032
022
1.0
5.0
20.0
56.2
2.2
1.0
5.0
2015
$5.0
011
0.5
54
15.3
55.7
23
0.0
5.6
253
53.3
Figure The train accuracy and test accuracy of FedAvg, FedAdam, and FedAdagrad on CIFAR-I00 (left) and Shakespeare (right) after training for 1 1500 communication rounds, for varying cohort sizes. The z-axis denotes the percentage of training clients in each cohort.
the larger cohorts uniformly lead to worse generalization. This resembles the findings of Keskar al. [30], who show that generalization issues of large-batch training can oceur even though the methods reach similar training losses. However, generalization issues do not occur uniformly. It often optimizer-dependent (as in CIFAR-ION 100) and does not occur on the EMNIST and Stack Overflow datasets (see Appendix B.2). Notably, CIFAR-100 and Shakespeare have many fewer clients overall. Thus, large-cohort training may reduce generalization, especially when the cohort size is large compared to the total number of clients.
3.4 Fairness Concerns
One critical issue in FL is fairness across clients, as minimizing (1) may disadvantage some clients [43, 56]. Intuitively, large-cohort training methods may be better suited for ensuring fairness, since greater fraction of the population is allowed to contribute to the model at each round. As coarse measure of fairness, we compute percentiles of accuracy of our trained models across test clients. Under many notions of fairness, this would lead to higher accuracy values for smaller percentiles. The percentiles for FedAdam on each task are given in Figure 4.
SPRIST
2.5
03
Colorr S
Cohor 522
Cohort
Cofort Sie
Figure 4: Accuracy of FedAdam after training for 1500 communication rounds using varying cohort sizes and tasks. The box plots show the 5th, 25th, 50th, 75th, and 95th percentiles of accuracy across test clients.
We find that the cohort size seems affect all percentiles in the same manner. For example, CIFAR- 100, M = 50 performs better for smaller percentiles and larger percentiles than larger M. This mirrors the CIFAR-100 generalization failures from Section 3.3. By contrast, for Stack Overflow we see increases in all percentiles as we increase M. While the accuracy gains are only slight, they consistent across percentiles. This suggests connection between the faimess of federated training algorithm and the fraction of test clients participating at every round. Notably, increasing M seems to have little effect on the spread between percentiles (such as the difference between the 75th and 25th percentiles) beyond a certain point. See Appendix B.4 for more results.
3.5 Decreased Data Efficiency
Despite issues such as diminishing returns and generalization failures, federated optimization methods see some benefit from larger cohorts. Large-cohort training, especially with adaptive optimizers, often leads to faster convergence to given accuracy thresholds. For example, in Figure 5, we see that the number of rounds FedAdam requires to reach certain accuracy thresholds generally decreases with the cohort size.
While it is tempting to say that large-cohort methods are "faster", this ignores the practical costs of large-cohort training. Completing a single communication round often requires more resources with larger cohorts. To showcase this, also plot the accuracy of FedAdam with respect to the number of

---

shokoporss rollcom
0.5
STACESPLOIG POLAZAT)
0.275
Slack Ovailonn Fodlico
Slace Ove Polldam
620
0.223
104
20.45
0.200
6.10
0.175
0.40
500
1050
3500
0.0
100
Teall Nombarai
30
500
2003
1300
Tebs
10"
10
10
Figure 5: Test accuracy of FedAdam on Shakespeare (left) and Stack Overflow (right) with various cohort sizes. We plot versus the number of communication rounds and the number of examples processed in total.
examples seen in Figure 5. This measures the data-efficiency of large-cohort training, and shows that large cohort-training requires significantly more examples per unit-aceuracy.
While data-inefficiency also occurs in large-batch training [5 IJ. it is especially important in federated learning. Large-cohort training faces greater limitations on parallelizability due to data-sharing constraints. Worse, in realistic cross cross-device settings client compute times can scale super-linearly with their amount of data, so clients with more data are more likely to become stragglers [7]. This straggler effect means that data- intta-inffficeent algorithms may require longer training times. To demonstrate this, we show in Appendix B.5 that under the probabilistic straggler runtime model from [38), large-cohort training can require significantly more compute time to converge.
4 Diagnosing Large-Cohort Challenges
We now examine the challenges in Section 3, and provide partial explanations for their occurrence. in (1)
One of the key differences between FedAvg and FedSCD is what the pseudo-gradient represents. In FedSGD, À is a stochastic gradient estimate (i.e., E[A] Vf. where the expectation is over all randomness in a given communication round). For special cases of Algorithm 1 where
clients perform multiple local training steps, is not an unbiased estimator of f f [9, 49, 60]. While increasing the cohort size should reduce the variance of A as an estimator of E it is unclear what this quantity represents.
To better understand A, we plot its norm on Stack Overflow in Figures 6a and 6b. For FedSGD, ||]] decreases slightly with M. but has high variance. By contrast, for FedAvg larger cohorts lead to smaller norms with little overlap. The decrease in norm obeys an inverse square root rule: Let A,, A2 be pscudo-gradients at some round for cohort sizes M1, M2. For FedAvg, | & /2 Mi We use this rule predict geuudogrradient norms for FedAvg in Figure 6c. After a small number of rounds, we obtain a remarkably good approximation. To explain this, we plot the average cosine similarity between client updates Aj at each round in Figure 6d, with M 50. For FedAvg, the client updates are on average almost orthogonal. This explains Figure 6b, as A is an average of nearly orthogonal vectors. As we show Appendix B.6, similar results hold for other tasks and optimizers.
Shack Oxerloo, FedSco
Stork Overfioe Artil Merr
Stack Ouerfoos Pred
Stank Owedfoo
Poubyconden hage 0.5
0.5
1.0
2..0
0.4
1350
0.6%
1936
1500
1500
500
1030
Communkation
communcatir
Reands
Rounds
(a)
(b)
(c)
(d)
Figure 6: The graddent norm of FedSGD (a) and FedAvg (b) on Stack Overflow with varying cohort sizes M. We also plot the predicted norm for FedAvg using an inverse square root scaling rule
relative to M = 50 (c) and the average cosine similarity of client updates for M 50 (d).
Implications for large-cohort training. This near-orthy of client updates is key to under- standing the challenges in Section 3. The diminishing returns in Section 3.2 occur in part because increasing M leads to smaller updates. This also sheds light on Section 3.5: In large-cohort training, we take an average of many pearly-otthogonal vectors, so each client's 's examples contribute little. The decreasing pscudo-gradien norms in Figure 6c also highlights an advantage of methods such as
7

---

FedAdam and FedAdagrad: Adaptive server optimizers employ a form of normalization that makes them somewhat scale-invariant, compensating for this norm reduction.
5 Designing Better Methods
We now explore an initial set of approaches aimed at improving large-cohort training, drawing inspiration where possible from large-batch training. Our solutions are designed to provide simple baselines for improving large-cohort training. In particular, our methods and experiments are intended to serve as a useful reference for future work in the area, not fully solve the challenges of large- cohort training.
5.1 Learning Rate Scaling
One common technique for large-batch training is to scale the learning rate according to the batch size. Two popular scaling methods are square root scaling [36] and linear scaling [20]. While such techniques have had clear empirical benefit in centralized training, there many different ways that they could be adapted to federated learning. For example, in Algorithm 1, the client and server optimization both use learning rates that could be scaled.
We consider the following scaling method for large-cohort training: We fix the client learning rate, and scale the server learning rate with the cohort size. Such scaling may improve convergence by compensating for the sgeudo-gradien norm reduction in Figure We use square root and linear scaling rules: Given a learning rate T/s tuned for M, for M' 2 M we use learning rate 1/, where
VM VM
M
!)s (square root scaling) OR 1/3-
M ")s (linear scaling).
We also version of the warmup strategy from [20]. For the first W communication rounds, linearly increase the server learning rate from y), to n//-- In our experiments, we set W = 100 and use reference server learning rate Tis tuned for M 50.
Our experiments show that server learning rate scaling rules have mixed efficacy in large-cohort training. Linear scaling is often too aggressive for federated learning, and caused catastrophic training failures beyond M = 100 even when using adaptive clipping (see Appendix B.7). By contrast, square root scaling did not cause catastrophic training failures. Its performance (Figure 7) varied widely across tasks. For example, it significantly improved train accuracy on Shakespeare, but reduced test accuracy. While it led to small accuracy improvements on Stack Overflow for some cohort sies, it degraded accuracy for the largest cohort sizes. In sum, we find that applying learning rate scaling at the server may not directly improve large-cohort training.
Ostrfow, Feldx)
Searn Far Stains
0.575
0.54
0.13
0.570
0.13
0.555
5.522
0.500
0.00
0.02
Co Caars Scaling
0.552
Soars Roet
0.00
Sqare Eostin
Coreet 5t00
Contint Size
Cohort 5
Figure 7: The train and test accuracy of FedAvg using square root scaling with warmup, versus no scaling. Results are given for Shakespeare (left) and Stack Overflow (right).
5.2 Layer-wise Adaptivity
Another popular technique for large-batch training is layer-wise adaptivity. Methods such as LARS [71] and Lamb [72] use layer-wise adaptive learning rates, which may allow the meth- ods to train faster than SGD with linear scaling and warmup large-batch settings [71, 72]. We propose two new federated versions of these optimizers, FedLARS and FedLamb. These are special cases of Algorithm 1, where the server uses LARS and Lamb, respectively. Given the difficulties of learning rate scaling above, FedLARS and FedLamb may perform better in large-cohort settings.
8

---

Strkssccor
Stack Overfoon
0.50
0.4
0.825
1055
Fediig
0.54
feckdom
0.522
redicsy
0.2
0.750
0.50
0.0
0.2
1.5
5/2
25.0
80.0
0.2
0,20
1:45
5:80
23:5
0.3
0..
54
15.3
5539
0.032
DIR2 3:30M
1.300 2000
Figure 8: The test accuracy of various methods, including FedLARS and FedLamb, after training for 1 1500 rounds, for varying cohort sizes and on varying tasks. The r-axis denotes percentage of training clients in each cohort.
In Figure 8 we present the test accuracy of various methods, including FedLARS and FedLamb, for varying cohort sizes. In most cases, we see that FedLamb performs comparably to FedAdam for large cohort sizes, but with slightly worse performance in intermediate stages. One notable exception is Stack Overflow, in which FedLamb performs well even for an eventual drop in test accuracy for M > 100. FedLARS has decidedly mixed performance, While it performs well on CIFAR-10, it does not do well on EMNIST or Shakespeare. While federated layer-wise adaptive algorithms can be better than coordinate-witg wise adaptive algorithms certain datasets in some large-cohort settings, our results do not indicate that they are universally better.
1. As Section 3.3, FedLamb sees
5.3 Dynamic Cohort Sizes
As we saw in Section 3.5, large-cohort training can reduce data efficiency. Part of this stems from the fact that larger cohorts may help very little for smaller accuracy thresholds (see Figure 5). In order to improve data efficiency, we may able to use smaller cohorts in earlier optimization stages, and increase the cohort size over time. This technique is parallel to "dynamic batch size" techniques used in large-batch training [66]. In order to test the efficacy of such techniques in large-cohort training, we start with an initial cohort size of . 50 and double the size every 300 rounds up to M 800 (or the maximum population size if smaller). This results in doubling the cohort size a maximum of times over the 1500 rounds of training we perform. We plot the results for FedAvg and FedAdam on CIFAR- 100 and Stack Overflow in Figure 9. See Appendix B.8 for results on all tasks.
C17A0-200 Fochys
07A8-200 Fockdum
Slace Overfan, Poddag
Slack Dest fle fosédem
18
20.4
Teak Acowecy
6.3
0.20
2..2
0.2
0:155
0.2
Tons Nimber Daamieee
Teta Munner Dxamp
TeDe Sumzer Doamples
Tedh Number of Examsles
Figure 9: Test accuracy of FedAvg and FedAdam on Shakespeare (left) and Stack Overflow (right), with respect to the total number of examples processed, using fixed and dynamic cohort sizes.
This dynamic strategy attains data efficiency closer to a fixed cohort size of M = 50, while still obtaining final accuracy closer to having used large fixed cohort size. While our initial findings promising, we note two important limitations. First, the accuracy of the dynamic strategy is bounded by the minimum and maximum cohort size used: It never attains better accuracy than M 800. Second, the doubling strategy still faces the generalization issues discussed in Section 3.3.
5.4 Normalized FedAvg
While the methods above show promise in resolving some of the issues of large-cohort training, they also introduce extra hyperparameters (such what type of learning rate scaling to use, how often to double the cohort size). Hyperparameter tuning can be difficult in federated learning, especially cross-device federated learning [27]. Even adaptive methods like FedAdam introduce number of new hyperparameters that can be challenging to contend with. We are therefore motivated to design a large-cohort training method that does not introduce any new hyperparameters.
Recall that in Section 4, we showed that for FedAvg, the client updates (Af, in Algorithm and Algorithm 2) are nearly orthogonal in expectation. By averaging nearly orthogonal updates in large-

---

cohort training, we get server pseudo-gradient A' 'hhtt is close to zero, To compensate, we propose a variant of FedAvg where rather than applying SGD the server gpeudo-gradient (as in Algorithm 1), we apply SGD to the normalized server pseudo-gradient That is, the server updates its model via
x' 1/. //.A/114 //
This method, which refer to normalized FedAvg. is a federated analog of normalized SGD methods used for centralized learning [57]. It introduces no new hyperparameters with respect to Algorithm 1. To test it, we present its training and test accuracy versus cohort size in Figure 10. Notably, we re-use the same learning rates tuned for (unnormalized) FedAvg. For full results, see Appendix B.9.
CPAN-AD2
ENIST
thanespzere
Sizce Oveflow
0.86-
0.24
70.84
8.50
02
0.02
8.50
1:30
0.00
Fedles
0.78
0.52
Felkro
2:30
Fibles Narmadenthrries
102
1020
15
200
200
2002
Calor Sizz
Calert
Conart siuc
Figure 10: The test accuracy of FedAvg and the normalized variant of FedAvg, after training for 1500 communication rounds. Results are given for various cohort sizes and tasks.
We find that for most cohort sizes and on most tasks, normalized FedAvg achieves better training accuracy for larger cohorts. Thus, this helps mitigate the diminishing returns issue in Section 3.2. We note two important exceptions: for EMNIST, the normalized FedAvg is slightly worse for all cohort sizes. For Stack Overflow, it obtains worse training accuracy for the largest cohort size. However, we see significant improvements on CIFAR-100 and all but the largest cohort sizes for Stack Overflow. We believe that the method therefore exhibits promising results, and may be improved in future work.
5.5 Hyperparameter tuning and other results.
The methods discussed above, including learning rate scaling and layer-wise adaptivity, can require significant tuning perform well [58]. To date, little work has been paid to how to tune hyperpa- rameters in federated learning. Such work may be vital to obtain optimal performance, especially given our observations in Section 4 and the client-server structure of federated algorithms, which gives rise to many more hyperparameters. In Algorithm 1, tuning could involve the client optimizer, the client batch size, the server optimizer, and the cohort size. In fact, the client batch size is key hyperparameter. Recall that clients perform E epochs of mini-batch SGD on their local datasets. Fixing E, the batch size dictates the number local training steps they perform. As we show in Appendix B. 10, this number of local steps is critical for achieving maximal performance, and may be necessary to tune according to the cohort size.
Limitations and Future Work
In this work we explore the benefits and limitations of large-cohort training in federated learning. As discussed in Sections 3.5 and 5, focusing on the number of communication rounds often obscures the data efficiency of a method. This in turn impacts many metries important to society, such as total energy consumption or total carbon emissions. While we show that large-cohort training can negatively impact such metries by reducing data-efficiency (see Section 3.5 and Appendix B.5), a more specialized focus on these issues is warranted. Similarly, we believe that an analysis of fairness in large-cohort settings going beyond Section 3.3 would be beneficial.
Future work also involves connecting large-cohor training to other important aspects of federated learning, and continuing to explore connections with growing lines of work in large-batch training. In particular, we wish to see whether noising strategies, especially differential privacy mechanisms, can help overcome the generalization issues of large-cohort training. Personalization may also help mitigate issues of generalization and fairness. Finally, although not a focus of our work, we note that some of the findings above may extend to cross-silo settings, especially if communication restrictions require subsampling clients.
10

---

References
[1] Martin Abadi, Andy Chu, lan Goodfellow, H. Brendan McMahan, llya Mironov, Kunal Talwar, and Li Zhang. Deep learning with differential privacy. In Proceedings of the 2016 ACM SIGSAC Conference on Computer and Communications Security, page 308-318, 2016.
[2] Galen Andrew, Om Thakkar, H Brendan McMahan, and Swaroop Ramaswamy. Differentially private learning with adaptive clipping. arXiv preprint arXiv: 00970787,, 2019.
[3] The TensorFlow Federated Authors. TensorFlow Federated Stack Overflow dataset, 2019. URL http::://wwww.ensorfiii... http://wwwooo...oooooeeeeeerrraa aatasett//staaoooerfll////aaaadd http://eeerraeeppppppoo////po///
[4] The TFF Authors. TensorFlow Federated, 2019. URL https://www.teesooffilo org/
federated.
[5] Debraj Basu, Deepesh Data, Can Karakus, and Suhas Diggavi. Qsparse-local-SGD; Dis- tributed SGD with quantization, sparsification and local computations. In Advances Neural Information Processing Systems, 2019.
[6] Andrea Bittau, Úlfar Erlingsson, Petros Maniatis, Ilya Mironov, Ananth Raghunathan, David Lie, Mitch Rudominer, Ushasree Kode, Julien Tinnes, and Bernhard Seefeld. PROCHLO: Strong privacy for analytics in the crowd, In Proceedings of the 26th Symposium on Operating Systems Principles, pages 441 1414459,40 2017.
[7] Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloé Kiddon, Jakub Konečný, Stefano Mazzocchi, Brendan McMahan, Timon Van Overveldt, David Petrou, Daniel Ramage, and Jason Roselander. Towards federated learning at scale: System design. In Proceedings of Machine Learning and Systems. Proceed- ings of MLSys, 2019.
[8] Sebastian Caldas, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. LEAF; A benchmark for federated settings. arXiv preprint arXiv: 112.011//, 2018.
[9] Zachary Charles and Jakub Konečný. Convergence and accuracy trade-offs in federated learn- ing and meta-learning. In Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, 2021.
[10] Wenlin Chen, Samuel Horvath, and Peter Richtarik. Optimal client sampling for federated learning. arXiv preprint S0td::20100132... 2020.
[11] Albert Cheu, Adam Smith, Jonathan Ullman, David Zeber, and Maxim Zhilyaev. Distributed differential privacy via shuffling. In Annual International Conference on the Theory and Applications of Cryptographic Techniques, pages 375-403, 2019.
[12] Yae Jee Cho, Jianyu Wang, and Gauri Joshi. Client selection in federated learning: Convergence analysis and power-of-choice selection strategies. arXiv preprint 2xXiv20010.01443 2020.
[13] Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. EMNIST: Extending MNIST to handwritten letters. In 20/7 International Joint Conference on Neural Networks (IJCNN, pages 2921-2926. IEEE, 2017.
[14] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc aare'aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, Quoe Le, and Andrew Ng. Large scale distributed deep networks. In Advances in Neural Information Processing Systems, 2012.
[15] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(ual):121-2155) 2011.
[16] Úlfar Erlingsson, Vitaly Feldman, Ilya Mironov, Ananth Raghunathan, Shuang Song, Kunal Talwar, and Abhradeep Thakurta. Encode, shuffle, analyze privacy revisited: Formalizations and empirical evaluation. arXiv preprint 2XXv:2001103616 2020.
11

---

[17] Antonious M Girgis, Deepesh Data, Suhas Diggavi, Peter Kairouz, and Ananda Theertha Suresh. Shuffled model of federated learning: Privacy, communication and accuracy trade-offs, arXiv preprini arXiv:2008.077 071000 2020.
[18] Jack Goetz, Kshitiz Malik, Due Bui, Seungwhan Moon, Honglei Liu, and Anuj Kumar. Active federated learning, arXiv preprint arXiv:/909 1909.1249,, 2019.
[19] Noah Golmant, Nikita Vemuri, Zhewei Yao, Vladimir Feinberg, Amir Gholami, Kai Rothauge, Michael W Mahoney, and Joseph Gonzalez. On the computational inefficiency of large batch sizes for stochastic gradient descent. arXiv preprint arXiv: 1211,1291, 2018.
[20] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangging Jia, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in I hour. arXiv preprint arXiv: 026777 2017.
[21] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and paltern recognition, pages 770-778, 2016.
[22] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, (8):1735-1780, 1997.
[23] Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. Advances in Neural Information Processing Systems, 2017.
[24] Kevin Hsieh, Amar Phanishayee, Onur Mutlu, and Phillip Gibbons. The non-IID data quagmire of decentralized machine learning. In Proceedings of the 37th International Conference Machine Learning. 2020.
[25] Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.0633: 2019.
[26] Zeou Hu, Kiarash Shaloudegi, Guojun Zhang, and Yaoliang Yu. FedMGDA+: Federated learning meets multi-objeccive optimization. arXiv preprint arXiv:2006. / 1489, 2020.
[27] Peter Kairouz, H. Brendan McMahan, and contributors. Advances and open problems in federated learning. Foundations and Trends in Machine Learning, 14(1), 2021. ISSN 1935-8237.
[28] Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. arXiv preprint arXiv:2008. 000::20000.0600. 2020.
[29] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Proceedings of the 37th International Conference on Machine Learning, 2020.
[30] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On larch-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
[31] Ahmed Khaled, Konstantin Mishchenko, and Peter Richtárik. First analysis of local GD on heterogeneous data. arXiv preprini (909.0:99900471: 2019.
[32] Ahmed Khaled, Konstantin Mishchenko, and Peter Richtarik. Tighter theory for local SGD on identical and heterogeneous data. In Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, 2020.
[33] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[34] Jakub Konečný, H Brendan McMahan, Daniel Ramage, and Peter Richtárik, Federated optimiza- tion: Distributed machine learning for on-device intelligence. arXiv preprint arXiv: 161000257. 2016.
12

---

[35] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
[36] Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. arXiv preprint arXiv; 1404,5944 2014.
[37] Yassine Laguel, Krishna Pillutla, Jerôme Malick, and Zaid Harchaoui. A superquantile ap- proach federated learning with heterogeneous devices. In 2021 55th Annual Conference on Information Sciences and Systems (CISS), 2021.
(38) Kangwook Lee, Maximilian Lam, Ramtin Pedarsani, Dimitris Papailiopoulos, and Kannan Ramchandran. Speeding up distributed machine learning using codes. IEEE Transactions on Information Theory, 64(3):1514-135 2017.
[39] Daliang Li and Junpu Wang. FedMD: Heterogenous federated learning via model distillation. arXiv preprint arXiv: 1910.0358, 2019.
[40] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smithy. FedDANE: A federated newton-type method. In 2019 53rd Asilomar Conference on Signals, Systems, and Computers, pages 1227-1231. IEEE, 2019.
[41] Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Chal- lenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50-60) 60, 2020.
[42] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems 2020, pages 429-450, 2020.
[43] Tian Li, Maziar Sanjabi, Ahmad Beirami, and Virginia Smith. Fair resource allocation in federated learning. In International Conference on Learning Representations, 2020.
[44] Wei Li and Andrew McCallum. Pachinko allocation: DAG-structured mixture models of topic correlations. In Proceedings of the 23rd International Conference on Machine Learning, 2006.
[45] Guanfeng Liang and Ulas C Kozat. Tofec: Achieving optimal throughput-delay trade-off of cloud storage using erasure codes. In IEEE INFOCOM 2014-IEEE Conference on Computer Communications, 2014.
[46] Tao Lin, Sebastian Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't 't use large use local SGD. In International Conference Learning Representations, 2019.
[47] Tao Lin, Lingjing Kong, Sebastian Stich, and Martin Jaggi. Extrapolation for large-batch training in deep learning. In Proceedings of the 37th International Conference on Machine Learning, 2020.
[48] Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of SGD in modern over-parametrized learning. In Proceedings of the 35th International Conference on Machine Learning, 2018
[49] Grigory Malinovskiy, Dmitry Kovalev, Elnur Gasanov, Laurent Condat, and Peter Richtarik. From local SGD local fixed-point methods for federated learning. In Proceedings of the 37th International Conference on Machine Learning, 2020.
[50] Dominic Masters and Carlo Luschi. Revisiting small batch training for deep neural networks. arXiv preprint arXiv: 1804.07612. 2018
[51] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model 2018.
of large-batch training. arXiv preprint
[52] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-cffiicc learning of deep networks from decentralized data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, volume 54, 2017.
[53] H. Brendan McMahan and Matthew J. Streeter. Adaptive bound optimization for online convex optimization. In COLT The 23rd Conference Learning Theory, 2010.
13

---

[54] H. Brendan McMahan, Daniel Ramage, Kunal Talwar, and Li Zhang. Learning differentially private recurrent language models. In International Conference on Learning Representations. 2018.
[55] Tomáš Mikolov, Martin Karafíát, Lukáš Burget, Jan Černocký, and Sanjeev Khudanpur. Recur- rent neural network based language model. In Eleventh Annual Conference of the International Speech Communication Association, 2010.
[56] Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In Proceedings of the 36th International Conference on Machine Learning, 2019.
[57] Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, 2019.
[58] Zachary Nado, Justin M Gilmer, Christopher J Shallue, Rohan Anil, and George E Dahl. A large batch optimizer reality check: Traditional, generic optimizers suffice across batch sizes. arXiv preprint SEKiv:2102.063330 2021.
[59] Takayuki Nishio and Ryo Yonetani. Client selection for federated learning with heterogeneous resources in mobile edge. In IEEE International Conference on Communications (ICC), 2019.
[60] Reese Pathak and Martin Wainwright. FedSplit: An algorithmic framework for fast federated optimization. In Advances in Neural Information Processing Systems, pages 7057-7066, 2020.
[61] Matthias Paulik, Matt Seigel, Henry Mason, Dominic Telaar, Joris Kluivers, Rogier van Dalen, Chi Wai Lau, Luke Carlson, Filip Granqvist, Chris Vandevelde, et al. Federated evaluation and tuning for on-device personalization: System design & applications. arXiv preprint arXiv:2102.08500 2021.
[62] Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Leurning Representations, 2021.
[63] Monica Ribero and Haris Vikalo. Communication-effienee federated learning via optimal client sampling. arXiv preprint arXiv:2007.15117, 2020.
[64] Christopher J Shallue, Jachoon Lee, Joseph Antognini. Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20:1-49 2019.
[65] Elaine Shi, TH Hubert Chan, Eleanor Rieffel, Richard Chow, and Dawn Song. Privacy- preserving aggregation of time-series data. In Proc. NDSS, volume 2. pages 1-17, 2011.
[66] Samuel L. Smith, Pieter-Jan Kindermans, and Quoc V. Le. Don't decay the learning rate, increase the batch size. In Internationall Conference on Learning Representations, 2018.
[67] Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khaz- aeni. Federated learning with matched averaging. In International Conference on Learning Representations, 2020.
[68] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European Conference on Computer Vision (ECCV), pages 3- 19, 2018.
[69] Haibo Yang. Minghong Fang, and Jia Liu. Achieving linear speedup with partial worker partici- pation in non-IID federated learning. In International Conference on Learning Representations, 2021.
[70] Dong Yin, Ashwin Pananjady, Max Lam, Dimitris Papailiopoulos, Kannan Ramchandran, and Peter Bartlett, Gradient diversity: a key ingredient for scalable distributed learning. In Proceedings the Twenty-First International Conference Artificial Intelligence and Statistics, 2018.
[71] Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint 03888, 2017.
14

---

[72] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020.
[73] Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. FedPD: A feder- ated learning framework with optimal rates and adaptivity to non-llD data. arXiv preprint 1144: 2020.
15


Reviews:
[{"review_id": "VuPGksWP425", "reviewer": "Reviewer_5Vr7", "summary": "This work introduces a wide empirical study of cohort sizes in local methods for federated learning problems. The authors explain the advantages and disadvantages of using a large cohort number. They provide connections with batch size in centralized optimization. They highlight 5 main aspects: Catastrophic Training Failures, Diminishing Returns, Generalization Failures, Fairness Concerns, and Decreased Data Efficiency. Also, the authors provide several solutions how to overcome issues led by a large cohort number. Authors provide several sequences of experiments with different models, datasets, and methods. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper raises many questions regarding using large cohort numbers in federated learning problems. They provide many insights that can be applied to many aspects of machine learning in general. These results are novel and relevant to the Federated Learning community. \n\nFirstly, the authors provide a detailed explanation and good motivation of considered problems. Problems, datasets, and methods are described well. \n\nAll experiments have error bars, authors use 5 trials. However, using a large number of trials can be useful. \n\nIn section 5 authors propose different ideas to overcome issues that can appear because of the large cohort number. Unfortunately, they do not provide any theoretical analysis. Firstly, the authors propose learning rate scaling. For a similar reason, we can use the clipping technique (https://arxiv.org/pdf/2005.10785.pdf). \n\nOverall, despite the lack of theoretical analysis, this empirical research has good quality. The authors cover many aspects of using large cohort numbers in federated learning. I believe that this paper is good enough to be published.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Authors can provide theoretical analysis for a subset of approaches that they provided in section 5.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "EBeSx3PwpgU", "reviewer": "Reviewer_SF2P", "summary": "The authors provide a comprehensive empirical evaluation of the effect of the cohort size (i.e., the number of clients sampled in each round) on the convergence and quality of the model learnt by federated learning algorithms. They investigate different versions of the FedOpt algorithm on 5 different model, dataset combinations. Similar to large mini-batch training in centralized settings they find that generalization properties can suffer from large cohort sizes, convergence speedups from increasing the cohort size are diminishing and the algorithm suffers worse data efficiency. In addition, specific to heterogeneous data settings, they show that larger cohort sizes increase the risk of divergence due to extreme outlier gradients.\n\nTowards addressing these challenges the authors propose to use gradient clipping for federated learning to achieve robustness to unbounded gradients, and they propose to consider layer-wise adaptive algorithms such as LARS or Lamb for server-side aggregation. \n\nFinally, they investigate the sources for these phenomena and show that gradients of individual agents are often orthogonal which they think could partially explain the diminishing return phenomena, the data inefficiency, and the advantage of adaptive step size schedules for large cohort training.\n", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well written, the setup is clear and easy to follow. I like the selection of datasets and tasks. Figures are nice, although they could be a bit larger and fill the entire column.\n\nI think the first observation that gradient clipping is important in federated learning to achieve robustness is interesting. I had not seen it being discussed previously. One question I have here is – how often do you observe these convergence issues? Does it also happen on other datasets, or is it just one particular client in the EMNIST training set that performs poorly? In any case I don’t see a reason why this should not be implemented as a default robustness feature in FL algorithms.\n\nRelated to the other claims of the paper I have a couple of questions and concerns that I hope the authors can clarify or comment on. \n\n1) You say that the failure rate increases from 0% to 80% when increasing the cohort size from 10 to 800. Is this evaluated for a fixed number of communication rounds, or for a target accuracy level? If it is a fixed number of rounds is there more going on here than the fact that if there is a small number of bad clients in the training set the probability of sampling it in the course of the algorithm is 80x higher in the larger cohort? (0% for 5 reps could as well be 1%).\n2) In L190 you state that the generalization failure do not occur on EMNIST and Stack Overflow because the cohort size is small compared to the total number of clients. I am not sure I understand this reasoning, you are picking the cohort size in your experiment. If you choose the participation on these tasks to 80% can you observe these effects too? \n4) I think it is interesting to visualize how the test accuracy is distributed across the clients in the test set, but I don’t think any conclusion about fairness issues arising from the cohort size is justified from your results. I would rather say that small cohort sizes have not yet converged after 1500 rounds, and that’s why the spread in accuracy values is smaller. For a fair comparison I think you would need to compare the distributions for different cohort sizes at a fixed test accuracy.\n5) I don’t understand the reasoning in L220. You say that clients with more data are more likely to become stragglers and thus large-cohort training requires more compute time. Isn’t the amount of data per client independent from the cohort size? I would rather think that the probability of having a straggler is larger with more parallel updates and this can increase runtime, but I don’t see the connection to the amount of local data.\n6) I don’t agree with the interpretation of the results in section 4. It is obvious from your results that increasing M leads to smaller updates. But I would argue that the true gradient of the training objective is very small (i.e. the heterogeneous updates from the clients approximately sum to zero which is to be expected once the algorithm is converging) and the norm of the gradient you plot is larger for small cohort sizes due to the variance in the individual terms of the sum. Your figures show that even after the algorithm has converged the norm of the update is almost constant and what we are left with is the fluctuations due to heterogeneity in the objectives. Do you think there is more going on here with the inverse square root rule you propose, or is it just reflecting this variance term? \n7) You say that the learning rate adjustment has mixed efficacy because the learning rate schedule is often too aggressive and you conclude that applying learning rate scaling at the server may not directly improve large-cohort training. Have you tried tuning the learning rate individually for the different cohort sizes to see if it is an issue of the schedule or whether there is no real potential in tuning the server-side learning rate?\n8) You talk about large cohort sizes, and that it might be good to dynamically increase it over time. I think it would be helpful to emphasize that this setting is quite different from mini-batch training because, as you said in the motivation, you cannot control client’s availability in federated learning. It seems that a more natural question would be to ask how you can dynamically react to a changing number of available clients.\n\nOverall, I appreciate the effort of the authors to conduct this extensive study. However, currently there are too many claims and conclusions in the paper that I am not convinced by. If the authors can pursuade me otherwise I am happy to increase my score.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"ethical_concerns": "NA", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "BvPg82_Ys6a", "reviewer": "Reviewer_Lnmy", "summary": "This paper investigates the impact of the size of the sampled client (called cohort size) on the quality of the learned model and the training dynamics in FL via extensive empirical study. It demonstrates several particular issues are were not observed in the centralized large-batch training. It also provides an empirical explanation and some heuristic approaches to solve the issues based on the characteristics of federated training dynamics.\n\n", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The problem of the cohort size in FL is well-motivated, as similar issues have also occurred in classical large-batch training. The experiments are solid. The near-orthogonality observations are interesting and shed light on the impact of adaptive optimizers.\n\nI hope to see more future work that solves the problem of large batch size training (still open even for centralized cases). \n \nI suggest the authors consider using the static Batch Normalization technique proposed in the following reference, which significantly outperforms other normalization methods (including the currently used GroupNorm) in FL settings.\n\"HeteroFL: Computation and Communication Efficient Federated Learning for Heterogeneous Clients\" by Diao et al.\n \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1 hour"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The work is a well written and well executed empirical study of phenomena associated with large cohort sizes when training practical cross-device FL models. 

All reviewers ultimately recommended to accept this paper for publication, and I concur. Moreover, I enjoyed the rich discussion between the reviewers and the authors. I have read the paper in detail myself, and have enjoyed the experience. However, since no theoretical explanation is offered, and because the scope of the experiments is necessarily limited, I am not fully convinced about the robustness of the conclusions. Because of this, I consider this paper to be a borderline accept. 

---

A few caveats:

- I believe that some of these observations may have natural theoretical explanations from existing/known theory (granted, under some constraining assumptions on the models and losses trained), which the authors are not exploring. 

- Given the experimental setup, whatever the results would be, *something* would certainly be observed, and a commentary on that something could always be made. Would that mean that a pattern or a phenomenon was discovered? Maybe, but not necessarily. Many more experiments across a vast array of model types and sizes would be needed for a more convincing argument to be made. It is very hard to say whether the presented observations are robust enough - would they be observed for other models and datasets? What would happen, for example, if you experimented with simple linear models that lead to convex optimization problems? What happens if local steps are removed and one uses compressed communication instead? 

Having said that, the observations, however fragile, are useful for further theoretical and empirical studies, and when replicated by other researchers in other settings, may serve as an inspiration to design practical mitigation strategies, such as those outlined in the work. 
Decision: Accept (Poster)