FedClassAvg: Local Representation Learning for Personalized Federated Learning on Heterogeneous Neural Networks

TL;DR

FedClassAvg is a communication-efficient personalized federated learning method that enables clients with heterogeneous neural network architectures to collaboratively learn classifiers and improve local feature representations without sharing private data.

Contribution

It introduces a novel classifier averaging approach that works with heterogeneous models and enhances local feature learning without extra data or complex optimization.

Findings

Outperforms state-of-the-art algorithms on heterogeneous tasks

Requires minimal communication by exchanging only classifier layers

Does not need auxiliary data or intensive computation

Abstract

Personalized federated learning is aimed at allowing numerous clients to train personalized models while participating in collaborative training in a communication-efficient manner without exchanging private data. However, many personalized federated learning algorithms assume that clients have the same neural network architecture, and those for heterogeneous models remain understudied. In this study, we propose a novel personalized federated learning method called federated classifier averaging (FedClassAvg). Deep neural networks for supervised learning tasks consist of feature extractor and classifier layers. FedClassAvg aggregates classifier weights as an agreement on decision boundaries on feature spaces so that clients with not independently and identically distributed (non-iid) data can learn about scarce labels. In addition, local feature representation learning is applied to stabilize the decision boundaries and improve the local feature extraction capabilities for clients. While the existing methods require the collection of auxiliary data or model weights to generate a counterpart, FedClassAvg only requires clients to communicate with a couple of fully connected layers, which is highly communication-efficient. Moreover, FedClassAvg does not require extra optimization problems such as knowledge transfer, which requires intensive computation overhead. We evaluated FedClassAvg through extensive experiments and demonstrated it outperforms the current state-of-the-art algorithms on heterogeneous personalized federated learning tasks.