On Large-Cohort Training for Federated Learning

1 Introduction . Federated learning ( FL ) [ 52 ] considers learning a 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 a large population of edge devices [ 27 , Table 1 ] . A distinguishing characteristic of cross-device FL is partial participation of the client population : Due to systems constraints such as network size , the server typically only communicates with a subset of the clients at a time1 . For example , in the popular FedAvg algorithm [ 52 ] , at each communication round the server broadcasts its current model to a subset of available clients ( referred to as a 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 population—possibly reducing the number of communication rounds needed to achieve a 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 as 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 1In contrast , cross-silo settings often have a small set of clients , most of which participate in each round [ 27 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 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 largebatch 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 a 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 large-batch training appear in large-cohort training , we also identify a number of new challenges introduced by the federated setting . Optimization for federated learning . Significant attention has been paid towards developing federated optimization techniques . Such work has focused on various aspects , including communicationefficiency [ 5 , 34 , 37 , 52 ] , data and systems heterogeneity [ 25 , 28 , 29 , 39–42 , 67 ] , and fairness [ 26 , 43 ] . We provide a 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 , 73 ] , which may be infeasible in cross-device 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 see whether increasing the cohort size leads to improved convergence in practical , communication-limited settings . Client sampling . A number of works have explored how to select cohorts of a fixed size in crossdevice 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 be able to choose which clients participate in a 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 : min x f ( x ) : = K∑ k=1 pkfk ( x ) , ( 1 ) where K is total number of clients , the pk are client weights satisfying pk ≥ 0 , and fk 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 a sampled cohort could send∇fk ( x ) to the server , and the 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 be desirable in communication-limited 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 a more general framework , FedOpt , introduced by Reddi et al . [ 62 ] that uses both client and server optimization . At each round , the server sends its model x to a cohort of clients C of size M . Each client ck ∈ C performs E epochs of training using mini-batch SGD with client learning rate ηc , producing a local model xk . Each client k ∈ C then communicates their client update ∆k to the server , where ∆k : = xk − x is the difference between the client ’ s local model and the server model . The server computes a weighted average ∆ of the client updates , and updates its own model via x′ = SERVEROPT ( x , ηs , ∆ ) , ( 2 ) where SERVEROPT ( x , ηs , g ) is some first-order optimizer , ηs is the server learning rate , and g is a gradient estimate . For example , if SERVEROPT is SGD , then SERVEROPT ( x , ηs , g ) = x − ηsg . The ∆ in ( 2 ) is referred to as a pseudo-gradient [ 62 ] . While ∆ may not be an unbiased estimate of ∇f , it can serve a somewhat comparable role ( though as we show in Section 4 , there are important distinctions ) . Full pseudo-code of FedOpt is given in Algorithm 1 . Algorithm 1 FedOpt framework Input : M , T E , x1 , ηc , ηs , SERVEROPT , { pk } Kk=1 for t = 1 , · · · , T do The server selects a cohort Ct of M clients uniformly at random , without replacement . The server sends xt to all clients in Ct. Each client k ∈ Ct performs E epochs of mini-batch SGD on fk with step-size ηc . After training , each k ∈ Ct has a local model xtk and sends ∆tk = xt − xtk to the server . The server computes a pseudo-gradient ∆t and updates its model via ∆t = ∑ k∈Ct pk∆ t k∑ k∈Ct pk , xt+1 = SERVEROPT ( xt , ηs , ∆t ) . Algorithm 1 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 is SGD , ηc = 1 , E = 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 1 across multiple datasets , models , and tasks . We discuss the key facets of our experiments below . Datasets , models , and tasks . We use four datasets : CIFAR-100 [ 35 ] , EMNIST [ 13 ] , 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 are partitioned by their author . For Shakespeare , speaking lines in Shakespeare plays are 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-100 , we train a 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 , dropout , and two dense layers . For Shakespeare , we train an RNN with two LSTM layers to perform next-characterprediction . For Stack Overflow , we perform next-word-prediction using an RNN with a single LSTM layer . For full details on the models and datasets , see Appendix A.1 . 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 1 where SERVEROPT 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 1 clients perform E = 1 epochs of training with mini-batch SGD . Their batch size is fixed per-task . We set pk to be the number of examples in client k ’ s dataset . We tune learning rates for all algorithms and models using a held-out validation set : We perform T = 1500 rounds of training with M = 50 , E = 1 for each algorithm and model , varying ηc , ηs over { 10i | − 3 ≤ i ≤ 1 } 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 open-source implementations of all simulations in TensorFlow Federated [ 4 ] 2 . All experiments were conducted using clusters of multi-core CPUs , though our results are independent of wall-clock time and amount of 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 .

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.