Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FedAlign: Federated Domain Generalization with Cross-Client Feature Alignment

Sunny Gupta, Vinay Sutar, Varunav Singh, Amit Sethi

TL;DR

FedAlign addresses Federated Domain Generalization under privacy constraints by introducing a privacy-preserving cross-client feature extension and a dual-stage alignment strategy. It employs a MixStyle-based cross-client augmentation to broaden domain exposure and a combination of representation and prediction alignment losses to learn domain-invariant features. The framework achieves state-of-the-art generalization to unseen domains across multiple benchmarks with low overhead and robust scalability to many clients. This work offers a practical, privacy-conscious path for robust FDG in real-world federated systems.

Abstract

Federated Learning (FL) offers a decentralized paradigm for collaborative model training without direct data sharing, yet it poses unique challenges for Domain Generalization (DG), including strict privacy constraints, non-i.i.d. local data, and limited domain diversity. We introduce FedAlign, a lightweight, privacy-preserving framework designed to enhance DG in federated settings by simultaneously increasing feature diversity and promoting domain invariance. First, a cross-client feature extension module broadens local domain representations through domain-invariant feature perturbation and selective cross-client feature transfer, allowing each client to safely access a richer domain space. Second, a dual-stage alignment module refines global feature learning by aligning both feature embeddings and predictions across clients, thereby distilling robust, domain-invariant features. By integrating these modules, our method achieves superior generalization to unseen domains while maintaining data privacy and operating with minimal computational and communication overhead.

FedAlign: Federated Domain Generalization with Cross-Client Feature Alignment

TL;DR

FedAlign addresses Federated Domain Generalization under privacy constraints by introducing a privacy-preserving cross-client feature extension and a dual-stage alignment strategy. It employs a MixStyle-based cross-client augmentation to broaden domain exposure and a combination of representation and prediction alignment losses to learn domain-invariant features. The framework achieves state-of-the-art generalization to unseen domains across multiple benchmarks with low overhead and robust scalability to many clients. This work offers a practical, privacy-conscious path for robust FDG in real-world federated systems.

Abstract

Federated Learning (FL) offers a decentralized paradigm for collaborative model training without direct data sharing, yet it poses unique challenges for Domain Generalization (DG), including strict privacy constraints, non-i.i.d. local data, and limited domain diversity. We introduce FedAlign, a lightweight, privacy-preserving framework designed to enhance DG in federated settings by simultaneously increasing feature diversity and promoting domain invariance. First, a cross-client feature extension module broadens local domain representations through domain-invariant feature perturbation and selective cross-client feature transfer, allowing each client to safely access a richer domain space. Second, a dual-stage alignment module refines global feature learning by aligning both feature embeddings and predictions across clients, thereby distilling robust, domain-invariant features. By integrating these modules, our method achieves superior generalization to unseen domains while maintaining data privacy and operating with minimal computational and communication overhead.
Paper Structure (24 sections, 14 equations, 4 figures, 1 table, 1 algorithm)

This paper contains 24 sections, 14 equations, 4 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Representation Alignment
  4. Style Transfer
  5. Federated Domain Generalization
  6. Methodology
  7. Preliminary
  8. Framework Overview
  9. Supervised Contrastive Loss ($L_{SC}$):
  10. Representation Consistency Loss ($L_{RC}$):
  11. Jensen–Shannon Divergence ($L_{JS}$):
  12. Cross-Client Feature Augmentation with MixStyle
  13. Clustering
  14. Probabilistic Sampling Weights
  15. Adversarial Training
  16. ...and 9 more sections

Figures (4)

  • Figure 1: Illustration of the typical scenario in FL. Each client contains data from a unique domain, and the test domain (Photo) differs from all domains present on the clients.
  • Figure 2: Overview of FedAlign: Clients share local model parameters and sample statistics with the server, which aggregates and redistributes them. Local training incorporates feature augmentation, representation alignment, and prediction alignment to enhance domain-invariant feature learning.
  • Figure 3: t-SNE visualization of the representation distribution using FedSR. The representations show domain-specific clusters with noticeable overlaps, highlighting the limitations of FedSR in learning robust domain-invariant features.
  • Figure 4: Average test accuracy (%) versus the number of participating clients.