F ED DAR: F EDERATED D OMAIN -A WARE R EPRESENTATION L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
Federated learning (FL) (McMahan et al., 2017a) is a machine learning approach that allows many clients(e.g.
mobile devices or organizations) to collaboratively train a model without sharing the data.
It has great potential to resolve the dilemma in real-world machine learning applications, especially in the domain of healthcare.
A robust and generalizable model in medical application usually requires a large amount of diverse data to train.
However, collecting a large-scale centralized dataset could be expensive or even impractical due to the constraints from regulatory, ethical and legal challenges, data privacy and protection (Rieke et al., 2020).
While promising, applying FL to real-world problems has many technical challenges.
One eminent challenge is data heterogeneity.
Data across the clients are assumed to be independently and identically distributed (iid) by many FL algorithms.
But this assumption rarely holds in the real world.
It has been shown that non-iid data distributions will cause the failure of standard FL strategies such as
FedAvg (Jiang et al., 2019; Sattler et al., 2020; Kairouz et al., 2019; Li et al., 2020).
As an ideal model that can perform well on all clients may not exist, it requires FL algorithms to personalize the model for different data distributions.
Prior theoretical work (Marfoq et al., 2021) shows that it is impossible to improve performances on all clients without making assumptions about the client’s data distributions.
Past works on personalized
FL methods (Marfoq et al., 2021; Sattler et al., 2020; Ghosh et al., 2020; Mansour et al., 2020; Deng et al., 2020) make their own assumptions and tailor their methods to those assumptions.
In this paper, we propose a new and more realistic assumption where each client’s data distribution is a mixture of several predefined domains.
We call our problem setting Domain-mixed FL .
It is inspired by the fact that the diversity of the medical data can be attributed to some known concept of domains, such as different demographic/ethnic groups of patients (Szczepura, 2005; Ranganathan & Bhopal, 2006;
NHS, 2004), different manufacturers or protocols/workflows of image scanners (Mårtensson et al., 2020; Ciompi et al., 2017), and so on.
It is necessary to address the ubiquitous issue of domain shifts among ethic groups (Szczepura, 2005; Ranganathan & Bhopal, 2006; NHS, 2004) or vendors (Yan et al., 2019; Garrucho et al., 2022; Guan & Liu, 2021) in healthcare data.
Despite of the domain shifts between data domains, same domain at different clients are usually considered to have the same distribution.
The data heterogeneity between FL clients actually comes from the distinct mixtures of diverse domains at clients.
These factors motivate us to personalize model for each domain instead of client.
Although our method is inspired by healthcare applications where the domain shifts issue is wellknown and domain labels are very basic and accessible, we believe that it can be generally applied to other domains like finance or recommendation systems where users/humans with different demography are involved (Ding et al., 2021; Asuncion & Newman, 2007).
However it would require a deep understanding of the data and background knowledge to verify the data distribution assumption as well as the accessibility of the domain label.
FedEM(Marfoq et al., 2021) and FedMinMax(Papadaki et al., 2021) makes similar assumption on data distribution as ours.
However, FedEM assumes the domains are unknown and tries to learn a linear combination of several shared component models with personalized mixture weights through an EM-like algorithm.
FedMinMax doesn’t acknowledge the domain shift between domains and still aims to learn one shared model across domains by adapting minmax optimization to FL setting .
Our Contributions.
We formulate the proposed problem setting, Domain-mixed FL .
Through our analysis, we find prior FL methods, both generic FL methods like FedAvg (McMahan et al., 2017a), and personalized FL methods like FedRep (Collins et al., 2021), are sub-optimal under our setting.
To address this issue, we propose a new algorithm, Federated Domain-Aware Representation
Learning (FedDAR) .
FedDAR learns a shared model for all the clients but embedded with domainwise personalized modules.
The model contains two parts: an shared encoder across all domains and a multi-headed predictor whose heads are associated with domains.
For an input from one specific domain, the model extracts representation via the shared encoder and then use the corresponding head to make the prediction.
FedDAR decouples the learning of the encoder and heads by alternating between the updates of the encoder and the heads.
It allows the clients to run many local updates on the heads without overfitting on domains with limited data samples.
This also leads to faster convergence and better performed model.
FedDAR also adapts different aggregation strategies for the two parts.
We use a weighted average operation to aggregate the local updates for the encoder.
With additional sample re-weighting, the overall training objective is equally weighted for each domain to encourage the fairness among domains.
While for the heads, we propose a novel second-order aggregation algorithm to improve the optimality of aggregated heads.
We theoretically show our method enjoys nice properties like linear convergence and small sample complexity in a linear case.
Through extensive experiments on both synthetic and real-world datasets, we demonstrate that FedDAR significantly improves performance over the state-of-the-art personalized FL methods.
To the best of our knowledge, our paper is among the first efforts in domain-wise personalized federated learning that achieve such superior performance.
R ELATED WORK
Besides the literature we have discussed above, other works on personalization and fairness in federated learning are also closely related to our work.
Personalized Federated Learning.
Personalized federated learning has been studied from a variety of perspectives: i) local fine-tuning (Wang et al., 2019; Yu et al., 2020) ii) meta-learning (Chen et al., 2018;
Fallah et al., 2020; Jiang et al., 2019; Khodak et al., 2019) iii) local/global model interpolation (Deng et al., 2020; Corinzia et al., 2019; Mansour et al., 2020).
iv) clustered FL that partition clients into clusters and learn optimal model for each cluster (Sattler et al., 2020; Mansour et al., 2020; Ghosh et al., 2020)(Zhu et al., 2021).
v) Multi-Task Learning(MTL) (Vanhaesebrouck et al., 2017; Smith et al., 2017; Zantedeschi et al., 2020)
(Hanzely & Richtárik, 2020; Hanzely et al., 2020; T Dinh et al., 2020; Huang et al., 2021; Li et al., 2021a) vi) local representations or heads for clients (Arivazhagan et al., 2019; Liang et al., 2020; Collins et al., 2021)(Luo et al., 2022).
vii) personalized model through hypernetwork or super model (Shamsian et al., 2021; Chen & Chao, 2021;
Xu et al., 2022).
The personalization module in our approach is similar to vi) and (Zhu et al., 2021) with a multi-branch network.
However, the targets we are personalizing the model for are the domains instead of clients.
Fairness in Federated Learning.
There are two commonly used definitions of fairness in existing
FL works.
One is client fairness, usually formulated as client parity (CP) , which requires clients to have similar performance.
A few works (Li et al., 2021a; 2019; Mohri et al., 2019; Yue et al., 2021;
Zhang et al., 2020) have studied on this.
Another is group fairness.
In the centralized setting, the fundamental tradeoff between group fairness and accuracy has been studied (Menon & Williamson, 2018;
Wick et al., 2019; Zhao & Gordon, 2019), and various fair training algorithms have been proposed(Roh et al., 2020; Jiang & Nachum, 2020; Zafar et al., 2017; Zemel et al., 2013; Hardt et al., 2016).
Since the notions of group fairness is the same in FL setting, most of existing FL works adapt methods from centralized setting (Zeng et al., 2021; Du et al., 2021;
Gálvez et al., 2021;
Chu et al., 2021; Cui et al., 2021).
In this work, our method is not designed specifically for certain group fairness notions like demographic parity.
Instead, we aim to achieve the best possible performance for each domain through personalization, admitting the difference between data domains.
Moreover, our concept of data domains is not limited as demographic groups.
It can also be applied to any other mixture of domain data, as long as our assumptions hold.
P ROBLEM : D OMAIN - MIXED F EDERATED L EARNING
Notations.
Federated learning involves multiple clients.
We denote number of clients as n .
We use i ∈ [ n ] ≜ { 1 , 2 , ..., n } to index each client.
Client i has a local data distribution D i which induces a local learning objective, i.e., the expected risk R i ( f ) =
E ( x i ,y i )
∼D i [ ℓ ( f ( x i ) , y i )] , where f : X → Y is the model mapping the input x ∈ X to the predicted label f ( x ) ∈ Y and ℓ : Y
× Y → R is a generic loss function.
In real practice, client i ∈ [ n ] has a finite number, say L i , of data samples, i.e.,
S i
= { ( x ji , y ji ) } L i j =1 .
L = (cid:80) ni
=1 L i denotes the total number of data samples.
Problem Formulation of Domain-mixed Federated Learning.
We introduce a new formulation of
FL problem by assuming each clients’ local data distribution is a weighted mixture of M domain specific distributions.
Specifically, we use { ˜ D m } Mm =1 to denote data distributions from M predefined domains.
For client i , its local data distribution is D i =
(cid:80) m
π i,m ˜ D m where the mixing coefficients
π i,m stand for the probabilities of client i ’s data sample coming from domain m .
Take medical application as an example, different hospitals are clients and different ethnic groups are domains.
Each ethnic group have different health data while each hospital’s data is a mix of ethnic group data.
Further, the domains of the data samples are assumed to be known.
We use a triplet of variables ( x , y, z ) to represent the input features, label and domain.
The goal of our problem is to learn a model f ( x , z ) that can perform well in every domain, as shown by the following learning objective,
[Equation]
where R m ( f ( · , m ))
=
E ( x ,m ) ∼ ˜ D m [ ℓ ( f ( x , m ) , y )] .
Our problem focuses on the setting that each domain have a different conditioned label distribution, i.e., P m ( y | x ) is different in each domain m .
3. C OMPARISON WITH P RIOR D OMAIN - UNAWARE FL P ROBLEM F ORMULATIONS
Our FL problem introduces the concept of the domain and focuses on the model’s performance in each domain.
Many prior FL formulations does not recognize the existence of the domains.
For example, the original federated learning algorithms like FedAvg (McMahan et al., 2017a), FedProx (Li et al., 2020) learn a globally shared model that via minimizing the averaged risk, i.e., min f 1 n (cid:80) i R i ( f ) .
Some variants consider the fairness across the clients.
To do so they optimize the worst client’s performance, instead of the averaged performance, i.e., min f max i R i ( f ) .
Further, personalized FL algorithms, such as FedRep (Collins et al., 2021), customize the model’s prediction for each client whose objective is min f i : i ∈ [ n ] 1 n (cid:80) ni
=1 R i ( f i ) .
All the FL algorithms mentioned above will lead sub-optimal solutions to our problem since they do not make domain specific predictions.
We illustrate this point by the following toy example of linear regression: We assume the data in m ’th domain is generated via the following procedure: x ∈ R d is i.i.d sampled from a distribution p ( x ) with mean zero and covariance I d .
The label y ∈ R obeys y
=
x ⊤ B ∗ w ∗ m where B ∗ ∈ R d × k is ground truth linear embedding shared by all domains, and
w ∗ m ∈
R k is the linear head specific to domain m .
Under this setting, ˜ D m stands for data ( x , y ) where x ∼ p ( x ) and y
=
x ⊤ B ∗ w ∗ m .
For each client, the local data D i is a mix of data from different domains with mixed coefficients, i.e., D i =
(cid:80) m
π i,m ˜ D m .
[Equation]
FedRep: learns shared representation B and separated heads w i for each clients i rather than for each domain m ,
[Equation]
FedDAR:
In contrast, in the linear case, our proposed method, FedDAR, which will be introduced next, learns a shared representation B and separate heads w m for each domain m ,
[Equation]
From the above formulations, we can see that FedAvd and FedRep are not able to achieve the zero error in our domain-mixed FL problem.
P ROPOSED M ETHOD : F ED DAR
To solve the Domain-mixed FL problem, we propose a new method called, Federated Domain-Aware
Representation Learning (FedDAR).
In the following, we first introduce the model, learning objective and the details of the federated optimization algorithm.
4. A LGORITHM O VERVIEW
Our model is made of a shared encoder ϕ ( · ; θ ) and M domain specific heads
h m ( · ; w m ) whose are parameterized by neural networks with the weights θ and w m , ∀ m ∈
[ M ] .
According to our problem formation in Equation 1, our algorithm aims to solve the following optimization,
[Equation]
We decouple the training between encoder and heads.
Specifically, we alternates the learning between the encoder and the heads.
The learning is done federatedly and has two conventional steps: (1) local updates; (2) aggregation at the server.
Algorithm 1 shows the pseudocode code.
Empirical Objectives with Re-weighting.
Empirically, the objectives are estimated via the finite data samples at each client.
We use S i,m to denote the set of samples from domain m in client i , with L i,m :=
|S i,m | denoting the sample size.
Further, L i :=
(cid:80)
Mm =1 L i,m is the number of samples in client i while L m :=
(cid:80) ni
=1 L i,m is the total number of samples belonging to domain m across all the clients.
We denote the empirical risk at client i specific to domain m as ˆ R i,m ( h m ◦ ϕ ) :=
L i,m (cid:80) ( x ,y ) ∈S i,m ℓ ( h m ◦ ϕ ( x ) , y ) .
The empirical risk at client i is designed as ˆ R i ( ϕ, h 1 , ..., h M ) =
(cid:80) m
L i,m L i u m ˆ R
i,m ( h m ◦ ϕ ) , where u m =
L L m M re-weights the risk for each domain.
Combining commonly used weighted average FL objective
ˆ R ( ϕ, h 1 , ..., h M ) = (cid:80)
ni
=
L i L
ˆ R i ( ϕ, h 1 , ..., h M ) , the overall empirical risk is derived as the following,
[Equation]
where ˆ R m ( h m ◦ ϕ ) := (cid:80) ni
=1 L i,m L m ˆ R i,m ( h m ◦ ϕ ) .
This is consistent with Equation 5.
4. L OCAL U PDATES AT C LIENTS
In each communication round, clients use gradient descent methods to optimize representation ϕ ( · ; θ ) and local heads h m ( · ; w m ) for m ∈
[ M ] alternately.
We use t to denote the current round.
For a module f , f t − 1 denotes its optimized version after t − 1 rounds.
Each round has multiple gradient descent iterations.
We use f t,s to denote the module in round t after s iterations.
Since the updates are made locally, clients maintain their own copies of both modules, we use subscripts i to index local copy at client i , e.g., f t,si .
We use GRD to denote a generic gradient-base optimization step which takes three inputs: objective function , variables , learning rate and maps them into a new module with updated variables.
For example, the vanilla gradient descent has the form
GRD ( L ( f w ) , f w , η ) =
f w − η ∇ w L ( f w ) .
For the heads , client i performs τ h local gradient-based updates to obtain optimal head given the current shared encoder ϕ t − 1 .
For s ∈ [ τ h ] , client i updates via h t,si,m ← GRD ( ˆ R i,m ( h t,s − 1 i,m ◦ ϕ t − 1 ) , h t,s − 1 i,m , η ) .
For the shared encoder , the clients executes τ ϕ local updates.
Specifically, for s ∈ [ τ ϕ ] , client i updates the local copy of the encoder via ϕ t,s i ← GRD ( ˆ R i ( ϕ t,s − 1 i , { h t m } M m =1 ) , ϕ t,s − 1 i , η ) .
The re-weighting mentioned in last section is implemented by re-weighting each sample with u m when calculating the loss function.
4. A GGREGATION AT S ERVER
We introduce two strategies: (1) weighted average (WA); (2) second-order aggregation (SA).
Weighted average means the aggregated model parameters are the average of the local model
’s parameters weighted by the number of data samples.
Specifically, for the shared encoder, we have θ t = (cid:80) ni
=1 L i L θ t − 1 .
Similarly for each head, we have w tm = (cid:80) ni
=1 L i,m L m w t − 1 m,i .
Second-order aggregation is a more complex strategy.
Ideally, we want the head aggregation generates the globally optimal model given a set of locally optimal model, as shown in the following,
[Equation]
where J i is i ’th client’s virtual objective, α i :
= L i /L is the importance of the client, L i is the number of data samples.
We call J i the virtual objective to distinguish it from the real learning objective R i .
The virtual objective is defined as an objective that the local updates give the optimal solution w.r.t it.
It is introduced since the local updates during two aggregated are not guaranteed to optimize the head to optimal w.r.t the real objective.
For example, if each local update is single step gradient descent with a learning rate η , i.e., w t +1 i =
w t − η ∇ w R
i ( w t ) .
Then the virtual objective becomes
J i ( w ) =
R i ( w t )
+( w − w t ) ⊤ ∇ w R
i ( w t )+ 12 η ∥ w − w t ∥ 22 which satisfies w t +1 i ∈ arg min w J i ( w ) .
Such a virtual objective leads the solution of problem 7 to w ∗ = 1 n (cid:80) ni
=1 w ∗ i which is the simple averaging strategy.
However, in real practice, the local updates are usually more complicated which makes the virtual objective closer to the true objective.
We consider the case that the virtual objective is the secondorder Taylor expansion of the true objective, i.e., J ( w ) =
R ( w t ) + ( w − w t ) ⊤ ∇ w R ( w t )
+ 12 ( w − w t )
⊤ H R ( w t )( w − w t ) where H R is the Hessian matrix.
Then each round of local update equivalent to a Newton-like step, w t +1 i =
w t − H R i ( w t ) − 1 ∇ w R
i ( w t ) .
While w t +1 =
w t − H R ( w t ) − 1 ∇ w R ( w t ) is the desired globally optima.
Leveraging the fact that, ∇ w R ( w ) = (cid:80) i ∈
[ n ] α i ∇ w R
i ( w ) and H R ( w ) =
(cid:80) i ∈ [ n ] α
i H
R
i ( w ) , we can get w t +1 from w t
+1 i via the following equation, which we call second-order aggregation,
[Equation]
Specifically, to implement second-order aggregation, in each round, the local clients first optimize the model locally for several epochs.
Then we compute the Hessian matrices for each local model and send them to the server for aggregation. Note that sending the Hessian takes a communication cost being quadratic to the size of the weight.
In real practice, the predictive head is usually small, e.g., a linear layer with hundreds of neurons.
Thus it is acceptable to aggregate the Hessian matrix of the head’s parameters.
In the following, we provide two instances of our second-order aggregation with a linear head.
Linear Regression where R i ( w ) =
1 L i (cid:80) L i j =1 ( w ⊤ x ji − y j ) 2 is quadratic itself.
Thus the second order Taylor expansion of the objective itself, i.e., J i ( w ) =
R i ( w ) .
In this case, H R i ( w ) =
X ⊤ i
X i where X
i =
[ x 1 i , · · · , x L i i ]
⊤ is the data matrix of client i . 2.
Binary Classification where R i ( w )
= − 1 L
i (cid:80) L i j =1 y ji log σ ( w ⊤ x ji )+(1 − y ji ) log(1 − σ ( w ⊤ x ji )) .
σ is the sigmoid function
Let µ ji ≜ σ ( w ⊤ x ji ) denote model’s output.
The gradient and the Hessian
are, ∇ w R i ( w ) = 1 L i (cid:80)
j ( µ ji − y ji ) x ji = 1 L i 1 ⊤ diag( µ i − y i ) X ⊤ i and H R
i w ) = 1 L
i X ⊤ i
SX
i where S ≜ diag( µ 1 i (1 − µ 1 i ) , · · · , µ L i i (1 − µ L i i )) .
Similar formulas can be derived for the multiclass classification.
Please refer to the text book (Murphy, 2022) for the exact equations.
Remark.
In practice, when the dimension of w is larger than the number of samples of a certain domain, the Hessian may have small singular values which cause numerical instability.
To mitigate this issue, one can either directly set the representation dimension k to some smaller number or add a (fully-connected) projection layer on top of a pretrained encoder to compress the representations to a lower dimensional space.
4. T HEORETICAL R ESULT OF F ED DAR
For a simplified linear regression setting as discussed in domain-mixed FL (4) (cf.
details in
Appendix A), we give below the sample complexity required for an adapted version of our algorithm (Algorithm 2 in the appendix) to enjoy linear convergence.
Due to the space limit, we only provide an informal statement to highlight the result.
The formal statement and the proof are deferred in the appendix.
Theorem 4.1 (Sample complexity of FedDAR convergence in linear case (informal)) .
Consider the linear setting for domain-mixed FL in (4).
At each iteration, suppose that the number of samples used by each of n clients to update the encoder, is ˜Ω( dk 2 n ) , and that the aggregate number of samples used in the update for the domain-specific heads, is ˜Ω( k ) .
Then, for a suitably chosen step size, the distance between the encoder B t Algorithm 2 outputs and the true encoder B ∗ converges at a linear rate.
Remark.
As our algorithm converges linearly to the true encoder, the per-iteration sample complexity of our algorithm gives a good estimate of the overall sample complexity.
Since we expect the output of the encoder to be significantly lower-dimensional than the input (i.e. k ≪ d ), our result indicates that Algorithm 2’s sample complexity is dominated by ˜Ω( dn ) , implying that the complexity reduces significantly as the number of clients n increases.
Moreover, a key implication of our result is the capacity for our algorithm to accommodate data imbalance across domains.
We note that our approach requires Ω( dk 2 ) samples per iteration for the update of the shared representation
B ∈ R d × k , whilst needing only Ω( k ) samples per iteration for the update of each domain head.
In particular, domains with more data can contribute disproportionately to the ˜Ω( dk 2 ) samples required to learn the common representation, whilst domains with fewer data need only provide ˜Ω( k ) samples to update its domain head during the course of the algorithm.
Whenever k ≪ d , which we believe is a reasonable assumption for many practical applications (e.g. medical imaging), the requirement of ˜Ω( k ) samples per domain is relatively mild.
Conversely, forgoing the shared representation structure would require each domain to learn a separate d -dimensional classifier, requiring ˜Ω( d ) samples per domain, which can pose a challenge in problems with domain data imbalance.
E XPERIMENTS
We validate our method’s effectiveness on both synthetic and real datasets.
We first experiment on the exact synthetic dataset described in our theoretical analysis to verify our theory.
We then conduct experiments on a real dataset, FairFace (Kärkkäinen & Joo, 2019), with controlled domain distributions to investigate the robustness of our algorithm under different levels of heterogeneity.
Finally, we compare our method with various baselines on a real federated learning benchmark,
EXAM (Dayan et al., 2021) with real-world domain distributions.
We also conduct extensive ablation studies on it to discern the contribution of each component of our method.
Full details of experimental settings can be found in Appendix B.
5. S YNTHETIC D ATA
We first run experiments on the linear regression problem analyzed in Appendix A.
We generate (domain, data, label) samples as the following, z i ∼ M ( π i ) , x i ∼ N (0 , I d ) , y i ∼ N ( w ∗
z
i ⊤ B ∗⊤ x i , σ ) where σ = 10 − 3 controls label observation errors, M ( π i ) is a multinomial domain distribution with parameter
π i
=
[ π i, 1 , ..., π i,M ]
∈
∆ M .
The hyper-parameters of domain distributions π i are drawn from a Dirichlet distribution, i.e., π i ∼ Dir ( α p ) , where p ∈ ∆ M is a prior domain distribution over M domains, and α > 0 is a concentration parameter controlling the heterogeneity of domain distributions among clients.
The largest domain distribution heterogeneity is achieved as α →where each client contains data only from a single randomly selected domain.
On the other hand, when α → ∞ , all clients have identical domain distributions that are equal to the prior p .
We generate ground-truth representation B ∗ ∈ R
d × k and domain specific heads w ∗ m , ∀ m ∈
[ M ] by sampling and normalizing Gaussian matrices.
Figure 5.2 shows result of our experiments where we set n = 100 clients, M = 5 domains, feature dimension k = 2 .
We vary the number of training samples per client from 5 to 20 .
The result shows that FedDAR-SA, achieves four orders of magnitude smaller errors than all the baselines: (1) Local-Only where each client train a model using its own data; (2) FedAvg which learns a single shared model; (3) FedRep which learns shared representation and client-specific heads. (4)Separate FedAvg which trains separate models for each domain using FedAvg.
The results demonstrate that our method overcomes the heterogeneity of domain distributions across clients.
FedDAR-WA fails to converge under this setting, confirming the effectiveness of the proposed second-order aggregation.
5. R EAL D ATA WITH C ONTROLLED D ISTRIBUTION
Dataset and Model.
We use FairFace (Kärkkäinen & Joo, 2019), a public face image dataset containing 7 race groups which are considered as the domains.
Each image is labeled with one ofage groups and gender.
We use the age label as the target to build a multi-class age classifier.
We created an FL setting by dividing training data into n clients without duplication.
Each client has a domain distribution π i ∼ Dir ( α p ) sampled from a Dirichlet distribution.
The total number of samples at each client L i = 500 is set to be the same in all experiments.
We control the heterogeneity of domain distributions by altering α .
The label distributions are uniform for all the clients. [Figure] [Figure] 64 k
A v e r a g e d o m a i n t e s t a cc u
[Figure] [Figure] [Table] Implementation and Evaluation.
We use Imagenet(Deng et al., 2009) pre-trained ResNet-34 (He et al., 2016) for all experiments on this dataset.
All the methods are trained for T = 100 communication rounds.
We use Adam optimizer with a learning rate of 1 × 10 − 4 for the first 60 rounds and 1 × 10 − 5 for the last 40 rounds.
Metrics and Results.
Our evaluation metrics are the classification accuracy on the whole validation set of FairFace for each race group.
We don’t have extra local validation set for each client since we assume the data distribution within each domain is consistent across the clients.
In Table 5.2, we report the accuracy averaged over the final 10 rounds of communication following the common practice (Collins et al., 2021).
The result shows our FedDAR achieved the best performance compared with the baselines.
Note that FedAvg + Multi-head also uses Equation 5 as objective for fair comparison.
Effect of k .
The limitation of using FedDAR-SA instead of FedDAR-WA is the need of tuning the dimension of representation k .
Figure 5.2 shows results of the average domain test accuracy with different k .
We can see that FedDAR-SA can achieve better accuracy with a properly chosen k .
We use k = 8 for all results with FedDAR-SA in Table 5.2.
Robustness to Varying Levels of Heterogeneity.
From the result with various α , we can observe that the performance of FedDAR-SA is very stable no matter how heterogeneous the domain mixtures are.
However, the baselines’ accuracy decrease when α becomes smaller.
5. R EAL D ATA WITH R EAL -W ORLD D ATA D ISTRIBUTION
Dataset and Model.
We use the EXAM dataset (Dayan et al., 2021), a large-scale, real-world healthcare FL study.
We use part of the dataset including 6 clients with a total of 7,681 cases.
We use race groups as domains.
The dataset is collected from suspected COVID-19 patients at the visit of the emergency department (ED), including both Chest X-rays (CXR) and electronic medical records (EMR).
We adopt the same data preprocessing procedure and the model as (Dayan et al., 2021).
Our task is to predict whether the patient received oxygen therapy higher than high-flow oxygen in 72 hours which indicates severe symptoms.
Baselines.
(1) methods that learn one global model, FedAvg(McMahan et al., 2017a), FedProx(Li et al., 2020), FedMinMax(Papadaki et al., 2021) along with their local fine-tuned variants; (2) train
M separate models with FedAvg; (3) train one global model with FedAvg first, then fine-tune on M domains separately with FedAvg; (4) client-wise personalized FL approaches, FedRep(Collins et al., 2021), FedPer(Arivazhagan et al., 2019), LG-Fedavg(Liang et
al., 2020), FedBN(Li et al., 2021b).
Implementation and Evaluation.
We apply 5-fold cross-validation.
All the models are trained for
T = 20 communication rounds with Adam optimizer and a learning rate of 1 ×
10 − 4 .
The models [Table] are evaluated by aggregating predictions on the local validation sets and then calculating the area under curve (AUC) for each domain.
The average AUCs on the local validation set of clients are also reported.
Average Performance Across Domains and Clients.
Table 3 shows the average of AUCs across domains and clients.
We can see that our methods, both FedDAR-WA and FedDAR-SA, achieve significantly better performance than all the baselines under both domain-wise and client-wise metrics.
The gap between our domain-wise personalized approach and other client-wise personalized baselines shows the validity of learning domain-wise personalized models facing diversity across domains.
The reason that fine-tuning methods induce worse results is mainly because of the imbalanced label distribution.
Each local training dataset doesn’t have enough positive cases to do proper fine-tuning.
Fairness Across Domains.
The AUCs of each specific domain in Table 3, show that our proposed
FedDAR method uniformly increases the AUC for each domain.
The column of the minimum AUC among domains also verifies that our method indeed improves the fairness across the domains.
[Table]
Ablation Studies.
i) re-weighting (RW [Table] in Table 2 shows adding sample re-weighting can significantly improve the fairness across the domains.
The minimum AUC among domains is improved by a large margin ( > 0 . 05 ); ii) multi-head (MH), domain as input feautre (DI) and alternating update (Alter): Comparing three blocks in Table 2, we can see that adding multi-head itself cannot bring any improvement.
We conjecture that alternating update prevents the overfitting of the heads with limited samples. This can also be reflected by the result in Table 5.2, where FedAvg+MH tends to perform badly on certain underrepresented domains especially when domain distributions are highly heterogeneous ( α is small).
Meanwhile, using domain labels directly as feature input is not as good as multi-head, and not compatible with alternating update; iii) projection (Proj) and aggregation method (AGG): Results in Table 2 shows that using second-order aggregation with the projection of the features gives the best result.
C ONCLUSIONS
In this paper, we propose a novel domain-aware personalized federated learning framework based on the assumption of the mixture of domain data distribution.
Our FedDAR approach is able to learn a global representation as well as domain-specific heads with balanced performance for each domain despite the heterogeneity of domain distributions across the clients.
We provide both theoretical and empirical justification for its effectiveness.
Our method is tested with a face recognition task and a real-world medical imaging FL dataset and can be easily extended to other complicated tasks like object detection and semantic segmentation due to its simplicity and flexibility.
The limitations of our method include: i) the domain information for all samples is required to be known; ii) the heterogeneity of label distributions is not considered; iii) the extra communication cost of sending Hessian matrices can be expensive, especially when the output dimension is big.
We plan to address these issues in future work.
Other future research directions include further boosting the fairness across domains and studying the setting where domains are structured, hierarchical, continuously indexed (Wang et al., 2020;
Nasery et al., 2021) or multi-dimensional (characterized by multiple factors)
(Wang et al., 2020).
A F ED DAR FOR L INEAR R EPRESENTATION A. S ETUP
We retain the setup for linear regression considered at the start of Section 3.1.
We additionally define
W ∗ ≜ [ w ∗ 1 , · · · , w ∗ M ] ⊤ ∈ R M × k as the concatenation of domain specific heads.
For notational convenience, we let ( x i,m , y i,m ) denote an (input, output) sample coming from client i and the m -th domain.
To measure the distance between any two matrices A , B with the same dimensions, we use the principal angle distance (Golub & Van Loan, 2013), given by dist( A , B ) ≜
∥ A ⊤⊥ B ∥ 2 , where
A ⊥ denotes a matrix whose columns form a basis for the orthogonal complement of the range of A .
To simplify the analysis, we further make the following assumptions.
Assumption A.1 (Sub-Gaussianilty) .
For each m ∈ [ M ] and i ∈ [ n ] , the samples x i,m ∈ R
d are independent, mean zero, have covariance I d , and has subgaussian norm 1, i.e. for every v ∈ R d , E [exp( v ⊤ x i,m )]
≤ exp( ∥ v ∥ 2 / 2) .
Assumption A.2 (Domain diversity) .
Let σ min , ∗ ≜ σ min ( 1 √ M W ∗ ) , i.e., σ min , ∗ is the minimum singular value of the head matrix.
Then σ min , ∗ > 0 .
Assumption A.3 (Ground truth normalization) .
The true domain parameters satisfy 12 √ k ≤ ∥ w ∗ m ∥ ≤ √ k for each m ∈
[ M ] , and B ∗ has orthonormal columns.
All the above assumptions aim to simplify the theoretical analysis whilst only imposing mild constraints on the data distribution and the parameters of the target functions.
Similar assumptions have also been adapted in prior work (Collins et al., 2021).
A. F ED DAR A DAPTED TO L INEAR R EGRESSION
We analyze an adapted version of our FedDAR algorithm.
Since the linear regression problem has an analytic solution, to ease analysis, we update the heads { w m } Mm =1 at the server in closed form using local gradient information.
Meanwhile, we update the representation B by taking a step using the averaged local gradients.
Algorithm 2 shows the procedure of this adapted version.
The local objective for i -th client in m -th domain at t -th iteration, f ti,m ( w m , B t ) is defined as the following,
[Equation]
where L ti,m is the number of samples from domain m at client i .
We assume in each iteration the data points { x ji,m , y ji,m } j ∈
[ L ti,m ] are all newly sampled from the distribution.
We denote L = (cid:80) m L ti,m .
Note that since the objective function has a quadratic form, thus its gradient w.r.t either w m or B has a linear form of A i,m w m − a i,n or C i,m B − c i,m which we write down explicitly in Appendix B. After every global update of the representation B , we apply an additional QR decomposition to normalize it to be column-wise orthogonal.
A. C ONVERGENCE A NALYSIS
We first provide a brief proof sketch. Overall, our approach largely follows that in Collins et al. (2021), with a few differences needed to handle the spreading of a domain’s data across different clients.
We note that we also tightened the analysis compared with Collins et al.
(2021), such that each domain only needs O ( k ) samples as opposed to O ( k 2 ) samples as in Collins et al. (2021) (where the requirement is for each client to have O ( k 2 ) samples since they considered the case where each client has a separate head). This can yield a significant improvement when k is moderately large and there is data imbalance.
First, in Lemma A.5, we show that our estimated weight matrix W t +1 ∈ R M × k (which is our estimation at time t + 1 of the true domain weights matrix W ∗
satisfies
the relationship W t +1 = W ∗ ( B ∗ ) ⊤ B t + F t ,
Input: Step size η ; number of rounds T
Client initialization: each agent i ∈ [ n ] collects L 0 samples, and sends Z i :=
(cid:80) L 0 i =
( y 0 ,j i ) 2 x 0 ,ji ( x 0 ,ji ) ⊤ to the server.
Server initialization: finds UDU ⊤ ← rank-k SVD ( 1 nL 0 (cid:80)
ni
Z i ) ; sets B 0 ← U .
for t = 0 , 1 , . . .
, T do
Server sends current B t to clients.
Client computation for W t +1 : for client i ∈
[ n ] do
Selects L new samples { ( x ji , y ji ) } .
Computes ∇ w m f ti,m ( w m , B t ) =
A ti,m w m − a ti,m for each domain m ∈ [ M ] .
Sends ( A ti,m , a ti,m , L ti,m ) back to server.
end for
Server update for W t +1 :
Server chooses w t +1 m ∈ (cid:110) w m ∈ R k : ∇ w m (cid:16) 1 (cid:80) i L ti,m (cid:80) ni
=1 f ti,m ( w m , B t ) (cid:17) = 0 (cid:111) , ∀ m ∈
[ M ] , i.e., w t +1 m that satisfies ( (cid:80) i
A ti,m ) w t
+1 m =
(cid:80) i a ti,m .
Sends W t +1 =
[ w 1 , · · · , w M ] ⊤ ∈ R M × k to clients.
Client computation for B t +1 : for client i ∈
[ n ] do
Selects L new samples { x ji , y ji } .
Computes ∇ B f t ′ i,m ( w t +1 m , B t ) =
C ti,m B t − c ti,m for each m ∈
[ M ] .
Sends ( ∇ B f t ′ i,m ( w t +1 m , B t ) , L t ′ i,m ) back to server.
for
Server t +1 :
B
[Equation]
← B
[Equation]
for where B t is our estimation of the true (high to low dimensional) representation embedding B ∗ at time t , and F t is an error term that we can show to be bounded (sufficiently small) in terms of dist( B t , B ∗ ) , with the scale of the error depending on the (random) number of samples L m seen at time t for each domain m ∈
[ M ] .
Bounding ∥ F t ∥ in our setting requires some care since the samples for each domain are spread over many clients. We refer the reader to Section A.4.2 for the details. Second, we show in Section A.4.3 that the update for B t satisfies the relationship (see equation 21)
[Equation]
where Q t := W t +1 ( B t ) ⊤ − W ∗ ( B ∗ ) ⊤ , and H Q denotes an error term which can be shown to be bounded (sufficiently small) in terms of dist( B t , B ∗ ) . Further simplifying, we have
[Equation]
By upper bounding ∥ H Q ∥ in terms of dist( B t , B ∗ ) and providing an appropriate lower bound on σ min ( W ⊤ t +1 W t +1 ) , we can then show by picking η sufficiently small and under other suitable assumptions, the quantity dist( B t , B ∗ ) decays at a linear rate (see Equation 29). Again, a difference from the analysis in Collins et al.
(2021) is our handling of bounding ∥ H Q ∥ , since the samples for each sample are spread over different clients. An issue created by the spreading of samples for a domain across different clients is how to pick an appropriate sample size such that the ∥ F t ∥ and ∥ H Q ∥ terms can be suitably bounded. In Lemma A.10, we prescribe a suitable sample size for each client such that each domain gets sufficient samples (with high probability) for the purposes of our analysis.
We now proceed with a detailed analysis. We first present a theorem that states our adapted FedDAR(Algorithm 2) enjoys linear convergence. The theorem is followed by multiple remarks which highlight key detailed points of our convergence result. Theorem A.4 (Algorithm 2 convergence) .
Define E 0 := 1 − dist 2 ( B 0 , B ∗ ) , ¯ σ max , ∗ : = σ max (cid:16) 1 √ M W ∗ (cid:17) , ¯ σ min , ∗ : = σ min (cid:16) 1 √ M W ∗ (cid:17) . Let κ := ¯ σ max , ∗ ¯ σ min , ∗ .
Suppose
[Equation] Then, for any T and any η ≤ 1 / (4¯ σ 2max , ∗ ) , with probability at least 1 − Te − 80 , [Equation]
B [Equation] , B [Equation]
Linear convergence speed:
The convergence of B T to B ∗ is linear, assuming that (1) σ min ( 1 √ M W ∗ ) > 0 and that (2) 1 − ηE 0 ¯ σ 2min ∈ (0 , 1) .
Initialization of B 0 :
For our convergence result to be meaningful, we need dist( B 0 , B ∗ ) to be close to 0.
We show in Appendix A that our algorithm’s choice of initial B 0 ensures that dist( B 0 , B ∗ ) is close enough to 0 whilst preserving privacy.
When the number of samples is uniform across the domains, this comes only at the cost of a logarithmic increase in sample complexity.
Sample complexity: The per-iteration sample complexity per client is L .
We note that in the requirement for L (9), we need that L ≥ Ω( dk 2 κ 4 /n ) ; this comes from the updates for B t ∈ R d × k .
While we expect that d could be large, a large number of clients n helps to mitigate the increase in sample complexity arising from d .
We also need L ≥ Ω( k 2 κ 4 (cid:80) mi =1 π i,m ) for every domain m ∈
[ M ] ; this requirement comes from the updates for w tm for each of the M domains.
A. P ROOF OF T HEOREM A. A.4. A NALYSIS FOR UPDATE OF W t +
Since we are analyzing the update step for any iteration t , unless necessary we drop all t superscripts.
Let L m =
(cid:80) ni
=1 L i,m denote the number of samples from domain m ∈
[ M ] across the n clients.
Then, we can express ∇ w m (cid:80) ni
{+=
+}
f i,m ( w m , B ) as ∇ w m n (cid:88) i =
[Equation]
Since
[Equation]
it follows that following Algorithm 2, 
[Equation]
Reexpressing, assuming G m is invertible [Equation] we have
[Equation]
Intuitively, assuming L m is large enough,
[Equation]
Hence,
[Equation]
This then implies that
[Equation]
where the m -th row of F is
[Equation]
Note the similarity of equation 18 to (17) in (Collins et al., 2021).
Following a similar analysis as (Collins et al., 2021), we should also be able to bound the Frobenius norm of F in terms of dist( B , B ∗ ) .
Below, we formalize the argument.
First, we have the following lemma.
Lemma A.5 (Update for W t +1 ) .
For each time t , let L tm := (cid:80) ni
=1 L ti,m denote the number of samples from domain m ∈
[ M ] across the n clients at time t .
For convenience, we drop the time index unless absolutely necessary.
We define the terms
[Equation]
Then, assuming that G m is invertible, the update for W takes the form
[Equation]
where the m -th row of F is
[Equation]
Proof.
We can express ∇ w m (cid:80) ni
{+=
+}
f i,m ( w m , B ) as
[Equation]
Since
[Equation]
it follows that following Algorithm 2, 
[Equation]
Reexpressing, assuming G m is invertible, we have
[Equation]
This then implies that
[Equation]
where the m -th row of F is
[Equation]
A.4. B OUNDING ∥ F ∥ F
We will proceed to bound the Frobenius norm of F .
We begin by showing that G − 1 m exists and (both lower and upper) bounding its spectral norm [Equation]
Lemma A.6.
Let L min :=
min m ∈ [ M ] L m .
Let δ k := 10 C √ k √ log( M ) √
L min for some absolute constant C .
Suppose that 0 ≤ δ k < 1 .
Then, with probability at least 1 − e − 80 k log( M ) , G − 1 m exists for each m ∈
[ M ] , and
[Equation]
[Equation]
Proof.
Note that
[Equation]
Let v ji,m :=
B ⊤ x ji,m .
Since B ⊤ B
=
I , it follows that each v ji,m is i.i.d 1-subgaussian.
Then, applying the same argument in Theorem 4.6.1 of Vershynin 2018, we have (cf.
equation (4.22) in
Vershynin 2018)
[Equation]
with probability at least 1 − e − z 2 for z ≥ 0 and some absolute constant C , assuming that 0 ≤ δ k,m ≤ 1 .
Consider the choice z = 9 √ k (cid:112) log( M ) .
Then,
[Equation]
Suppose we choose L min ≥ 1 such that δ k,m < 1 .
Then, taking a union bound, with probability at least 1 − Me − z 2 =
[Equation]
Therefore, with probability at least 1 − exp( − 80 k log( M )) , G − 1 m exists for every m ∈
[ M ] , and in addition, ∥ G − 1 m ∥ 2 ≤ [Equation] − δ
[Equation]
m ∈ [ M ] . We next bound the operator norm of term B ⊤ X m
( I − BB ⊤ ) B ∗ .
Lemma A.7.
Let L min :=
min m ∈ [ M ] L m .
Let δ k := 10 C √ [Equation] min for some absolute constant C .
Suppose L min is such that 0 ≤ δ k < 1 .
Then, with probability at least 1 − e − 50 k log M , ∥ B ⊤ X m
[Equation] ≤ dist( B ∗ , B ) δ k .
Proof.
We will use an ϵ -net argument, similar to the proof of Theorem 4.6.1 in (Vershynin, 2018).
First, by Corollary 4.2.13 in (Vershynin, 2018), there exists an 1 / 4 -net N of the unit sphere S k −with cardinality N ≤ 9 k .
Using Lemma 4.4.1 in (Vershynin, 2018), we have that
[Equation]
To prove our result, by applying a union bound over m ∈ [ M ] , it suffices to show that with the probability at least 1 − e − 100 k 2 log M ,
[Equation]
where we recall that
[Equation]
We will assume that min m L m :=
L min ≥ 1 is chosen large enough such that δ k,m ≤ 1 .
For a fixed z ∈ S k − 1 , observe that
[Equation]
, where we defined u ji,m := B ⊤ x ji,m , and v ji,m = ( B ∗ ) ⊤ ( I − BB ⊤ ) x ji,m .
Since each x ji,m is 1-subgaussian, ∥ B ∥ 2 = 1 , and ∥ ( I − BB ⊤ ) B ∗ ∥ 2 = dist( B ∗ , B ) , it follows that z ⊤ u ji,m is subgaussian with norm at most 1, and ( v ji,m ) ⊤ z is subgaussian with norm at most dist( B ∗ , B ) .
Thus, the random variable α ji,m :=
( z ⊤ u ji,m )(( v ji,m ) ⊤ z ) (for a fixed unit z ) is subexponential with sub-exponential norm at most dist( B ∗ , B ) .
Moreover, note that α ji,m is mean-zero, since
[Equation]
as x ji,m is assumed to have identity covariance.
Thus, the α ji,m ’s are i.i.d mean-zero subexponential variables each with subexponential norm at most dist( B ∗ , B ) .
Hence, by Bernstein’s inequality (cf.
Corollary 2.8.3 in (Vershynin, 2018)),
[Equation]
Above we used the assumption that δ k,m ≤ 1 to simplify the minimum operator in the exponent.
Taking a union bound over each z ∈ N , it follows that
[Equation]
where the last inequality follows by picking C large enough (but still it is an absolute constant).
(Here the choice of 51 in the exponent is somewhat arbitrary; any choice smaller than 81 should work).
By applying a union bound over the domains m ∈ [ M ] , this then completes our proof.
We are now finally ready to bound ∥ F ∥ F .
Lemma A.8.
Let L min :=
min m ∈ [ M ] L m .
Let δ k := 10 C √ k √ log( M ) √
L min for some absolute constant C .
Suppose that 0 ≤ δ k < 1 .
Then, with probability at least 1 − 2 e − 50 k log( M ) ,
[Equation]
≤
[Equation]
∗ ∥ F .
Proof.
By Lemma A.6 and Lemma A.7, we have that with probability at least 1 − 2 e − 50 k log M ,
[Equation]
B
[Equation]
≤
[Equation]
dist
[Equation]
The proof then follows by recalling that the m -th row, F ⊤ m , takes the form
[Equation]
A.4. A NALYSIS OF UPDATE FOR B t +
Similarly to (Collins et al., 2021), we define
Q
[Equation]
= W
[Equation]
B [Equation] −
[Equation]
W [Equation] [Equation] B
[Equation]
⊤ . Below, we drop the time index and use B , Q, W to denote B t , Q t , and W t +1 respectively.
Based on algorithm 2, we have that
[Equation]
[Equation]
m
[Equation]
M
[Equation]
m
[Equation]
L m
[Equation]
[Equation]
( A ji
[Equation]
m )
⊤ W,
A ji
[Equation]
m :
[Equation]
e m
[Equation]
x ji
[Equation]
m ) ⊤ =
B − η M
M (cid:88) m
[Equation]
L m n (cid:88) i =
L i
[Equation]
m (cid:88)
j = (cid:16)(cid:68)
A ji,m , Q (cid:69)(cid:17) ( A ji,m ) ⊤ W
= B − η M
M (cid:88) m =
L m n (cid:88)
i [Equation] Above, we define q m ∈ R d to denote the m -th row of Q (viewed as a column vector).
Note again that since
[Equation]
the term H Q in equation 21 can be appropriately bounded.
Note the resemblance of equation 21 to (53) in (Collins et al., 2021); the crucial difference is that we will need to lower bound 1 m σ 2min ( W ∗ ) , instead of 1 n σ 2min ( W ∗ ) as in (Collins et al., 2021).
Thus we should be able to carry out the rest of the analysis in a similar way to the outline in (Collins et al., 2021) and derive an analogous result to
Theorem 1 in (Collins et al., 2021).
We first bound the error term H Q .
Lemma A.9.
Let
[Equation]
Let γ k := 20 k √ d c √ nL for some absolute constant c .
Suppose that 0 ≤ γ k < k .
Then, for any t , with probability at least 1 − exp( − 90 d ) − 2 e − 50 k log M ,
[Equation]
Proof.
As before, we may omit the time superscript t in cases where it is clear for notational convenience.
The proof is based on the argument in Lemma 5 in (Collins et al., 2021).
Again, the main tool is an ϵ -net argument.
We first bound ∥ q m ∥ 2 and ∥ w m ∥ 2 .
Bounding q m : With probability at least 1 − 2 e − 50 k log M , for each m ∈
[ M ] , we have that
[Equation]
Above, we utilized the assumption that ∥ w ∗ m ∥ 2 ≤ √ k , the orthonormality of B t (which was derived as the orthogonal matrix from a Gram-Schmidt procedure), the assumption that 0 < δ k ≤ 1 / 2 , as well Lemma A.8 which bounds ∥ F m ∥ 2 (for all m ) with probability at least 1 − 2 e − 50 k log M .
Bounding w m : Note that for notational convenience, we let w m denote w t
+1 m .
For each t and every m ∈ [ M ] , we have that
[Equation]
with probability at least 1 − 2 e − 50 k log M , where again we used Lemma A.8 to handle ∥ F m ∥ 2 , the assumption that δ k < 1 / 2 , and the fact that dist( B t , B ∗ ) ≤ 2 .
For the rest of the proof, we condition on the event
[Equation]
ϵ -net argument to bound H Q :
Again, note that there exists an 1 / 4 -net N k of the unit sphere
S k − 1 and an 1 / 4 -net N d of the unit sphere S d − 1 with cardinalities less than or equal to 9 k and 9 d respectively.
Note now that by Equation 4.13 in (Vershynin, 2018), we have
[Equation] [Equation]
Fix now a u ∈ N [Table] Note now that
[Table] [Equation] [Table]
( x j i,m ) ⊤ q m w ⊤ m v (cid:17) is subexponential with norm less than or equal to ∥ q m ∥
2 ∥ w m ∥ 2 ≤ 6 k dist( B t , B ∗ ) , since it is the product of two subgaussian variables u ⊤ x ji,m and ( x ji,m ) ⊤ q m w ⊤ m v with subgaussian norms bounded by 1 and ∥ q m ∥
∥ w m ∥ 2 respectively.
Note also that
[Equation]
Thus, by Bernstein’s inequality, carrying on from equation 22, we have that
[Equation]
where we will choose ρ such that ρ k dist( B t , B ∗ ) ≤ 1 to simplify the exponent in the way we did, and c is an absolute constant that may change from line to line.
Above, we also used the fact that (cid:80)
Mm =
[Equation]
m = nL (recall that L is the total number of samples per agent and there are n agents).
Consider the choice
[Equation]
Then,
[Equation]
≤ [Equation] d [Equation]
Taking a union bound over all u ∈ N d and v ∈ N k , it follows then that
[Equation]
where above we used the fact that d ≥ k .
A.4. C OMBINING EARLIER ARGUMENT : CONVERGENCE OF F ED DAR
As seen in Lemma A.8, we require that L min :=
min m ∈ [ M ] L m to be lower bounded.
However, since
L m is a stochastic variable, we are unable to directly lower bound it.
Below, we provide a result that converts a lower bound on each client’s sample size L (a deterministic quantity we can control) to a high-probability lower bound on L min .
Lemma A.10.
Let L min :=
min m ∈ [ M ] L m .
For any α > 0 , suppose that for each m ∈
[ M ] ,
[Equation]
Then, with probability at least 1 − exp( − 90) ,
L min ≥ α.
Proof.
Note that
[Equation]
, which is a sum of nL independent random variables bounded between 0 and 1.
Moreover,
E [ L m ] = n (cid:88) i =
π i,m L, where π i,m is the probability that a datapoint comes from domain m for client i .
Note finally that
[Equation]
Hence, by Bernstein’s inequality, it follows that for any s > 0 ,
P
[Equation]
Since we wish to perform union bound over the M domains, we seek to choose s and L such that exp (cid:32) − s 2 /(cid:80)
ni =1 (cid:80)
Lj =1 π i,m + s/(cid:33) ≤ exp ( − 91 log M ) , so that
M exp (cid:32) − s 2 /(cid:80)
ni =1 (cid:80)
Lj =1 π i,m + s/(cid:33) ≤ M exp ( − 91 log M ) ≤ exp ( − 90 log M ) .
To this end, note that we need s 2 /(cid:80) ni =
[Equation]
(cid:80) Lj =1 π i,m
[Equation]
[Equation]
≥ 91 log M ⇐⇒ s 2 ≥ 2 · 91 log M   n (cid:88) i = L (cid:88)
j = π i,m + s/ 
⇐⇒ s ≥ (cid:118)(cid:117) (cid:117)(cid:116) 182 log M n (cid:88) i
=
π i,m L + (cid:18) 182 log M (cid:19)+ 182 log M
Suppose we pick L such that n (cid:88)
i =
π i,m L ≥ 182 log M, so that
[Equation]
Then, by picking s = 2 (cid:112)(cid:80) ni =1 π i,m L , it follows that
[Equation]
such that for each m ∈
[ M ] ,
[Equation]
By choosing L such that
[Equation]
it follows that
[Equation]
The result now follows by choosing L such that it also satisfies
[Equation]
for each m .
Lemma A.11 (Descent lemma) .
Define E 0 := 1 − dist 2 ( B 0 , B ∗ ) and ¯ σ max , ∗ :=
σ max (cid:16) 1 √ M W ∗ (cid:17) and ¯ σ min , ∗ :
= σ min (cid:16) 1 √ M W ∗ (cid:17) .
Let κ := ¯ σ max , ∗ ¯ σ min , ∗ .
Consider any iteration t .
Suppose that
[Equation]
where c > 0 is absolute constant.
Suppose also that
[Equation]
which by Lemma A.10, ensures that with probability at least 1 − e − 90 ,
[Equation]
[Equation] where L t min = min m ∈ [ M ] L t m denotes the minimum number of samples from any domain at iteration t , and C > 0 is an absolute constant.
Then, for any η ≤ 1 / (4¯ σ 2max , ∗ ) , we have dist
[Equation]
with probability at least 1 − e − 80 .
Proof.
We begin with the observation that
[Equation]
where
[Equation]
and
[Equation]
Above ¯ B t +1 denotes the estimate of B before we perform the QR decomposition.
We note that the updates for W and B are exactly analogous to the updates for W and B as seen in the proof of
Lemma 6 in (Collins et al., 2021).
The only two differences are 1.
The definitions of F in our paper and (Collins et al., 2021) are slightly different.
However, in both cases,
[Equation]
for some term δ k ≤ 1 / 2 with high probabilities.
In our case, this event holds with probability at least 1 − 2 e − 50 k log M , whilst in (Collins et al., 2021), the event holds with probability at least 1 − exp( − 110 k 2 log n ) . The update for ¯ B t +1 in (Collins et al., 2021) takes the form
[Equation]
where 0 ≤ r ≤ 1 is a ratio term used in (Collins et al., 2021), and m above represents the number of samples used by each learner in (Collins et al., 2021)
(which is different from our use of m as an index over the domains).
However, we note that with high probabilities,
[Equation]
where the definition of γ k in both papers differ but both satisfy the assumption that γ k ≤ k.
Due to these similarities in the updates for W t +1 and B t +1 with the update in (Collins et al., 2021), the proof of this lemma follows naturally from the proof of Lemma 6 in (Collins et al., 2021), by plugging in ηM ( Q t )
⊤ W t +1 in the update for ¯ B t +1 in place of ηrn ( Q t )
⊤ W t +1 as in (Collins et al., 2021).
In particular, following the same analysis as in (Collins et al., 2021), we see that on the events in Lemma A.8 and Lemma A.9, following the equation immediately after Equation (84) in (Collins et al., 2021), we have
[Equation]
where in our case ¯ δ k
= δ k + γ k .
Then, by choosing
[Equation]
k <
[Equation]
[Equation]
it follows that ¯ δ k < 1 / 5 , and so
[Equation]
as in equation (85) in (Collins et al., 2021), such that
[Equation]
It remains for us to understand what the constraint on ¯ δ k spelt out in equation 25, and the constraints on δ k and γ k (in Lemmas A.8 and A.9 respectively) mean in our choice of the sample size L for each agent, and the domain size L m at each iteration.
Observe that we need
[Equation]
where c, C > 0 are absolute constants.
By choosing
[Equation]
we ensure that the requirements in equation 26, equation 27 and equation 28 are all satisfied.
The final result then follows by applying Lemma A.10.
This then yields the following convergence result, which is a more complete statement of A.4.
Theorem A.12 (Convergence result for Algorithm 2) .
Define E 0 := 1 − dist 2 ( B 0 , B ∗ ) and ¯ σ max , ∗ :=
σ max (cid:16) 1 √ M W ∗ (cid:17) and ¯ σ min , ∗ :
= σ min (cid:16) 1 √ M W ∗ (cid:17) .
Let κ := ¯ σ max , ∗ [Equation]
Suppose that
[Equation]
where c > 0 is absolute constant.
Suppose also that
[Equation]
Then, for any η ≤ 1 / (4¯ σ 2max , ∗ ) , we have
[Equation]
, B
[Equation]
dist
[Equation] B
[Equation]
B
[Equation]
, with probability at least 1 − e − 80 . Then for any T and any η ≤ 1 / (4 σ 2max , ∗ ) , we have
[Equation]
with probability at least 1 −
Te − 80 .
By assuming that σ 2min , ∗ > 0 , the bound in Theorem 1 decays exponentially.
We note that the total number of samples required per client scales with L log(1 /ϵ ) .
In addition, in order for the result to be meaningful, we implicitly assume that E 0 is close to 1 such that
[Equation]
To do so, we note it is possible to choose B 0 such that dist( B 0 , B ∗ ) is close enough to 0, with only a logarithmic increase in sample complexity when the number of samples is uniform across the domains.
The argument follows the proof of Theorem 3 in (Tripuraneni et al., 2021).
Theorem A.13.
Suppose Assumptions A.1, A.2, A.3 all hold.
Suppose also that x 0 ,ji ∼ N (0 , I d ) independently for all i ∈ [ n ] .
Suppose each client i sends the server Z i :
= (cid:80) L
j =
( y 0 ,j i ) 2 x ji ( x ji ) ⊤ , as well as the integer value of L i , such that the server can compute Z := 1 nL 0 (cid:80) ni
=1 Z i .
Then, the server computes UDU ⊤ ← rank-k SVD ( Z ) , and sets B 0 := U .
Let
[Equation]
where m ( i, j ) denotes the sample of the j -th sample from the i -th client.
Let σ min , ∗ :
= σ min (¯Λ) , and let σ max , ∗ :
= σ max (¯Λ) .
Suppose that L 0 ≥ c polylog( d, nL 0 ) σ max , ∗ dk 2 / ( nσ 2min , ∗ ) .
Then, with probability at least 1 − ( nL 0 ) − 100 , we have that
[Equation]
In particular, when the number of samples is uniform across the domains, we have that
[Equation]
where we recall that κ := ¯ σ max , ∗ / ¯ σ min , ∗ , and
[Equation]
Proof.
We omit the proof since it is a slight variant of Theorem 3 in (Tripuraneni et al., 2021).
For completeness, note that in the case when the number of samples is uniform across the domains, some algebra shows that
[Equation]
However, since k/ 4 M ≤ [Equation] [Equation] , we have that
[Equation]
which proves the last statement in the theorem.
B A DDITIONAL E XPERIMENTAL R ESULTS B. E XPERIMENTS ON F AIR F ACE DATASET FOR GENDER CLASSIFICATION
[Table]
We also conduct experiments for gender classification on FairFace with the same settings.
The best representation dimension is k = 2 for this task, probably due to the smaller diversity across the domains.
We can see that the results shown in Table 4 have a similar trend as the results in Table 5.
[Table]
B. E XPERIMENTS ON DIGITS DATASET
[Table]
We perform additional experiments on the digits dataset with five data domains with feature shift (Li et al., 2021b).
Details are described in the following paragraphs.
From Table 5, we can see that FedDAR-WA outperforms FedAvg consistently except in the case where domain distributions are extremely heterogeneous ( α = 0 . 1 ).
In this case, each client tends to have data from only one domain.
It is difficult for the proposed method to learn a good domain-specific head for the domain with the most different data (more obvious feature shift) under this circumstance.
For other levels of heterogeneity, although the min and max domain accuracies are similar between FedAvg and FedDAR-WA, the average accuracies are improved as a result of the domain-wise personalized model.
On the other hand, without an alternative update of the head and representation, FedAvg + Multi-head will overfit quickly.
We don’t include the results of FedDAR-SA here because using representation dimension k ≥ 64 causes numerical instability during head aggregation and failure to converge.
While using representation dimension k ≤ 32 leads to lower accuracy.
Datasets.
We use the same digits dataset containing five different data domains as (Li et al., 2021b).
Specifically, we use SVHN (Netzer et al., 2011), USPS (Hull, 1994), SynthDigits (Ganin &
Lempitsky, 2015), MNIST-M (Ganin & Lempitsky, 2015) and MNIST (LeCun et al., 1998) as five data domains.
Similarity to the experiments on FairFace datraset, the training data is divided into n clients without duplication.
Each client has a domain distribution π i ∼ Dir ( α p ) sampled from a
Dirichlet distribution.
Implementation Details.
We adapt the codebase from (Li et al., 2021b).
A 6-layer CNN withconvolutional layers and 3 fully-connected layers is used, with the last layer as domain-specific head.
We use SGD optimizer with learning rate 10 − 2 and cross-entropy loss.
The batch size is set to 32 , and the total communication rounds is set to 100 .
For each method, we first train the model forrounds with 1 local epoch using FedAvg as warmup.
The accuracy shown is the average over the last ten communication rounds.
We repeat experiment for each setting three times with different random seeds and report the averages.
B. F URTHER EXPERIMENTAL DETAILS B.3. S YNTHETIC D ATA
For the synthetic data experiments, we adapt the code from (Collins et al., 2021) and follow a similar protocol.
The ground-truth matrices W ∗ ∈ R M × k and B ∗ ∈
R
d × k are generated following the same way as (Collins et al., 2021) by sampling each element from i.i.d.
standard normal distribution and taking the QR factorization.
The same L samples are used for each client during the whole training process.
Test samples are generated in the same way as the traning samples but without noise.
For all the methods, models are initalized with ramdom Gaussian samples.
We set α = 0 .
4 for experiments in Figure 5.2.
B.3. R EAL DATA WITH CONTROLLED DISTRIBUTION
Implementation details.
We use Imagenet(Deng et al., 2009) pre-trained ResNet-34 (He et al., 2016) for all experiments on this dataset.
All the methods are trained for T = 100 communication rounds, with 20 rounds of FedAvg as warmup.
For FedDAR-WA and FedDAR-SA, 5 epochs of local updates are executed for both heads and representation at each round.
For the baselines, 5 epochs of local updates are executed at each round for fair comparison.
We use Adam optimizer with a learning rate of 1 × 10 − 4 for the first 60 rounds and 1 × 10 − 5 for the last 40 rounds.
The images are resized to 224 × 224 with only random horizontal flip for augmentation.
The learning rate and the number of local epochs is tuned by grid search with a fixed batch size of 64.
We tuned the projection dimension k for FedDAR-SA among {4,8,16,32,64} with α = 1 . 0 and used k = 8 for all other α .
Our evaluation metrics are the classification accuracy on the whole validation set of FairFace for each race group.
We don’t have extra local validation set to each client since we assume the data distribution within each domain is consistent across the clients.
The numbers reported are the average over the final 10 rounds of communication following the standard practice in (Collins et al., 2021), and the average of three independent runs with different random seeds.
B.3. R EAL DATA WITH REAL -W ORLD DATA DISTRIBUTION
Dataset details.
The detailed statistics of the partial EXAM dataset is summarized in Table 6.
The "Other" category includes American Indian or Alaska native, native Hawaiian or other Pacific islander and patients with more than one race or unknown race.
≥ HFO % means the percentage of cases with positive labels (receiving oxygen therapy higher or equal to high-flow oxygen with 72 hours).
[Table]
Implementation details.
We apply 5-fold cross validation.
The input of the model is one chest x-ray image resized to 224x224 paired with a 22-dimensional electronic health record(EHR) data, the representation dimension is 278 if it is not projected.
All the models are trained for T =communication rounds with Adam optimizer and a learning rate of 1 × 10 − 4 .
For each round we do 1 local epoch for all the methods.
For all the methods, the models are initialized with the same pretrained model as in (Dayan et al., 2021) without any warmup.
For FedDAR-SA and FedDARWA, we execute 5 epochs of update for heads on each round, and set representation dimension k = 16 for FedDAR-SA.
Hyperparameters including learning rate, number of epochs for head update and representation dimension are tuned through grid search with a fixed batch size of 36.
For FedRep,FedDARand FedPer.
For LG-FedAvg, we treated the last fully-connected layer as the global parameters and all other layers as local representation.
For FedMinMax, multiple local iterations are executed during each round instead of one step of GD for reasonable comparison.
For FedProx we tuned µ among { 0 . 05 , 0 . 1 , 0 . 25 , 0 . 5 } and used µ = 0 .
For the fine-tuning methods, we only fine-tune the global trained model locally with Adam optimizer and learning rate of 5 e − 5 forepoch since more epochs of fine-tuning leads to worse results.
The models are evaluated by aggregating predictions on the local validation sets then calculating the area under curve (AUC) for each domain. The average AUCs on local validation set of clients are also reported. The AUC shown is first averaged over the last five communication rounds, and then averaged over five runs of 5-fold cross validation. C D ISCUSSION ON COMMUNICATION AND PRIVACY
Communication For FedDAR-WA, the only communication overhead comes from the extra parameters of multiple heads for different domains, which only slightly increase the communication cost.
For FedDAR-SA, we need to send a Hessian with k 2 × N 2 parameters from each client to the server at each round.
This might be costly when both representation dimension k and output dimension N are large.
However, compared to sending millions of parameters of neural network, the extra communication cost is acceptable.
Privacy For the FedDAR-WA, there is no extra parameters shared compared to FedAvg.
So there is no additional privacy risk introduced.
Privacy techniques like homomorphic encryption (Cheon et al., 2017) or differential privacy (McMahan et al., 2017b; Kairouz et al., 2021) that apply to FedAvg also works for FedDAR-WA.
In fact, the multi-head design of different domains makes it harder to perform gradient based attack(Zhu et al., 2019) targeting our method.
Because the attacker need to first figure out which domain the sample comes from.
For FedDAR-SA, the only extra parameters shared is the Hessian matrices, which are aggregated results from all the local data.
Recovering the information for a specific sample from Hessian is extremely difficult.
Under the worst circumstance, what the attacker can recover from the Hessian is the label and the features at last layer, which hardly ease the difficulty of recovering original input.