FedAlign: Federated Domain Generalization with Cross-Client Feature Alignment

TL;DR

FedAlign is a novel federated learning framework that enhances domain generalization by expanding feature diversity and aligning features across clients, all while preserving privacy and minimizing overhead.

Contribution

It introduces a cross-client feature extension and dual-stage alignment modules to improve domain invariance in federated learning.

Findings

Achieves superior generalization to unseen domains.

Maintains data privacy with minimal computational overhead.

Enhances feature diversity and domain invariance.

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.