Federated Submodel Optimization for Hot and Cold Data Features

TL;DR

This paper introduces FedSubAvg, a federated learning algorithm designed to efficiently handle sparse, non-i.i.d. data by averaging only relevant submodels, with proven convergence and superior empirical performance.

Contribution

The paper proposes FedSubAvg, a novel federated optimization method that effectively manages sparse data features and provides theoretical convergence guarantees.

Findings

FedSubAvg outperforms FedAvg on multiple datasets.

Theoretical convergence rate established under element-wise gradient norm.

Significant speedup in training with sparse data features.

Abstract

We study practical data characteristics underlying federated learning, where non-i.i.d. data from clients have sparse features, and a certain client's local data normally involves only a small part of the full model, called a submodel. Due to data sparsity, the classical federated averaging (FedAvg) algorithm or its variants will be severely slowed down, because when updating the global model, each client's zero update of the full model excluding its submodel is inaccurately aggregated. Therefore, we propose federated submodel averaging (FedSubAvg), ensuring that the expectation of the global update of each model parameter is equal to the average of the local updates of the clients who involve it. We theoretically proved the convergence rate of FedSubAvg by deriving an upper bound under a new metric called the element-wise gradient norm. In particular, this new metric can characterize the convergence of federated optimization over sparse data, while the conventional metric of squared gradient norm used in FedAvg and its variants cannot. We extensively evaluated FedSubAvg over both public and industrial datasets. The evaluation results demonstrate that FedSubAvg significantly outperforms FedAvg and its variants.