[Paper Review] Federated Learning with Domain Generalization

Liling Zhang, Xinyu Lei+3 more|arXiv (Cornell University)|Nov 20, 2021

Privacy-Preserving Technologies in DataComputer Science29 references28 citations

TL;DR

The paper presents FedADG, a federated adversarial domain generalization framework that align source-domain distributions to a dynamically learned reference distribution in a class-wise, privacy-preserving manner to improve unseen-domain generalization.

ABSTRACT

Federated Learning (FL) enables a group of clients to jointly train a machine learning model with the help of a centralized server. Clients do not need to submit their local data to the server during training, and hence the local training data of clients is protected. In FL, distributed clients collect their local data independently, so the dataset of each client may naturally form a distinct source domain. In practice, the model trained over multiple source domains may have poor generalization performance on unseen target domains. To address this issue, we propose FedADG to equip federated learning with domain generalization capability. FedADG employs the federated adversarial learning approach to measure and align the distributions among different source domains via matching each distribution to a reference distribution. The reference distribution is adaptively generated (by accommodating all source domains) to minimize the domain shift distance during alignment. In FedADG, the alignment is fine-grained since each class is aligned independently. In this way, the learned feature representation is supposed to be universal, so it can generalize well on the unseen domains. Intensive experiments on various datasets demonstrate that FedADG has comparable performance with the state-of-the-art.

Motivation & Objective

• Motivate domain generalization in federated learning settings where client data cannot be centralized.
• Propose a federated adversarial domain generalization (FedADG) framework to align source-domain distributions to a dynamic reference.
• Enable class-wise alignment to improve discrimination and generalization to unseen target domains.
• Preserve data privacy by keeping data local and only exchanging model parameters.
• Demonstrate competitive performance against state-of-the-art domain generalization methods in FL setups.

Proposed method

• Introduce Federated Adversarial Domain Generalization (FedADG) consisting of a feature extractor, discriminator, distribution generator, and classifier on each client.
• Use an adversarial learning network (ALN) to align features from real data to a generated reference distribution.
• Adopt a dynamically generated reference distribution that centers around all source-domain distributions to reduce alignment distortion.
• Incorporate class-wise alignment by using label information (one-hot vector) during adversarial learning.
• Train with a combined loss L_FedADG = L_adv_d + L_adv_g + lambda_0 L_adv_f + lambda_1 L_err, where L_adv_d, L_adv_g, L_adv_f are adversarial losses and L_err is cross-entropy with label smoothing.
• Implement FedALN to enable federated training of the adversarial network without centralized access to client data.
• Provide a server–client training protocol with periodic aggregation of (F, C, G) and local updates to D.

Experimental results

Research questions

• RQ1 How can domain generalization be achieved in federated learning without sharing raw data?
• RQ2 Can a dynamically generated reference distribution improve alignment and reduce feature distortion across source domains?
• RQ3 Does class-wise alignment yield better generalization to unseen target domains compared to global alignment?
• RQ4 How does FedADG perform relative to centralized DG methods and existing FL baselines on standard DG benchmarks?
• RQ5 What is the impact of adaptively generated references versus fixed references on target-domain performance?

Key findings

• FedADG achieves competitive domain generalization performance compared to state-of-the-art methods on standard DG datasets.
• On VLCS with AlexNet backbone, FedADG (distributed) achieves 75.09 average accuracy, and with ResNet18 backbone 76.28 average accuracy.
• FedADG outperforms several centralized and distributed baselines (e.g., FedAvg) in average accuracy on VLCS.
• Adaptive reference distributions yield smaller alignment distortion and more discriminative, class-wise aligned features than fixed references.
• Class-wise alignment and adversarially learned domain-invariant features contribute to improved unseen-domain generalization.
• A representative example shows FedADG’s average accuracy across target domains improving from distributed FedAvg to FedADG (e.g., 54.79 vs 75.09 on AlexNet in VLCS).