Federated Learning with Personalization Layers

TL;DR

The paper tackles personalization under federated learning by introducing FedPer, a base+personalization layer architecture where a globally shared base is trained via FedAvg and client-specific personalization layers are trained locally. This separation combats statistical heterogeneity across devices, enabling better performance on personalized tasks and reducing cross-client variance. Empirical results on non-identical CIFAR partitions and the Flickr-AES dataset show FedPer outperforming FedAvg, with insights into the roles of base versus personalization layers. The work demonstrates the practicality of deep federated personalization and provides a reproducible experimental setup across multiple architectures and datasets.

Abstract

The emerging paradigm of federated learning strives to enable collaborative training of machine learning models on the network edge without centrally aggregating raw data and hence, improving data privacy. This sharply deviates from traditional machine learning and necessitates the design of algorithms robust to various sources of heterogeneity. Specifically, statistical heterogeneity of data across user devices can severely degrade the performance of standard federated averaging for traditional machine learning applications like personalization with deep learning. This paper pro-posesFedPer, a base + personalization layer approach for federated training of deep feedforward neural networks, which can combat the ill-effects of statistical heterogeneity. We demonstrate effectiveness ofFedPerfor non-identical data partitions ofCIFARdatasetsand on a personalized image aesthetics dataset from Flickr.

Figures (18)

• Figure 1: Pictorial view of proposed federated personalization approach. All user devices share a set of base layers with same weights (colored blue) and have distinct personalization layers that can potentially adapt to individual data. The base layers are shared with the parameter server while the personalization layers are kept private by each device.
• Figure 2: Performance of FedAvg vs FedPer on non-identical 10 CIFAR CIFAR-10 partition ($k \in \IfNoValueTF{-NoValue-}{ \mleft\{4,8,10\mright\} } {\mathopen{-NoValue-\{}4,8,10\mathclose{-NoValue-\}}}$)
• Figure 3: Variation in v1 MobileNet MobileNet-v1 performance across clients for FedAvg vs FedPer on non-identical 10 CIFAR CIFAR-10 partition ($k=4$)
• Figure 4: Performance of FedPer on non-identical 100 CIFAR CIFAR-100 partition w.r.t. #personalization layers
• Figure 5: Effect on the performance of v1 MobileNet MobileNet-v1 on 100 CIFAR CIFAR-100 before and after replacing the base layers with linear layers