Federated Learning with Incrementally Aggregated Gradients
Aritra Mitra, Rayana H. Jaafar, George J. Pappas, Hamed Hassani · 2021 60th IEEE Conference on Decision and Control (CDC) · 2021
We consider the standard federated learning (FL) framework where a set of clients coordinate with a central server to train a statistical model. In a single-machine centralized setting, it is well-known that for smooth and strongly convex finite-sum optimization problems, one can design algorithms that guarantee exact linear convergence to the global minimum without computing full (batch) gradients at every iteration. Despite its popularity, an analog of the above result does not exist in FL. Motivated by this gap, we consider a setting where the local loss function of each client can be expressed as a finite sum of smooth component functions. For this setting, we propose a novel computationally-efficient FL algorithm called FedTrack that rests on two key ideas: (i) using the most recently communicated versions of the clients’ gradients in the local update rule, and (ii) incrementally aggregating gradients of the component functions of each client. While the first idea serves to overcome the effect of heterogeneity across the clients’ local loss functions, the second helps to significantly reduce the overall number of gradient computations. For both strongly convex and non-convex local loss functions, we prove that the convergence guarantees of FedTrack match their centralized counterparts (up to constants). In particular, for the strongly convex setting, we show that FedTrack guarantees exact linear convergence to the global minimum deterministically.