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

TL;DR

The paper tackles S-FDG, where clients hold unlabeled data and the server has limited labels, under practical domain shifts. It introduces the Unified Alignment Protocol (UAP), a two-stage alternating training framework that first aligns the server feature distribution to a known parametric form and then guides client features to align with that distribution, all without additional communication. By enforcing a Gaussian-like server feature model with a class-conditioned CDD loss and a covariance regularizer, UAP yields domain-invariant representations and improved unseen-domain generalization. Experiments on five domain-generalization benchmarks demonstrate state-of-the-art performance in the S-FDG setting across diverse architectures, highlighting a practical, privacy-preserving path to robust cross-domain learning in real-world federated systems.

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.

Figures (2)

• Figure 1: A comparative illustration of Federated Domain Generalization (FDG) FedGMAFedDGFedDG_GAfedsr, Semi-Supervised Federated Learning (SSFL) jeong2020federated9671693semiFL, and our proposed Semi-Supervised Federated Domain Generalization (S-FDG). FDG (left) assumes clients have labeled data, with domain shift occurring during testing. SSFL (right) assumes clients have unlabeled data, the server has limited labeled data, but training and testing occur within a single domain. In contrast, S-FDG (middle) models a more realistic scenario where clients have unlabeled data, the server has limited labeled data, and domain shift occurs during testing.
• Figure 2: Overview of our proposed UAP, where the server model is trained to learn and align feature with a parametric distribution (Sever Feature Alignment). Then, the server conveys both the model and its feature distribution parameters (no communication overhead) to the client by embedding them into the model parameters. Clients then leverage the server feature distribution knowledge to align their features (Client Feature Alignment) accordingly.