On Convergence of Federated Averaging Langevin Dynamics
Abstract: We propose a federated averaging Langevin algorithm (FA-LD) for uncertainty quantification and mean predictions with distributed clients. In particular, we generalize beyond normal posterior distributions and consider a general class of models. We develop theoretical guarantees for FA-LD for strongly log-concave distributions with non-i.i.d data and study how the injected noise and the stochastic-gradient noise, the heterogeneity of data, and the varying learning rates affect the convergence. Such an analysis sheds light on the optimal choice of local updates to minimize the communication cost. Important to our approach is that the communication efficiency does not deteriorate with the injected noise in the Langevin algorithms. In addition, we examine in our FA-LD algorithm both independent and correlated noise used over different clients. We observe there is a trade-off between the pairs among communication, accuracy, and data privacy. As local devices may become inactive in federated networks, we also show convergence results based on different averaging schemes where only partial device updates are available. In such a case, we discover an additional bias that does not decay to zero.
Submission Length: Regular submission (no more than 12 pages of main content)
Changes Since Last Submission: In the second revised version, we have included discussions on the challenges and motivations pertaining to the FA-LD algorithm, as suggested by reviewer 64eW (highlighted using purple color). Furthermore, we have addressed the notation issue raised by reviewer zn8u, which has been marked with orange color.
In the first revised version, we have included discussions, highlighted in red, regarding the scale-invariance property, emphasizing the comparisons with [3]. Additionally, we have provided more comprehensive comparisons with [1] and explored the (sub)-optimal choice of local steps in relation to [2].
[1] Federated Learning with a Sampling Algorithm under Isoperimetry. arXiv:2206.00920v2. 2022
[2] Analysis of Langevin Monte Carlo via Convex Optimization. JMLR. 2019.
[3] User-friendly guarantees for the Langevin Monte Carlo with inaccurate gradient. arXiv:1710.00095v3. Stochastic Processes and their Applications. 2019
Assigned Action Editor: ~Peter_Richtarik1
Submission Number: 1183
Loading