PLayer-FL: A Principled Approach to Personalized Layer-wise Cross-Silo Federated Learning

TL;DR

PLayer-FL introduces a novel, principled method to identify which layers in a neural network should be federated in personalized federated learning, improving performance and fairness across clients.

Contribution

The paper proposes a new federation sensitivity metric and a layer-wise federated learning approach that adaptively selects layers to federate, overcoming architecture-specific limitations of prior methods.

Findings

PLayer-FL outperforms existing FL algorithms on various tasks.

The federation sensitivity metric correlates well with generalization measures.

PLayer-FL achieves more uniform client performance.

Abstract

Non-identically distributed data is a major challenge in Federated Learning (FL). Personalized FL tackles this by balancing local model adaptation with global model consistency. One variant, partial FL, leverages the observation that early layers learn more transferable features by federating only early layers. However, current partial FL approaches use predetermined, architecture-specific rules to select layers, limiting their applicability. We introduce Principled Layer-wise-FL (PLayer-FL), which uses a novel federation sensitivity metric to identify layers that benefit from federation. This metric, inspired by model pruning, quantifies each layer's contribution to cross-client generalization after the first training epoch, identifying a transition point in the network where the benefits of federation diminish. We first demonstrate that our federation sensitivity metric shows strong correlation with established generalization measures across diverse architectures. Next, we show that PLayer-FL outperforms existing FL algorithms on a range of tasks, also achieving more uniform performance improvements across clients.