SCAFFOLD: Stochastic Controlled Averaging for On-Device Federated Learning

Federated learning is a key scenario in modern large-scale m achine learning. In that scenario, the training data remains distributed over a larg e number of clients, which may be phones, other mobile devices, or network sensors and a centr alized model is learned without ever transmitting client data over the network. The standar d optimization algorithm used in this scenario is Federated A veraging (FedA vg). However, when client data is heterogeneous, which is typical in applications, FedA vg does not a dmit a favorable convergence guarantee. This is because local updates on clients can drif t apart, which also explains the slow convergence and hard-to-tune nature of FedA vg in pract ice. This paper presents a new Stochastic Controlled A veraging algorithm ( SCAFFOLD) which uses control variates to reduce the drift between different clients. We prove that the algorithm requires significantly fewer rounds of communication and benefits from favorable co nvergence guarantees.