Overcoming label shift in targeted federated learning

Adam Breitholtz, Edvin Listo Zec, Fredrik D. Johansson

Keywords: federated learning, label shift, distribution shift

TL;DR: We propose a method of aggregation, based on central knowledge of label marginals, to deal with shifts between the aggregated label distribution of clients and a fixed target. Our method outperforms baselines convincingly on a multitude of tasks.

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 client and target 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 targeted aggregation as offered by FedPALS.

Submission Number: 12