[Paper Review] FedDAR: Federated Domain-Aware Representation Learning

Aoxiao Zhong, Hao He+3 more|arXiv (Cornell University)|Sep 8, 2022

Machine Learning in Healthcare4 citations

TL;DR

FedDAR proposes a novel federated learning framework for domain-mixed non-i.i.d. data, where each client's data is a mixture of predefined domains. It decouples representation learning and prediction by sharing an encoder across domains while using domain-specific heads, enabling domain-aware personalization through alternating optimization and second-order aggregation, achieving linear convergence and superior performance on synthetic and real-world medical datasets.

ABSTRACT

Cross-silo Federated learning (FL) has become a promising tool in machine learning applications for healthcare. It allows hospitals/institutions to train models with sufficient data while the data is kept private. To make sure the FL model is robust when facing heterogeneous data among FL clients, most efforts focus on personalizing models for clients. However, the latent relationships between clients' data are ignored. In this work, we focus on a special non-iid FL problem, called Domain-mixed FL, where each client's data distribution is assumed to be a mixture of several predefined domains. Recognizing the diversity of domains and the similarity within domains, we propose a novel method, FedDAR, which learns a domain shared representation and domain-wise personalized prediction heads in a decoupled manner. For simplified linear regression settings, we have theoretically proved that FedDAR enjoys a linear convergence rate. For general settings, we have performed intensive empirical studies on both synthetic and real-world medical datasets which demonstrate its superiority over prior FL methods.

Motivation & Objective

• To address data heterogeneity in cross-silo federated learning by modeling clients' data as mixtures of predefined domains, rather than assuming i.i.d. or client-specific distributions.
• To overcome the limitations of existing FL methods that personalize per client but ignore underlying domain structures, especially in healthcare where domains like ethnicity or scanner protocols cause distribution shifts.
• To develop a method that learns a shared representation across domains while enabling domain-specific prediction heads, improving model robustness and fairness.
• To theoretically and empirically validate that domain-wise personalization outperforms client-wise personalization under domain-mixed FL settings.

Proposed method

• FedDAR decomposes the model into a shared encoder and domain-specific prediction heads, decoupling representation learning from prediction for improved generalization.
• It alternates between updating local heads with multiple local epochs and aggregating encoder weights via a weighted average, with sample re-weighting to ensure domain fairness.
• For head aggregation, FedDAR employs a novel second-order optimization strategy to improve convergence and optimality of the global heads.
• The method uses a two-stage training process: a FedAvg warmup followed by alternating local updates and global aggregation of encoder and heads.
• It applies domain-aware loss re-weighting to balance contributions from underrepresented domains during training.
• The framework supports both linear and non-linear models, with evaluation on CNNs and ResNet-34 in real medical applications.

Experimental results

Research questions

• RQ1 Can a federated learning method that explicitly models domain mixtures outperform standard and personalized FL baselines in non-i.i.d. medical data settings?
• RQ2 Does decoupling representation learning from domain-specific prediction improve convergence and performance compared to joint optimization?
• RQ3 How does domain-aware aggregation of prediction heads affect model fairness and generalization across diverse client data compositions?
• RQ4 What theoretical guarantees, such as convergence rates, can be established for the proposed method in linear settings?
• RQ5 To what extent does FedDAR reduce bias and improve performance on underrepresented domains in real-world medical datasets?

Key findings

• FedDAR achieves linear convergence in the linear regression setting, demonstrating strong theoretical convergence properties.
• On synthetic data, FedDAR outperforms FedAvg, FedProx, FedPer, and FedRep in terms of test accuracy across all domains.
• On the FairFace dataset, FedDAR achieves higher average accuracy across all racial groups compared to state-of-the-art methods, with improved fairness and reduced bias.
• In the real-world EXAM dataset, FedDAR achieves an average AUC of 0.89 on local validation sets, significantly outperforming FedProx, FedRep, and FedMinMax.
• The ablation study confirms that the second-order head aggregation and domain-aware re-weighting are critical for performance, especially on imbalanced domains.
• FedDAR maintains strong performance even when one client has only 100 samples per domain, demonstrating robustness to low-data domains.