[Paper Review] Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning

Sai Praneeth Karimireddy, Martin Jaggi+5 more|arXiv (Cornell University)|Aug 8, 2020

Privacy-Preserving Technologies in Data68 references94 citations

TL;DR

Mime provides a general framework to adapt centralized optimization algorithms to cross-device federated learning, reducing client drift and enabling asymptotic speed-ups with local updates, including momentum-based variance reduction.

ABSTRACT

Federated learning (FL) is a challenging setting for optimization due to the heterogeneity of the data across different clients which gives rise to the client drift phenomenon. In fact, obtaining an algorithm for FL which is uniformly better than simple centralized training has been a major open problem thus far. In this work, we propose a general algorithmic framework, Mime, which i) mitigates client drift and ii) adapts arbitrary centralized optimization algorithms such as momentum and Adam to the cross-device federated learning setting. Mime uses a combination of control-variates and server-level statistics (e.g. momentum) at every client-update step to ensure that each local update mimics that of the centralized method run on iid data. We prove a reduction result showing that Mime can translate the convergence of a generic algorithm in the centralized setting into convergence in the federated setting. Further, we show that when combined with momentum based variance reduction, Mime is provably faster than any centralized method--the first such result. We also perform a thorough experimental exploration of Mime's performance on real world datasets.

Motivation & Objective

• Mitigate client drift in cross-device federated learning.
• Adapt arbitrary centralized optimization algorithms (e.g., momentum, Adam) to FL.
• Provide convergence guarantees that Mime preserves or improves convergence compared to centralized baselines.
• Show that momentum-based variance reduction with Mime yields faster convergence than server-only methods.
• Validate Mime and a lighter variant MimeLite with empirical experiments on real datasets.

Proposed method

• Introduce Mime, a framework that uses server-side optimizer state and SVRG-style corrections to make each local client update mimic the corresponding centralized update on i.i.d. data.
• Represent a base centralized optimizer as a pair of components: an update function U(g, s) and an optimizer-state update V(g, s).
• Compute global optimizer state at the server using full gradients at server parameters, and keep this state fixed during local client updates.
• During local updates, apply the base optimizer’s update U to biased or corrected gradients to reduce drift, using either SVRG-like correction (Mime) or plain local minibatch gradients (MimeLite).
• Provide a generic convergence reduction: if the base optimizer converges robustly, Mime inherits convergence with local steps; and analyze specific momentum-based variance reduction (MVR) to beat server-only lower bounds under Hessian dissimilarity.

Experimental results

Research questions

• RQ1 Can arbitrary centralized optimizers be faithfully transferred to the cross-device FL setting without inducing client drift?
• RQ2 Under what conditions do Mime and MimeLite preserve or improve convergence relative to centralized methods and server-only FL baselines?
• RQ3 Does incorporating momentum-based variance reduction (MVR) in Mime yield asymptotic speed-ups over server-only methods in federated learning?
• RQ4 How do heterogeneity measures (G^2, δ) affect the convergence rates and potential speed-ups of Mime/MimeLite?
• RQ5 What is the empirical performance of Mime and MimeLite on real federated datasets compared to FedAvg and other FL baselines?

Key findings

• Mime reduces client drift and translates the convergence of a generic centralized algorithm to the federated setting.
• Mime with MVR can beat lower bounds for server-only methods when Hessian dissimilarity δ is small compared to the smoothness L.
• Mime and MimeLite converge faster than FedAvg in empirical experiments on real-world datasets.
• Applying momentum locally at every client update yields stronger performance gains than server-side momentum alone.
• MimeLite closely matches Mime in practice and can achieve similar acceleration benefits with simpler correction.
• Theoretical results include a reduction theorem and conditions under which Mime speeds up convergence asymptotically in cross-device FL.