PaperlayerView exact revision on arXiv

Source-linked AI summary

Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning

Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank J. Reddi, Sebastian U. Stich, Ananda Theertha Suresh

arXiv:2008.03606v2cs.LGcs.DCmath.OCstat.ML

TL;DR

Cross-device federated learning must optimize across heterogeneous, transient clients, where local steps create client drift and prior methods do not uniformly improve on centralized training. MIME adapts arbitrary centralized optimizers using server-level state and local control-variate corrections, and its momentum-variance-reduction variant is proved to beat server-only lower bounds. Experiments report faster convergence than FEDAVG, while performance can suffer on tasks with sparse client gradients.

  • Problem

    Cross-device FL faces client drift from heterogeneous local data, while existing approaches lack a uniform improvement over centralized methods and do not fully fit transient, stateless clients.

  • Method

    MIME adapts arbitrary centralized optimizers by computing optimizer state at the server and applying it locally with an SVRG-style correction during client updates.

  • Results

    MIME with momentum-based variance reduction is proved to beat centralized lower bounds, and experiments show MIME and MIMELITE converge faster than FEDAVG.

  • Takeaways & Limitations

    Local steps can yield asymptotic communication improvements when client Hessian dissimilarity δ is sufficiently smaller than smoothness L.

  • Takeaways & Limitations

    MIME’s SVRG correction hinders performance on tasks with sparse client gradients, and its theoretical speed-up depends on Hessian similarity.

Abstract

from arXiv · show

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.

1 Introduction

Cross-device federated learning faces client drift and communication barriers, while existing methods do not fully support its transient, stateless clients. MIME addresses this by adapting centralized optimizers through server-level state and local corrections, with theoretical and empirical evidence of improved convergence.

  • Problem: Cross-device FL involves extremely many, resource-constrained clients, often preventing a full pass over client data during training.Clients communicate over unreliable networks, creating challenges distinct from cross-silo FL.
  • Problem: FEDAVG’s multiple local SGD steps can overfit atypical client data, causing client drift, slower convergence, and potentially higher communication.The same setting lacks principled ways to incorporate momentum, adaptivity, and clipping into client updates.
  • Method: MIME adapts arbitrary centralized optimizers by combining server-level optimizer state with an SVRG-style correction during every local client update.The state remains fixed across local steps to avoid overfitting to individual clients.
  • Theory: MIME’s convergence theory translates convergence of a generic centralized algorithm into the federated setting while reducing client drift.The framework is designed for cross-device clients that may be stateless or participate only once.
  • Results: MIME with momentum-based variance reduction is proved to beat centralized lower bounds, while experiments report faster convergence than FEDAVG for MIME and MIMELITE.The paper presents this as the first such asymptotic speed-up from local steps in federated learning.

2 Problem setup

The paper formalizes cross-device FL as optimization over heterogeneous, transient clients under a communication-round constraint. Its assumptions characterize gradient and Hessian dissimilarity, along with smoothness and lower-bounded objectives.

  • Setting: Cross-device FL has many transient devices, few local samples, heterogeneous data, and communication rounds as a key optimization metric.Clients may connect at most once, so client-side state cannot be stored.
  • Objective: The objective is to minimize the global loss using as few client-server communication rounds as possible.Each round samples a subset of clients and updates the server model.
  • Assumptions: Assumption A1 bounds inter-client gradient variance by G2, capturing gradient dissimilarity across clients.The analysis also considers stochastic gradients sampled from local data.
  • Assumptions: Assumption A2 bounds Hessian variance by δ, whereas standard smoothness bounds client Hessians by L.The paper notes A2 is weaker than L-smoothness and expects δ ≪ L when clients are similar.
  • Assumptions: The global objective is assumed L-smooth and bounded below by f⋆.These conditions support the convergence analysis.

3 Mime framework

MIME decomposes a centralized optimizer into parameter and state updates, computes global optimizer statistics at the server, and applies them during local client steps. MIME uses an SVRG-style correction, while MIMELITE uses local minibatch gradients alone.

  • Framework: The framework combines a base centralized optimizer, a server-state computation rule, and local client updates.It supports base algorithms with internal state such as momentum.
  • Base algorithm: The base optimizer separates parameter updates U from optimizer-state updates V, allowing algorithms such as momentum and Adam to be represented generically.U is kept linear in the gradient, while V may be more complex.
  • Base algorithm: For momentum SGD, the parameter and state updates both combine the current gradient with the previous momentum state.The momentum state is denoted mt and uses coefficient β.
  • Server state: MIME computes optimizer state from gradients at the server parameters and keeps that state fixed throughout local client updates.This keeps the state unbiased and representative of the global objective.
  • Local updates: MIMELITE uses local minibatch gradients, whereas MIME adds an SVRG-style correction using a full local-batch gradient at the server parameters.The correction reduces minibatch sampling noise and has faster theoretical rates, while MIMELITE closely matches MIME empirically in deep learning.
  • Implementation: Practical implementations weight client averages by local datapoint counts and perform K epochs rather than K steps.

4 Theoretical analysis of Mime

MIME preserves centralized convergence under suitable assumptions and, with momentum-based variance reduction, can surpass server-only communication lower bounds. The speed-up depends on Hessian similarity, while empirical comparisons show strong performance but identify scope limitations.

  • Generic reduction: MIME and MIMELITE preserve convergence of a centralized base optimizer when its update is robust to small perturbations.The induced perturbation is described as O(˜η2), enabled by linearity of the parameter update.
  • Generic rates: With sufficient local steps K ≥ G2/σ2, MIME matches server-only SGD rates, while MIMELITE matches the asymptotic term and becomes FEDAVG with SGD.
  • Generic rates: MIME and MIMELITE are analyzed for Adam under bounded coordinate gradients, with β1 = 0 and a sufficiently large client batch S ≥ G2/ϵ.The Adam regularization parameter ε0 is distinct from the target accuracy ϵ.
  • Empirical evidence: MIME is theoretically fastest or matches optimal statistical rates, while experiments report MIME ≈ MIMELITE > FEDAVG and about 7× faster performance than FEDAVG in one comparison.The comparison also notes that SCAFFOLD can be disadvantaged by stale client control variates in cross-device FL.
  • Caveats: The theory has scope boundaries: some bounds may be vacuous with very large client populations, Adam needs large batches, and speed-ups rely on Hessian similarity.In the worst case δ ≈ L, all methods are comparable.
  • Momentum variance reduction: MimeMVR and MimeLiteMVR require K ≥ L/δ local steps and can beat server-only lower bounds when δ is sufficiently smaller than L.Their rates depend on Hessian dissimilarity δ rather than smoothness L.
  • Momentum variance reduction: When δ ≈ 0, the methods require O(1/ϵ) communication rounds, showing that local steps can help when heterogeneity is smaller than optimization difficulty.
  • Momentum variance reduction: MimeMVR uses momentum β of order 1 − O(TG2)^−2/3, approaching 1 as T increases.The paper presents this as a theoretical demonstration of the usefulness of large momentum in federated learning.

5 Experimental analysis on real world datasets

Experiments on real-world federated datasets show that Mime and MimeLite generally outperform FedAvg, while momentum and adaptivity improve performance. Results depend on task structure, communication constraints, and gradient sparsity.

  • Overall findings: Mime and MimeLite consistently outperform FedAvg across experiments, with momentum and adaptivity improving performance.The study evaluates EMNIST62, Shakespeare, and StackOverflow tasks using multiple base optimizers.
  • Ablation and comparative study: 7× faster than FedAvg, Mime and MimeLite achieve nearly identical performance on EMNIST62 MLP training with SGD+momentum.The comparison uses 10 local epochs and 1,000 training rounds.
  • Ablation and comparative study: Mime benefits more from momentum than FedAvg, whose improvement is only slight in the EMNIST62 MLP comparison.Mime applies momentum during every client update rather than only at the server level.
  • Ablation and comparative study: Mime ≈ MimeLite > FedAvg > SCAFFOLD > FedProx in the EMNIST62 MLP comparison.SCAFFOLD is slower in this cross-device setup because its client control variates become stale, while FedProx regularization slows convergence.
  • Large scale comparison with equal server and client communication: Under equalized communication, Mime and MimeLite are competitive with FedAvg and achieve the best results on EMNIST62 image-classification tasks with Adam.FedAvg appears more unstable with Adam, while MIME methods use fewer sampled clients to offset additional optimization-state communication.
  • Limitations and discussion: Mime can underperform on sparse-gradient tasks because its SVRG correction hinders performance.The discussion identifies sparse-gradient adaptation as an avenue for improvement.

6 Conclusion

The paper formalizes cross-device federated optimization through MIME and provides theoretical and empirical evidence for its effectiveness. MimeMVR additionally achieves asymptotic improvements over centralized lower bounds.

  • Conclusion: MIME provides a theoretically justified framework for addressing client heterogeneity in cross-device federated learning.The framework is presented as a formal study of the cross-device setting.
  • Conclusion: MIME adapts arbitrary centralized algorithms such as Adam without additional hyper-parameters.The conclusion identifies this adaptability as a central property of the framework.
  • Conclusion: MimeMVR beats centralized lower bounds, showing that additional local steps can yield asymptotic improvements for the first time.The conclusion connects this result to the analysis of non-convex and decentralized optimization.

A How momentum can help reduce client drift

FEDAVG’s multiple local updates and biased server momentum can drive the model toward an incorrect point. MIME instead applies fixed global momentum locally to reduce update bias while retaining multiple local steps.

  • Server-only momentum: Server-only momentum uses unbiased gradients but can be communication-intensive because it takes no local steps.The strategy may spend substantial time waiting for communication rather than computing gradients.
  • FEDAVG and client drift: FEDAVG’s multiple local SGD steps overfit atypical client data, producing client drift and potentially worse performance than server-only training.Its pseudo-gradient and resulting server momentum are biased because local updates optimize f_i(x) rather than the global objective f(x).
  • MIME approach: MIME computes unbiased global momentum and applies it at every local update to address bias in both momentum and client updates.The momentum remains fixed during local steps; there is no client-local momentum.
  • MIME approach: By combining local gradients with fixed global momentum, MIME preserves the communication benefits of multiple local steps while preventing client drift.Figure 2 illustrates the contrast between FEDAVG and MIME for two clients, three local steps, and β = 0.5.

B Proof sketch

The proof sketch separates MIME’s convergence analysis into statistical variance, optimization bias, and SVRG-correction components. It links these components to momentum, local-step usefulness under Hessian similarity, and reduced within-client variance.

  • Statistical term: Momentum reduces the variance of local updates by incorporating past gradients computed on different clients.The analysis states that the variance can decrease from G^2 to aG^2 under the described momentum construction.
  • Optimization term: More local updates can improve the optimization term when client Hessians are similar, with the suggested choice of local updates proportional to L/δ.The proof sketch contrasts this with the server-only optimization term and notes that δ ≪ L can make the local-step term smaller.
  • Optimization term: Only the first gradient evaluated at the server parameter is unbiased; later local gradients are biased because the local iterate depends on the same client loss.This bias is analyzed to explain the effect of local updates.
  • Optimization term: Hessian dissimilarity controls the bias and therefore the usefulness of local updates.The proof sketch formalizes this relationship through Lemma 3.
  • SVRG correction: When many clients use minibatches, the SVRG correction reduces within-client variance while retaining variance of order G^2/S.This conclusion uses smoothness and assumes the local iterate remains close to the server iterate within one round.

C.2 Description of large scale experiments

The experiments evaluate FEDAVG, MIME, and MIMELITE across realistic federated datasets and models, with extensive hyperparameter sweeps. The study emphasizes communication rounds while acknowledging MIME’s higher per-round communication and computational costs.

  • Datasets and models: The study runs five simulations across three real-world federated datasets: EMNIST62, Shakespeare, and StackOverflow.Tasks include linear classification, MLP and CNN image models, a Shakespeare charRNN, and a StackOverflow LSTM language model.
  • Datasets and models: Client partitions use original-author metadata, creating naturally heterogeneous datasets with StackOverflow closest to the cross-device setting.Shakespeare is described as closest to the cross-silo setting based on average client participation.
  • Practicality: The experiments optimize for communication rounds, although MIME requires twice the bits per round and server-only methods require less computation.The authors justify this focus by treating connection latency as the dominant cross-device cost, while identifying compression as future work.
  • Algorithms and tuning: The evaluation includes light and heavy hyperparameter sweeps, with the light setup using a 9× learning-rate grid and the heavy setup using a 567× sweep.Additional grids tune client and server learning rates and, for adaptive methods, ε0.
  • Comparison with previous results: Under low-tuning defaults, FedAdam often underperforms FedSGDm, whereas MimeAdam and MimeLiteAdam are described as stable and often better than their SGD counterparts.The paper attributes FedAdam’s behavior in this setup to insufficient tuning of parameters such as ε.

D Stability of methods to hyper-parameters

The experiments examine how method performance changes with server learning rates and optimizer settings, while the theory develops assumptions and lemmas used to analyze federated updates. Momentum is less sensitive to server learning rate than SGD, and default server learning rate 1 works well across the non-adaptive methods tested.

  • Hyper-parameter stability: Momentum methods are more insensitive to server learning rate than SGD methods, with server learning rate 1 working well for all methods.Figure 4 compares FedAvg, Mime, and MimeLite using SGD and momentum with β = 0.9 on EMNIST62 with CNN.
  • Assumptions: The analysis distinguishes inter-client gradient dissimilarity from intra-client gradient variance and treats them as separate assumptions.The paper identifies G2 as inter-client gradient dissimilarity and σ2 as intra-client variance, noting that typically σ2 ≤ G2.
  • Assumptions: Hessian dissimilarity is formalized through δ, which bounds differences in client and global Hessians and yields gradient-difference bounds between points.The paper rewrites the condition using Ψi(z) = fi(z; ζ) − f(z), obtaining ∥∇2Ψi(z)∥ ≤ δ and a corresponding gradient-smoothness relation.
  • Limitations: The weaker Hessian-similarity notion remains an open direction because the analysis faces technical issues when the comparison point is random and client-dependent.The paper explicitly identifies extending results to this weaker notion and proving tight non-convex lower bounds as future theoretical challenges.
  • Proof tools: The update analysis separates mean and variance using relaxed triangle inequalities, variance decompositions, and martingale-difference arguments.These lemmas support bounding accumulated local-update deviations and exploit zero-mean cross terms when the noise forms a martingale difference sequence.
  • Federated update structure: Each round samples a fixed client set, communicates server parameters and an average gradient statistic, performs K local updates, and aggregates the resulting client parameters.The average gradient statistic is required by Mime but not MimeLite and is assumed unbiased over the sampled clients.

F.1 Proof of Theorem I (generic reduction)

The generic reduction proof expresses Mime and MimeLite server updates as centralized base-algorithm updates plus a local-step error. Smoothness and variance assumptions then bound this error and support convergence analyses for multiple optimizers.

  • Generic reduction: For a linear updater, the Mime and MimeLite server updates can be written as centralized updates perturbed by a defining error term.The proof introduces et from the difference between gradients evaluated at local iterates and at the round’s server parameters.
  • Error bounds: Mime’s SVRG correction improves the client-drift error bound relative to MimeLite under smoothness, variance, and gradient-dissimilarity assumptions.The proof separately bounds MimeLite error and then states that the additional correction in Mime improves the bound for Mime.
  • Generic reduction: Mime and MimeLite mimic the centralized base algorithm with perturbation magnitude O(˜η2), enabling centralized convergence analyses to transfer to federated versions.The paper attributes the reduction to linearity of the parameter update and the separation of parameter and optimizer-state updates.
  • SGD convergence: For SGD-based variants, one-round progress bounds combine smoothness with PL-inequality assumptions and include error terms proportional to L˜η2G2 or 3L˜η2G2.The corresponding bounds are stated for MimeLiteSGD and MimeSGD with ˜η ≤ 1/(12BL).
  • Momentum variance reduction: The momentum-variance-reduction extension is designed to reduce update variance, and its correction becomes small when successive parameters are close.When all client functions are identical, the update reduces to ordinary gradient descent, while the method otherwise maintains an exponential moving average of variance terms.
  • Bias in local updates: Client-dependent local iterates make the expected local gradient biased, and controlling this bias is central to the analysis of client drift.The paper identifies dependence of local gradients on the sampled client as the underlying source of client drift.

G.5 Final convergence rates

The final theory derives convergence guarantees for MimeMVR and MimeLiteMVR under progressively different smoothness and Hessian-variance assumptions. The proofs select step sizes and momentum parameters to simplify the resulting rates.

  • MimeMVR: Theorem V establishes a convergence result for MimeMVR under assumptions (A1) and (A2).The theorem is stated for MimeMVR with a specified step size, followed by a proof that unrolls one-round progress over T rounds.
  • Proof setup: The MVR analysis controls initialization through either time-varying step sizes or a large initial momentum accumulation.The paper follows the latter approach for simplicity and notes that both choices preserve the total complexity up to additional proof complexity.
  • Rate simplification: Choosing the prescribed step size and momentum parameter simplifies the convergence rates obtained for MimeMVR.The paper presents simplified rates after selecting these parameters and then concludes the theorem proof.
  • MimeLiteMVR: Theorem VI establishes a convergence result for MimeLiteMVR under assumptions (A1) and (A2*).The proof states that MimeLiteMVR follows the MimeMVR argument while replacing the G2/S term with G2 + σ2 and requiring a weaker Hessian variance condition.
Loading 2008.03606v2…