Federated Hyperparameter Tuning: Challenges, Baselines, and Connections to Weight-Sharing

1 Introduction . Federated learning ( FL ) is a popular distributed computational setting where training is performed locally or privately [ 30 , 36 ] and where hyperparameter tuning has been identified as a critical problem [ 18 ] . Although general hyperparameter optimization has been the subject of intense study [ 3 , 16 , 26 ] , several unique aspects of the federated setting make tuning hyperparameters especially challenging . However , to the best of our knowledge there has been no dedicated study on the specific challenges and solutions in federated hyperparameter tuning . In this work , we first formalize the problem of hyperparameter optimization in FL , introducing the following three key challenges : 1 . Federated validation data : In federated networks , as the validation data is split across devices , the entire dataset is not available at any one time ; instead a central server is given access to some number of devices at each communication round , for one or at most a few runs of local training and validation . Thus , because the standard measure of complexity in FL is the number of communication rounds , computing validation metrics exactly dramatically increases the cost . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 2 . Extreme resource limitations : FL applications often involve training using devices with very limited computational and communication capabilities . Furthermore , many require the use of privacy techniques such as differential privacy that limit the number times user data can be accessed . Thus we can not depend on being able to run many different configurations to completion . 3 . Evaluating personalization : Finally , even with non-federated data , applying common hyperparameter optimization methods to standard personalized FL approaches ( such as finetuning ) can be costly because evaluation may require performing many additional training steps locally . With these challenges1 in mind , we propose reasonable baselines for federated hyperparameter tuning by showing how to adapt standard non-federated algorithms . We further study the challenge of noisy validation signal due to federation , and show that simple state-estimation-based fixes do not help . Our formalization and analysis of this problem leads us to develop FedEx , a method that exploits a novel connection between hyperparameter tuning in FL and the weight-sharing technique widely used in neural architecture search ( NAS ) [ 4 , 34 , 40 ] . In particular , we observe that weight-sharing is a natural way of addressing the three challenges above for federated hyperparameter tuning , as it incorporates noisy validation signal , simultaneously tunes and trains the model , and evaluates personalization as part of training rather than as a costly separate step . Although standard weightsharing only handles architectural hyperparameters such as the choice of layer or activation , and not critical settings such as those of local stochastic gradient descent ( SGD ) , we develop a formulation that allows us to tune most of these as well via the relationship between local-training and fine-tuningbased personalization . This make FedEx a general hyperparameter tuning algorithm applicable to many local training-based FL methods , e.g . FedAvg [ 36 ] , FedProx [ 31 ] , and SCAFFOLD [ 19 ] . In Section 4 , we next conduct a theoretical study of FedEx in a simple setting : tuning the client step-size . Using the ARUBA framework for analyzing meta-learning [ 20 ] , we show that a variant of FedEx correctly tunes the on-device step-size to minimize client-averaged regret by adapting to the intrinsic similarity between client data . We improve the convergence rate compared to some past meta-learning theory [ 20 , 25 ] while not depending on knowing the ( usually unknown ) task-similarity . Finally , in Section 5 , we instantiate our baselines and FedEx to tune hyperparameters of FedAvg , FedProx , and Reptile , evaluating on three standard FL benchmarks : Shakespeare , FEMNIST , and CIFAR-10 [ 5 , 36 ] . While our baselines already obtain performance similar to past hand-tuning , FedEx further surpasses them in most settings examined , including by 2-3 % on Shakespeare . Related Work To the best of our knowledge , we are the first to systematically analyze the formulation and challenges of hyperparameter optimization in the federated setting . Several papers have explored limited aspects of hyperparameter tuning in FL [ 7 , 23 , 38 ] , focusing on a small number of hyperparameters ( e.g . the step-size and sometimes one or two more ) in less general settings ( studying small-scale problems or assuming server-side validation data ) . In contrast our methods are able to tune a wide range of hyperparameters in realistic federated networks . Some papers also discussed the challenges of finding good configurations while studying other aspects of federated training [ 41 ] . We argue that it is critical to properly address the challenges of federated hyperparameter optimization in practical settings , as we discuss in detail in Section 2 . Methodologically , our approach draws on the fact that local training-based methods such as FedAvg can be viewed as optimizing a surrogate objective for personalization [ 20 ] , and more broadly leverages the similarity of the personalized FL setup and initialization-based meta-learning [ 6 , 11 , 17 , 25 ] . While FedEx ’ s formulation and guarantees use this relationship , the method itself is general-purpose 1A further challenge we do not address is that of the time-dependency of federated evaluation , c.f . [ 9 ] . and applicable to federated training of a single global model . Many recent papers address FL personalization more directly [ 13 , 29 , 35 , 43 , 47 ] . This connection and our use of NAS techniques also makes research connecting NAS and meta-learning relevant [ 10 , 33 ] , but unlike these methods we focus on tuning non-architectural parameters . In fact , we believe our work is the first to apply weightsharing to regular hyperparameter search . Furthermore , meta-learning does not have the data-access and computational restrictions of FL , where such methods using the DARTS mixture relaxation [ 34 ] are less practical . Instead , FedEx employs the lower-overhead stochastic relaxation [ 8 , 28 ] , and its exponentiated update is similar to the recently proposed GAEA approach for NAS [ 27 ] . Running NAS itself in federated settings has also been studied [ 12 , 15 , 46 ] ; while our focus is on non-architectural hyperparameters , in-principle our algorithms can also be used for federated NAS . Theoretically , our work makes use of the average regret-upper-bound analysis ( ARUBA ) framework [ 20 ] to derive guarantees for learning the initialization , i.e . the global model , while simultaneously tuning the step-size of the local algorithm . The step-size of gradient-based algorithms has also been tuned on its own in the settings of data-driven algorithm design [ 14 ] and of statistical learning-to-learn [ 45 ] . 2 Federated Hyperparameter Optimization . In this section we formalize the problem of hyperparameter optimization for FL and discuss the connection of its personalized variant to meta-learning . We also review FedAvg [ 36 ] , a common federated optimization method , and present a reasonable baseline approach for tuning its hyperparameters . Global and Personalized FL In FL we are concerned with optimizing over a network of heterogeneous clients i = 1 , . . . , n , each with training , validation , and testing sets Ti , Vi , and Ei , respectively . We use LS ( w ) to denote the average loss over a dataset S of some w-parameterized ML model , for w ∈ Rd some real vector . For hyperparameter optimization , we assume a class of algorithms Alga hyperparameterized by a ∈ A that use federated access to training sets Ti to output some element of Rd . Here by “ federated access '' we mean that each iteration corresponds to a communication round at which Alga has access to a batch of B clients 2 that can do local training and validation . Specifically , we assume Alga can be described by two subroutines with hyperparameters encoded by b ∈ B and c ∈ C , so that a = ( b , c ) and A = B × C. Here c encodes settings of a local training algorithm Locc that take a training set S and initialization w ∈ Rd as input and outputs a model Locc ( S , w ) ∈ Rd , while b sets those of an aggregation Aggb that takes the initialization w and outputs of Locc as input and returns a model parameter . For example , in standard FedAvg , Locc is T steps of gradient descent with step-size η and Aggb takes a weighted average of the outputs of Locc across clients ; here c = ( η , T ) and b = ( ) . As detailed in the appendix , many FL methods can be decomposed this way , including well-known ones such as FedAvg [ 36 ] , FedProx [ 31 ] , SCAFFOLD [ 19 ] , and Reptile [ 39 ] as well as more recent methods [ 1 , 2 , 29 ] . Our analysis and our proposed FedEx algorithm will thus apply to all of them , up to an assumption detailed next . Starting from this decomposition , the global hyperparameter optimization problem can be written as min a∈A n∑ i=1 |Vi|LVi ( Alga ( { Tj } n j=1 ) ) ( 1 ) In many cases we are also interested in obtaining a device-specific local model , where we take a model trained on all clients and finetune it on each individual client before evaluating . A key assumption we make is that the finetuning algorithm will be the same as the local training algorithm Locc used by Alga . This assumption can be justified by recent work in meta-learning that shows that algorithms that aggregate the outputs of local SGD can be viewed as optimizing for personalization using local SGD [ 20 ] . Then , in the personalized setting , the tuning objective becomes min a= ( b , c ) ∈A n∑ i=1 |Vi|LVi ( Locc ( Ti , Alga ( { Tj } n j=1 ) ) ( 2 ) Our approach will focus on the setting where the hyperparameters c of local training make up a significant portion of all hyperparameters a = ( b , c ) ; by considering the personalization objective we will be able to treat such hyperparameters as architectural and thus apply weight-sharing . Algorithm 1 : Successive halving algorithm ( SHA ) applied to personalized FL . For the non-personalized objective ( 1 ) , replace LVti ( wi ) by LVti ( wa ) . For random search ( RS ) with N samples , set η = N and R = 1 . Input : distribution D over hyperparameters A , elimination rate η ∈ N , elimination rounds τ0 = 0 , τ1 , . . . , τR sample set of ηR hyperparameters H ∼ D [ ηR ] initialize a model wa ∈ Rd for each a ∈ H for elimination round r ∈ [ R ] do for setting a = ( b , c ) ∈ H do for comm . round t = τr−1 + 1 , . . . , τr do for client i = 1 , . . . , B do send wa , c to client wi ← Locc ( Tti , wa ) send wi , LVti ( wi ) to server wa ← Aggb ( wa , { wi } Bi=1 ) sa ← ∑B i=1 |Vti|LVti ( wi ) / ∑B i=1 |Vti| H ← { a ∈ H : sa ≤ 1η -quantile ( { sa : a ∈ H } ) } Output : remaining a ∈ H and associated model wa Tuning FL Methods : Challenges and Baselines In the non-federated setting , the objective ( 1 ) is amenable to regular hyperparameter optimization methods ; for example , a random search approach would repeatedly sample a setting a from some distribution overA , run Alga to completion , and evaluate the objective , saving the best setting and output [ 3 ] . With a reasonable distribution and enough samples this is guaranteed to converge and can be accelerated using early stopping methods [ 26 ] , in which Alga is not always run to completion if the desired objective is poor at intermediate stages , or by adapting the sampling distribution using the results of previous objective evaluations [ 44 ] . As mentioned in the introduction , applying such methods to FL is inherently challenging due to 1 . Federated validation data : Separating data across devices means we can not immediately get a good estimate of the model ’ s validation performance , as we only have access to a possibly small batch of devices at a time . This means that decisions such as which models to flag for early stopping will be noisy and may not fully incorporate all the available validation signal . 2 . Extreme resource limitations : As FL algorithms can take a very long time to run in-practice due to the weakness and spotty availability of devices , we often can not afford to conduct many training runs to evaluate different configurations . This issue is made more salient in cases where we use privacy techniques that only allow a limited number of accesses to the data of any individual user . 3 . Evaluating personalization : While personalization is important in FL due to client heterogeneity , checking the performance of the current model on the personalization objective ( 2 ) is computationally intensive because computing may require running local training multiple times . In particular , while regular validation losses require computing one forward pass per data point , personalized losses require several forward-backward passes , making it many times more expensive if this loss is needed to make a tuning decision such as eliminating a configuration from consideration . Despite these challenges , we can still devise sensible baselines for tuning hyperparameters in FL , most straightforward of which is to use a regular hyperparameter method but use validation data from a single round as a noisy surrogate for the full validation objective . Specifically , one can use random search ( RS ) —repeatedly evaluate random configurations—and a simple generalization called successive halving ( SHA ) , in which we sample a set of configurations and partially run all of them for some number of communication rounds before eliminating all but the best 1η fraction , repeating until only one configuration remains . Note both are equivalent to a “ bracket ” in Hyperband [ 26 ] and their adaptation to FL is detailed in Algorithm 1 . 2For simplicity the number of clients per round is fixed , but all methods can be easily generalized to varying B . As shown in Section 5 , SHA performs reasonably well on the benchmarks we consider . However , by using validation data from one round it may make noisy elimination decisions , early-stopping potentially good configurations because of a difficult set of clients on a particular round . Here the problem is one of insufficient utilization of the validation data to estimate model performance . A reasonable approach to use more is to try some type of state-estimation : using the performance from previous rounds to improve the noisy measurement of the current one . For example , instead of using only the most recent round for elimination decisions we can use a weighted sum of the performances at all past rounds . To investigate this , we study a power decay weighting , where a round is discounted by some constant factor for each time step it is in the past . We consider factors 0.0 ( taking the most recent performance only , as before ) , 0.5 , and 1.0 ( taking the average ) . However , in Figure 2 we show that incorporating more validation data this way than is used by Algorithm 1 by default does not significantly affect results . Thus we may need a better algorithm to use more of the validation signal , most of which is discarded by using the most recent round ’ s performance . We next propose FedEx , a new method that does so by using validation on each round to update a client hyperparameters distribution used to sample configurations to send to devices . Thus it alleviates issue ( 1 ) above by updating at each step , not waiting for an elimination round as in RS or SHA . By simultaneously training the model and tuning ( client ) hyperparameters , it also moves towards a fully single-shot procedure in which we only train once ( we must still run multiple times due to server hyperparameters ) , which would solve issue ( 2 ) . Finally , FedEx addresses issue ( 3 ) by using local training to both update the model and to estimate personalized validation loss , thus not spending extra computation on this more expensive objective .

Summary: This paper proposes a general hyperparameter tuning framework for federated learning based on the weight-sharing technique. In contrast to tuning architectural parameters as in NAS, this paper focus on tuning non-architectural parameters. The proposed framework is general enough to be applied to many FL algorithms including FedAvg, FedProx, and SCAFFOLD. A theoretical study is offered for tuning client step-size. Empirical experiments on Shakespeare, FEMNIST, and CIFAR-10 show promising results.