Unified Alignment Protocol: Making Sense of the Unlabeled Data in New Domains

TL;DR

This paper introduces the Unified Alignment Protocol (UAP), a novel framework for semi-supervised federated learning that enhances model generalization across unseen domains by aligning feature distributions without extra communication overhead.

Contribution

The paper proposes UAP, a two-stage training framework that aligns server and client features to improve domain generalization in SSFL, addressing practical domain shift challenges.

Findings

UAP achieves state-of-the-art results on domain generalization benchmarks.

UAP effectively aligns features without additional communication overhead.

Experimental results demonstrate improved generalization across multiple datasets.

Abstract

Semi-Supervised Federated Learning (SSFL) is gaining popularity over conventional Federated Learning in many real-world applications. Due to the practical limitation of limited labeled data on the client side, SSFL considers that participating clients train with unlabeled data, and only the central server has the necessary resources to access limited labeled data, making it an ideal fit for real-world applications (e.g., healthcare). However, traditional SSFL assumes that the data distributions in the training phase and testing phase are the same. In practice, however, domain shifts frequently occur, making it essential for SSFL to incorporate generalization capabilities and enhance their practicality. The core challenge is improving model generalization to new, unseen domains while the client participate in SSFL. However, the decentralized setup of SSFL and unsupervised client training necessitates innovation to achieve improved generalization across domains. To achieve this, we propose a novel framework called the Unified Alignment Protocol (UAP), which consists of an alternating two-stage training process. The first stage involves training the server model to learn and align the features with a parametric distribution, which is subsequently communicated to clients without additional communication overhead. The second stage proposes a novel training algorithm that utilizes the server feature distribution to align client features accordingly. Our extensive experiments on standard domain generalization benchmark datasets across multiple model architectures reveal that proposed UAP successfully achieves SOTA generalization performance in SSFL setting.