Federated Automatic Differentiation

TL;DR

This paper introduces Federated Automatic Differentiation (FAD), a novel framework enabling derivative computation across communication boundaries in federated learning, facilitating improved algorithm adaptation and performance without compromising privacy.

Contribution

The paper proposes FAD, a framework for computing derivatives in federated learning across communication boundaries, compatible with existing privacy-preserving primitives, and enabling dynamic algorithm learning.

Findings

FAD allows derivatives to be computed across client-server communication boundaries.

FAD can be implemented using various accumulation modes with different trade-offs.

Using FAD, federated algorithms like FedAvg can adaptively improve performance.

Abstract

Federated learning (FL) is a general framework for learning across an axis of group partitioned data (heterogeneous clients) while preserving data privacy, under the orchestration of a central server. FL methods often compute gradients of loss functions purely locally (ie. entirely at each client, or entirely at the server), typically using automatic differentiation (AD) techniques. We propose a federated automatic differentiation (FAD) framework that 1) enables computing derivatives of functions involving client and server computation as well as communication between them and 2) operates in a manner compatible with existing federated technology. In other words, FAD computes derivatives across communication boundaries. We show, in analogy with traditional AD, that FAD may be implemented using various accumulation modes, which introduce distinct computation-communication trade-offs and systems requirements. Further, we show that a broad class of federated computations is closed under these various modes of FAD, implying in particular that if the original computation can be implemented using privacy-preserving primitives, its derivative may be computed using only these same primitives. We then show how FAD can be used to create algorithms that dynamically learn components of the algorithm itself. In particular, we show that FedAvg-style algorithms can exhibit significantly improved performance by using FAD to adjust the server optimization step automatically, or by using FAD to learn weighting schemes for computing weighted averages across clients.