Timezone: »

 
Poster
FedAvg with Fine Tuning: Local Updates Lead to Representation Learning
Liam Collins · Hamed Hassani · Aryan Mokhtari · Sanjay Shakkottai

Wed Nov 30 09:00 AM -- 11:00 AM (PST) @ Hall J #219
in Poster Session 3 »

The Federated Averaging (FedAvg) algorithm, which consists of alternating between a few local stochastic gradient updates at client nodes, followed by a model averaging update at the server, is perhaps the most commonly used method in Federated Learning. Notwithstanding its simplicity, several empirical studies have illustrated that the model output by FedAvg leads to a model that generalizes well to new unseen tasks after a few fine-tuning steps. This surprising performance of such a simple method, however, is not fully understood from a theoretical point of view. In this paper, we formally investigate this phenomenon in the multi-task linear regression setting. We show that the reason behind the generalizability of the FedAvg output is FedAvg’s power in learning the common data representation among the clients’ tasks, by leveraging the diversity among client data distributions via multiple local updates between communication rounds. We formally establish the iteration complexity required by the clients for proving such result in the setting where the underlying shared representation is a linear map. To the best of our knowledge, this is the first result showing that FedAvg learns an expressive representation in any setting. Moreover, we show that multiple local updates between communication rounds are necessary for representation learning, as distributed gradient methods that make only one local update between rounds provably cannot recover the ground-truth representation in the linear setting, and empirically yield neural network representations that generalize drastically worse to new clients than those learned by FedAvg trained on heterogeneous image classification datasets.

Keywords: Representation Learning · federated learning

Author Information

Liam Collins (The University of Texas at Austin)
Hamed Hassani (UPenn)
Aryan Mokhtari (UT Austin)
Sanjay Shakkottai (University of Texas at Austin)

More from the Same Authors