Calibre: Towards Fair and Accurate Personalized Federated Learning with Self-Supervised Learning

TL;DR

Calibre improves personalized federated learning by calibrating self-supervised learning representations, balancing generic and client-specific features to enhance accuracy and fairness across heterogeneous data sources.

Contribution

The paper introduces Calibre, a novel framework that calibrates SSL representations in federated learning using a client-specific prototype loss and prototype-guided aggregation, improving personalization and fairness.

Findings

Achieves state-of-the-art accuracy in non-i.i.d. settings.

Balances fairness and accuracy effectively across clients.

Outperforms existing methods in diverse heterogeneity scenarios.

Abstract

In the context of personalized federated learning, existing approaches train a global model to extract transferable representations, based on which any client could train personalized models with a limited number of data samples. Self-supervised learning is considered a promising direction as the global model it produces is generic and facilitates personalization for all clients fairly. However, when data is heterogeneous across clients, the global model trained using SSL is unable to learn high-quality personalized models. In this paper, we show that when the global model is trained with SSL without modifications, its produced representations have fuzzy class boundaries. As a result, personalized learning within each client produces models with low accuracy. In order to improve SSL towards better accuracy without sacrificing its advantage in fairness, we propose Calibre, a new personalized federated learning framework designed to calibrate SSL representations by maintaining a suitable balance between more generic and more client-specific representations. Calibre is designed based on theoretically-sound properties, and introduces (1) a client-specific prototype loss as an auxiliary training objective; and (2) an aggregation algorithm guided by such prototypes across clients. Our experimental results in an extensive array of non-i.i.d.~settings show that Calibre achieves state-of-the-art performance in terms of both mean accuracy and fairness across clients. Code repo: https://github.com/TL-System/plato/tree/main/examples/ssl/calibre.