FetchSGD: Communication-Efficient Federated Learning with Sketching

TL;DR

FetchSGD is a federated learning algorithm that uses Count Sketch compression and sketch mergeability to reduce communication costs, improve convergence with sparse client participation, and maintain high model accuracy.

Contribution

It introduces a novel sketch-based compression method that enables efficient aggregation and convergence guarantees in federated learning.

Findings

Achieves high compression rates with maintained accuracy.

Demonstrates effective training of residual networks and transformers.

Provides theoretical convergence guarantees.

Abstract

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.