PaperlayerView exact revision on arXiv

Source-linked AI summary

FetchSGD: Communication-Efficient Federated Learning with Sketching

Daniel Rothchild, Ashwinee Panda, Enayat Ullah, Nikita Ivkin, Ion Stoica, Vladimir Braverman, Joseph Gonzalez, Raman Arora

arXiv:2007.07682v2cs.LGstat.ML

TL;DR

Federated learning needs communication-efficient optimization that remains effective with stateless clients and non-i.i.d. data. FetchSGD compresses client gradients with linear Count Sketches, allowing the aggregator to perform momentum and error accumulation centrally. The paper proves convergence guarantees and reports empirical validation across image-recognition and language-modeling tasks.

  • Problem

    Federated training is constrained by slow communication, clients that may participate only once, and non-i.i.d. data.

  • Method

    FetchSGD sends Count Sketches of client gradients to the aggregator, where sketch linearity enables centralized momentum and error accumulation before sparse updates are extracted.

  • Results

    FetchSGD has convergence guarantees for non-i.i.d. L-smooth non-convex objectives and was empirically validated on two image-recognition tasks and one language-modeling task.

  • Takeaways & Limitations

    FetchSGD supports communication-efficient federated optimization without local client state while retaining momentum and handling small local datasets.

  • Takeaways & Limitations

    Theorem 2 gives a worse convergence rate than uncompressed SGD, whereas Theorem 1 matches uncompressed SGD’s rate.

Abstract

from arXiv · show

Existing approaches to federated learning suffer from a communication bottleneck as well as convergence issues due to sparse client participation. In this paper we introduce a novel algorithm, called FetchSGD, to overcome these challenges. FetchSGD compresses model updates using a Count Sketch, and then takes advantage of the mergeability of sketches to combine model updates from many workers. A key insight in the design of FetchSGD is that, because the Count Sketch is linear, momentum and error accumulation can both be carried out within the sketch. This allows the algorithm to move momentum and error accumulation from clients to the central aggregator, overcoming the challenges of sparse client participation while still achieving high compression rates and good convergence. We prove that FetchSGD has favorable convergence guarantees, and we demonstrate its empirical effectiveness by training two residual networks and a transformer model.

1. Introduction

Federated learning must train high-quality models despite slow communication, stateless clients, and non-i.i.d. data. FetchSGD addresses these constraints by compressing client gradients with Count Sketches while moving momentum and error accumulation to the aggregator, with theoretical and empirical validation.

  • Federated learning coordinates training without sending clients’ local data to a central server.
  • Slow connections, stateless clients, and non-i.i.d. data are three constraints that complicate high-quality federated training.Clients may connect at approximately 1Mbps, and some clients participate only once.
  • FetchSGD compresses locally computed gradients with Count Sketch before sending them to the central aggregator.
  • The aggregator maintains momentum and error accumulation sketches, then extracts sparse model updates from the accumulated sketch.
  • FetchSGD requires no local client state and has convergence guarantees for non-i.i.d. L-smooth non-convex objectives.The paper also reports empirical validation on two image-recognition tasks and one language-modeling task using models with 6–125 million parameters.
  • FetchSGD can perform momentum as if on local gradients before compression and scales to small local datasets.

2. Related Work

Prior federated optimization methods trade off communication, client state, non-i.i.d. convergence, or compression quality. FetchSGD complements them by reducing per-round communication while retaining compatibility with stateless clients and non-i.i.d. data.

  • FedAvg: FedAvg reduces total training communication through multiple local SGD steps but provides no per-client savings when a client participates only once.
  • FedAvg: FedAvg’s many local steps can degrade convergence on non-i.i.d. data, with a reported rate of O(K/T) under stated assumptions.
  • FedAvg: FedAvg requires downloading an entire model and uploading an entire model update each communication round.Upload difficulty is amplified by residential connections whose download speeds exceed upload speeds.
  • Gradient Compression: Unbiased gradient compression faces a compression–variance tradeoff that can slow convergence and has had limited empirical success.
  • Gradient Compression: Biased compression methods benefit from error accumulation, but implementing it locally requires client state that may be infeasible when clients participate only once.

3. FetchSGD

FetchSGD sends Count Sketches of client gradients to a central aggregator, which combines them and performs momentum, error accumulation, and sparse update extraction. Count Sketch linearity makes server-side accumulation equivalent to client-side accumulation while avoiding local state.

  • Each participating client computes a stochastic gradient and sends its Count Sketch to the aggregator as a model update.
  • A Count Sketch compresses vectors through randomized projections while allowing high-magnitude elements to be approximately recovered.
  • Count Sketch linearity lets the server compute the sketch of the summed minibatch gradient by aggregating client sketches.
  • The aggregator applies momentum and error accumulation within sketches, then unsketches the accumulated representation and extracts a Top-k update.
  • Server-side error accumulation is equivalent to client-side accumulation because sketching and accumulation use linear operations.

4. Theory

FetchSGD’s theory analyzes convergence under two assumptions: an opaque contraction condition and a more interpretable sliding-heavy-hitters condition. The first matches uncompressed SGD’s convergence rate, while the second supports convergence with sliding-window error accumulation but yields a weaker rate.

  • Theory overview: The theory studies FetchSGD under a strong, opaque contraction assumption and a weaker, more interpretable gradient assumption.The latter assumes heavy signal may appear across a sliding window rather than at every iteration.
  • Scenario 1: Contraction Holds: Under the contraction assumption, FetchSGD converges for L-smooth non-convex functions with bounded stochastic gradients.The theorem uses a step size depending on momentum, smoothness, and the iteration horizon.
  • Scenario 1: Contraction Holds: For large T, Theorem 1’s convergence rate matches that of uncompressed SGD.The expectation in the theorem is over minibatch-sampling randomness.
  • Scenario 1: Contraction Holds: The contraction analysis requires a sufficiently heavy coordinate in the descent-direction and error-accumulation vector, an assumption the paper notes may be difficult to justify.The condition involves gradients, momentum, and error accumulation, motivating the second scenario.
  • Scenario 2: Sliding Window Heavy Hitters: Under the sliding-heavy-hitters assumption, FetchSGD uses sliding-window error accumulation to capture signal spread across up to I gradients.The implementation can maintain I sketches, while a logarithmic number of sketches is possible; experiments use a single error-accumulation sketch.
  • Guarantees and scope: The paper’s guarantees apply to non-i.i.d. data and bound the objective gradient norm; Theorem 1 matches uncompressed SGD, whereas Theorem 2 has a worse rate.The proof uses a virtual-sequence argument and can extend to other function classes through careful averaging.

5. Evaluation

FetchSGD is evaluated against FedAvg and local top-k across highly non-i.i.d. and small-client-data settings, including CIFAR, FEMNIST, and PersonaChat. It generally remains competitive or superior, while central sketching supports momentum and avoids reliance on persistent client state.

  • CIFAR: On CIFAR10 and CIFAR100, FedAvg and local top-k frequently diverge or achieve low accuracy, while FetchSGD performs especially well at higher compression.Many FedAvg and local top-k runs were excluded because they failed to converge or achieved very low accuracy.
  • FetchSGD’s linear compression operator enables momentum and error accumulation at the central aggregator rather than on clients.The implementation uses momentum parameter 0.9 in all experiments.
  • FetchSGD significantly outperforms competing methods with very small local datasets and non-i.i.d. data.
  • FEMNIST: FetchSGD remains competitive with FedAvg and local top-k on FEMNIST, a setting with relatively large local datasets and a more nearly i.i.d. distribution.
  • PersonaChat: On PersonaChat, FetchSGD achieves 3.9× compression without accuracy loss relative to uncompressed SGD and lower perplexity than similarly compressed alternatives.Most compressed methods outperform the uncompressed baseline early but later saturate as compression error hinders training.
  • PersonaChat: Sketching outperforms local top-k for all but the highest compression levels because local top-k requires local error-feedback state unavailable to intermittently participating clients.

6. Discussion

FetchSGD complements prior federated-learning work by reducing communication per round rather than only reducing the number of rounds, while retaining the setting’s other constraints.

  • FetchSGD reduces communication required each round, complementing prior approaches that primarily reduce the total number of rounds needed for convergence.

A. Experimental Details

The experiments were run on commercially available NVIDIA Pascal, Volta, and Turing architecture GPUs.

  • All experiments used commercially available NVIDIA Pascal, Volta, and Turing architecture GPUs.

A.1. CIFAR

The CIFAR experiments use fixed training schedules, highly fragmented single-class client data, and Pareto-frontier reporting across compression methods. Convergence failures were observed for several highly compressed or momentum-based runs.

  • Non-FedAvg CIFAR experiments train for 24 epochs with 1% client participation per round and a triangular learning-rate schedule.
  • CIFAR10 and CIFAR100 are partitioned across 10,000 and 50,000 clients, giving each client 5 or 1 image from a single class.
  • FetchSGD and local top-k are tuned over k and sketch or sparsity dimensions, while FedAvg uses reduced iterations to obtain compression.
  • For CIFAR10, one FetchSGD run, three local top-k runs, and all FedAvg runs with global momentum diverged; CIFAR100 had one divergent local top-k run and divergent momentum FedAvg runs.

A.2. FEMNIST

The FEMNIST experiments evaluate FetchSGD across compression settings, using a 40M-parameter ResNet101 and hyperparameter searches over sketch and sparsification choices. The supplied passages also note performance variation with momentum factor masking.

  • Training setup: The experiment trains a 40M-parameter ResNet101 with layer normalization, an average batch size of approximately 600, augmentation, and a triangular learning-rate schedule.The schedule uses a peak learning rate of 0.01 and a pivot epoch of 0.2 for one-epoch training.
  • Hyperparameters: FetchSGD searches k in [50, 100, 200] ×10^3 and sketch columns in [1, 2, 5, 10] ×10^6, while local top-k and FedAvg receive separate hyperparameter sweeps.The local top-k sweep covers k from 10 ×10^3 to 20,000 ×10^3; FedAvg varies global epochs, local epochs, and batch size.
  • Results: For larger k, momentum factor masking degrades performance.The supplied passage also states that performance increases over the uncompressed baseline in the discussed setting.

A.3. PersonaChat

The PersonaChat evaluation studies compression methods on a naturally non-i.i.d. GPT2-small language-modeling task. It reports perplexity and examines compression trade-offs and true top-k behavior across different k values.

  • Task: PersonaChat partitions 17,568 clients by assigned personality and evaluates a pretrained 124M-parameter GPT2-small model.The task uses conversations between Mechanical Turk workers with faux personalities.
  • Evaluation: Perplexity is reported as the average per-word branching factor, with standard deviations over three random seeds for representative compression-method runs.The modified optimizer, learning rate, and dataset augmentation prevent direct comparison with Wolf et al. (2019).
  • Evaluation: Local top-k consistently performs worse on this task when using global momentum, so the reported local top-k results omit momentum.Local momentum is not possible because each client participates only once.
  • True top-k: True top-k outperforms the uncompressed baseline for intermediate k values, while performance decreases for large k because momentum factor masking inhibits momentum.Figure 10 examines validation perplexity on PersonaChat across a range of k values.
  • Compression analysis: Compression plots report upload, download, and overall compression, with CIFAR10 and CIFAR100 shown by dataset and Pareto-frontier runs shown for readability.All converged runs are shown separately in Figure 7.

B.1. Scenario 1:

Scenario 1 analyzes FetchSGD under a contraction assumption and establishes convergence for non-convex federated optimization with vanilla Count Sketch error accumulation. The analysis also covers client sampling, momentum, sketch recovery, and sketch-size requirements.

  • Assumption and theorem: Scenario 1 assumes a contraction property throughout training and proves convergence using vanilla Count Sketch error accumulation while recovering the k highest-magnitude elements.This matches the experimental recovery procedure described in the supplied passage.
  • Client sampling: Each iteration samples W clients uniformly from C clients, whose stochastic gradients are combined into an unbiased gradient estimator with bounded norm.The estimator is described as stochastic gradients with respect to the empirical mixture across clients.
  • Sketch update: FetchSGD’s compressed update uses Top-k after Count Sketch unsketching, while error accumulation remains represented by sketches rather than explicitly stored vectors.The proof introduces a virtual sequence and reduces the update to an SGD-like update with a scaled learning rate.
  • Sketch guarantees: Count Sketch must approximate vector norms within a multiplicative (1 ± τ) error and recover all τ-heavy coordinates with high probability.The required sketch memory follows from these norm-approximation and heavy-hitter-recovery goals.
  • Consequences: Setting momentum ρ = 0 recovers a no-momentum convergence guarantee, and the proof concludes convergence after substituting the established bounds.The supplied proof passages state this as Corollary 1 and finish the theorem argument after T iterations.

B.2. Scenario 2

Scenario 2 addresses gradient signal spread across multiple iterations, where vanilla error accumulation can let sketch noise grow and hinder heavy-hitter recovery. FetchSGD instead uses a sliding-window Count Sketch to discard old noise while retaining recent signal, supporting convergence under the stated heavy-hitter assumption.

  • Limitation of vanilla accumulation: Under Assumption 2, useful gradient signal may be spread across at most I iterations, while vanilla error accumulation noise can grow with t and hinder convergence when t ≫ I.The assumption is more interpretable than Scenario 1’s contraction assumption but is still a condition on the encountered gradient sequence.
  • Sliding-window method: A sliding-window Count Sketch recovers signal up to iteration t − I and discards remaining noise, preventing noise from growing as O(t) without losing useful signal.The construction maintains I individual Count Sketches in the straightforward implementation.
  • Recovery guarantee: The sliding-window structure recovers all ℓ2-heavy signal spread over up to I iterations with probability 1 − δ when using the individual sketches.The data structure is built so any sequence of up to I gradients appears in one of the maintained sketches.
  • Implementation variants: The implementation can use fewer than I individual sketches by maintaining only log(I) sketches, and it may recover τ-heavy rather than highest-magnitude k elements for simpler analysis.These are described as more efficient implementations and presentation simplifications, respectively.
  • Warm-up case: For I = 1, the sliding-window method reduces to vanilla Count Sketch compression without cross-iteration error accumulation.The proof analyzes this warm-up case before extending the argument to general I.
  • General case: For I > 1, sliding-window sketches capture signal spread across multiple iterations even when it is not ℓ2-heavy in any single iteration.The general-case proof largely follows the I = 1 proof because of the similarity between the two sketch structures.

B.3. Are these assumptions necessary?

The analysis establishes convergence without assumptions on the gradient sequence, but its worst-case variance dependence on dimension motivates additional assumptions to explain practical performance.

  • The unsketched estimate is unbiased, so Top-k, error feedback, and momentum can be analyzed as operating on stochastic gradients.
  • The estimator’s worst-case variance is Θ(d), yielding an optimal convergence rate under standard SGD analysis for these stochastic gradients.
  • The resulting guarantee holds without assumptions on the sequence of gradients encountered during optimization.
  • Count Sketch supports heavy-coordinate recovery through hashed bins, while AMS-style sketches provide norm estimates with controllable variance and failure probability.
  • Count Sketch linearity enables combining client sketches and maintaining error accumulation and momentum within the sketch.
Loading 2007.07682v2…