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 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.