Optimization in Federated Learning

Federated learning (FL) is an emerging branch of machine learning (ML) research, that is examining the methods for scenarios, where individual nodes possess parts of the data, and the task is to form a single common model that fits to the whole distribution. In FL, we generally use mini batch gradient descent for optimizing weights of the models that appears to work very well for federated scenarios. For traditional machine learning setups, a number of modifications has been proposed to accelerate the learning process and to help to get over challenges posed by the high dimensionality and nonconvexity of search spaces of the parameters. In this paper we present our experiments on applying different popular optimization methods for training neural networks in a federated manner. 1 Federated Learning Federated learning (FL) [1] is a new paradigm in Machine Learning (ML), that is dealing with an increasingly important distributed optimization setting, that came into view with the spread of small user devices and applications written for them that can profit from ML. The domain of ML models is often the data collected on the devices, thus, to train these models, one should incorporate the knowledge contained into the learning process. The traditional way for this would be to transfer the information gathered at the users to data centers, where the training takes place, then send back the trained models to the users. That, apart from the obvious privacy concerns, can incur a huge communication overhead along with the need for significant storage and computational resources at the place of centralized training. The idea proposed in [1] is that, instead of moving the training data to centralized location, one could exploit the computational power residing at the user devices and distribute the training process across the participating nodes. 1.1 Distributed Optimization In ML, the goal is to find a model for the training data that minimizes a loss function f that defines how our learned model distribution differs from the empirical distribution. Copyright ©2019 for this paper by its authors. Use permitted under Creative Commons License Attribution 4.0 International (CC BY 4.0). This measure in general case can be formalized as a negative log likelihood. f =−Ex∼pdata [log pmodel(x)] (1) That is, if a given example x is drawn from the training data distribution, what is the probability that it will be present in the same form in the model distribution as well. If the model is for predicting some value(s) y based on a vector of some attributes x this can be rewritten as f (x,y,w) =− log p(y|x;w) (2) , The problem we want to solve is to minimize the loss function f with respect to model parameters w, that is an aggregation of the losses over all available data point as follows: min w∈Rd f (w), where f (w) = 1 n n ∑ i=1 fi(w), (3) where fi(w) def = f (xi,yi,w) denotes the loss on data point i given the parametrization w. In the setup of FL the characteristics of data distribution from which our training examples (xi,yi) will be drawn, are the following: 1. Massively Distributed. Data points are stored across a large number K of nodes. In particular, the number of nodes can be much bigger than the average number of training examples stored on a given node (n/K). 2. Non-IID. Data on each node may be drawn from a different distribution, i.e. the data points available locally are far from being a representative sample of the overall distribution. 3. Unbalanced. Different nodes may vary by orders of magnitude in the number of training examples they hold. If we adapt the objective function (see Eq. 3) to these characteristics, our problem can be defined as introduced in the following paragraphs. We have K nodes and n data points, a set of indices Pk (k ∈ {1, . . . ,K}) of data stored at node k, and nk = |Pk| is the number of data points at Pk. We assume that Pk ∩ Pl = / 0 whenever l 6= k, thus ∑k=1 nk = n. We can then define the local loss for node as Fk(w) def = 1 nk ∑i∈Pk fi(w) Thus the problem to be minimized will become: minw∈Rd f (w) = K ∑ k=1 nk n Fk(w). (4) To solve the problem in (4) the simplest algorithm is FederatedSGD introduced in [2], that is equivalent to minibatch gradient descent over all data, and it is a simple application of distributed synchronous Stochastic Gradient Descent (SGD) [3] for the described setup. Algorithm 1 FederatedSGD 1: procedure SERVER 2: initialize w0 3: for t = 0;1;2; . do 4: for all k in the K nodes in parallel do 5: wk t+1← ClientUpdate(k,wt) 6: end for 7: wt+1 = ∑k=1 nk n w k t+1 8: end for 9: end procedure 10: procedure CLIENTUPDATE(k,w) 11: B← split Pk to set of batches 12: for all b ∈B do 13: w← w−η∇ f (w,b) 14: end for 15: return W 16: end procedure In neural network (NN) optimization, due to the non convexity of the loss functions, the most used methods for optimization of network parameters are gradient based, more specifically the versions of SGD [4]. Gradient descent methods take derivatives of loss function according to the parameters of the model, then move the parameter values in the negative of the gradient. The pure form of SGD samples a random function (e.g a random training data point) it ∈ 1,2, . ,n in iteration t and performs the update: wt+1 = wt −ηt∇ fit (wt), (5) where ηt denotes the learning rate, which is, in the base case, decaying during the learning to enforce convergence. Intuitively, SGD works because evaluating the gradient at a single training example gives an unbiased estimation of derivative of the error function over all the training examples: E[∇ fit (w)] = ∇ f (w). In practice, instead of applying the gradient for w at each example, usually an average of gradients over b randomly chosen examples is used, that are evaluated at the same w. This method is called minibatch gradient descent (MBGD), that better exploits parallel computational capabilities of the hardware. (MBGD is still commonly referred to as SGD) Though SGD/MBGD in the above form is very popular in optimization, the basic approach can sometimes result in very slow learning. To tackle the challenges incurred by high curvature and noisy gradients of the loss function of NN, a range of method has been proposed based on exponentially decaying average of the gradients or on adapting learning rates. [5] In this paper, we investigate the effects of these methods on the performance , of federated training of artificial neural networks. 1.2 Momentum techniques Momentum based techniques [6] use a velocity term during learning, that is an exponentially weighted average over the gradients of the past. v← βv−∇w( 1 m m ∑ i=1 fi(w))