Communication-Efficient Federated Hypergradient Computation via Aggregated Iterative Differentiation

TL;DR

This paper introduces AggITD, a communication-efficient method for federated hypergradient computation that reduces communication rounds while maintaining sample complexity, especially effective under data heterogeneity.

Contribution

The paper proposes AggITD, a novel federated hypergradient estimator that improves communication efficiency and is simpler to implement compared to existing methods.

Findings

AggITD achieves the same sample complexity as AID-based methods.

AggITD significantly reduces communication rounds in federated settings.

Experiments confirm the effectiveness and communication efficiency of AggITD.

Abstract

Federated bilevel optimization has attracted increasing attention due to emerging machine learning and communication applications. The biggest challenge lies in computing the gradient of the upper-level objective function (i.e., hypergradient) in the federated setting due to the nonlinear and distributed construction of a series of global Hessian matrices. In this paper, we propose a novel communication-efficient federated hypergradient estimator via aggregated iterative differentiation (AggITD). AggITD is simple to implement and significantly reduces the communication cost by conducting the federated hypergradient estimation and the lower-level optimization simultaneously. We show that the proposed AggITD-based algorithm achieves the same sample complexity as existing approximate implicit differentiation (AID)-based approaches with much fewer communication rounds in the presence of data heterogeneity. Our results also shed light on the great advantage of ITD over AID in the federated/distributed hypergradient estimation. This differs from the comparison in the non-distributed bilevel optimization, where ITD is less efficient than AID. Our extensive experiments demonstrate the great effectiveness and communication efficiency of the proposed method.