Federated Learning Using Variance Reduced Stochastic Gradient for Probabilistically Activated Agents

This paper proposes an algorithm for Federated Learning (FL) with a two-layer structure that achieves both variance reduction and a faster convergence rate to an optimal solution in the setting where each agent has an arbitrary probability of selection in each iteration. The first layer of our algorithm corresponds to the model parameter propagation across agents done by the server. In the second layer, each agent does its local update with a stochastic and variance-reduced technique called Stochastic Variance Reduced Gradient (SVRG). We leverage the concept of variance reduction from stochastic optimization when the agents want to do their local update step to reduce the variance caused by stochastic gradient descent (SGD). The special attention in this paper is on FL operation where the agents’ participation in the update process in each round is probabilistic and non-uniform. We provide a convergence bound for our algorithm which improves the rate from $O\left( {\frac{1}{{\sqrt K }}} \right)$ to $O\left( {\frac{1}{K}} \right)$ by using a constant step-size when the cost is strongly convex. For non-convex costs, we establish a $O\left( {\frac{1}{{\sqrt[3]{{{K^2}}}}}} \right)$ to a stationary point using a vanishing stepsize. We demonstrate the performance of our algorithm using numerical simulations.