Overcoming label shift with target-aware federated learning

TL;DR

This paper introduces FedPALS, a federated learning method that adapts to label shifts across clients to improve model performance in the target domain, especially under label distribution differences.

Contribution

The paper presents FedPALS, a novel aggregation scheme that accounts for label shifts in federated learning, enhancing robustness and generalization across diverse client data.

Findings

FedPALS outperforms baseline methods in image classification tasks.

Conventional federated learning suffers under extreme label sparsity.

Target-aware aggregation improves robustness to label distribution differences.

Abstract

Federated learning enables multiple actors to collaboratively train models without sharing private data. Existing algorithms are successful and well-justified in this task when the intended target domain, where the trained model will be used, shares data distribution with the aggregate of clients, but this is often violated in practice. A common reason is label shift -- that the label distributions differ between clients and the target domain. We demonstrate empirically that this can significantly degrade performance. To address this problem, we propose FedPALS, a principled and practical model aggregation scheme that adapts to label shifts to improve performance in the target domain by leveraging knowledge of label distributions at the central server. Our approach ensures unbiased updates under federated stochastic gradient descent which yields robust generalization across clients with diverse, label-shifted data. Extensive experiments on image classification tasks demonstrate that FedPALS consistently outperforms baselines by aligning model aggregation with the target domain. Our findings reveal that conventional federated learning methods suffer severely in cases of extreme label sparsity on clients, highlighting the critical need for target-aware aggregation as offered by FedPALS.