T OWARDS U NDERSTANDING AND M ITIGATING D IMENSIONAL C OLLAPSE IN H ETEROGENEOUS F EDERATED L EARNING
A BSTRACT
Federated learning aims to train models collaboratively across different clients without sharing data for privacy considerations.
However, one major challenge for this learning paradigm is the data heterogeneity problem, which refers to the discrepancies between the local data distributions among various clients.
To tackle this problem, we first study how data heterogeneity affects the representations of the globally aggregated models.
Interestingly, we find that heterogeneous data results in the global model suffering from severe dimensional collapse , in which representations tend to reside in a lower-dimensional space instead of the ambient space. Moreover, we observe a similar phenomenon on models locally trained on each client and deduce that the dimensional collapse on the global model is inherited from local models.
In addition, we theoretically analyze the gradient flow dynamics to shed light on how data heterogeneity result in dimensional collapse for local models.
To remedy this problem caused by the data heterogeneity, we propose F ED D ECORR , a novel method that can effectively mitigate dimensional collapse in federated learning.
Specifically, F ED D ECORR applies a regularization term during local training that encourages different dimensions of representations to be uncorrelated.
F ED D ECORR , which is implementation-friendly and computationally-efficient, yields consistent improvements over baselines on standard benchmark datasets.
Code: https://github.com/Yujun-Shi/FedCLS.
I NTRODUCTION
With the rapid development deep learning and the availability of large amounts of data, concerns regarding data privacy have been attracting increasingly more attention from industry and academia.
To address this concern, McMahan et al. (2017) propose Federated Learning —a decentralized training paradigm enabling collaborative training across different clients without sharing data.
One major challenge in federated learning is the potential discrepancies in the distributions of local training data among clients, which is known as the data heterogeneity problem.
In particular, this paper focuses on the heterogeneity of label distributions (see Fig. 1(a) for an example).
Such discrepancies can result in drastic disagreements between the local optima of the clients and the desired global optimum, which may lead to severe performance degradation of the global model.
Previous works attempting to tackle this challenge mainly focus on the model parameters, either during local training (Li et al., 2020; Karimireddy et al., 2020) or global aggregation (Wang et al., 2020b).
However, these methods usually result in an excessive computation burden or high communication costs (Li et al., 2021a) because deep neural networks are typically heavily over-parameterized.
In contrast, in this work, we focus on the representation space of the model and study the impact of data heterogeneity. To commence, we study how heterogeneous data affects the global model in federated learning in Sec. 3.1.
Specifically, we compare representations produced by global models trained under different degrees of data heterogeneity. Since the singular values of the covariance matrix provide a comprehensive characterization of the distribution of high-dimensional embeddings, we use it to
[Figure]
H o m o .
D a t a
H e t e r o .
D a t a ( [Figure] ) Data Heterogeneity study the representations output by each global model.
Interestingly, we find that as the degree of data heterogeneity increases, more singular values tend to evolve towards zero.
This observation suggests that stronger data heterogeneity causes the trained global model to suffer from more severe dimensional collapse , whereby representations are biased towards residing in a lower-dimensional space (or manifold).
A graphical illustration of how heterogeneous training data affect output representations is shown in Fig.
1(b-c).
Our observations suggest that dimensional collapse might be one of the key reasons why federated learning methods struggle under data heterogeneity. Essentially, dimensional collapse is a form of oversimplification in terms of the model, where the representation space is not being fully utilized to discriminate diverse data of different classes. Given the observations made on the global model, we conjecture that the dimensional collapse of the global model is inherited from models locally trained on various clients.
This is because the global model is a result of the aggregation of local models.
To validate our conjecture, we further visualize the local models in terms of the singular values of representation covariance matrices in Sec. 3.2. Similar to the visualization on the global model, we observe dimensional collapse on representations produced by local models. With this observation, we establish the connection between dimensional collapse of the global model and local models.
To further understand the dimensional collapse on local models, we analyze the gradient flow dynamics of local training in Sec. 3.3. Interestingly, we show theoretically that heterogeneous data drive the weight matrices of the local models to be biased to being low-rank, which further results in representation dimensional collapse. Inspired by the observations that dimensional collapse of the global model stems from local models, we consider mitigating dimensional collapse during local training in Sec. 4.
In particular, we propose a novel federated learning method termed F ED D ECORR .
F ED D ECORR adds a regularization term during local training to encourage the Frobenius norm of the correlation matrix of representations to be small.
We show theoretically and empirically that this proposed regularization term can effectively mitigate dimensional collapse (see Fig. 1(d) for example).
Next, in Sec. 5, through extensive experiments on standard benchmark datasets including CIFAR10, CIFAR100, and TinyImageNet, we show that F ED D ECORR consistently improves over baseline federated learning methods.
In addition, we find that F ED D ECORR yields more dramatic improvements in more challenging federated learning setups such as stronger heterogeneity or more number of clients.
Lastly, F ED D ECORR has extremely low computation overhead and can be built on top of any existing federated learning baseline methods, which makes it widely applicable.
Our contributions are summarized as follows.
First, we discover through experiments that stronger data heterogeneity in federated learning leads to greater dimensional collapse for global and local models.
Second, we develop a theoretical understanding of the dynamics behind our empirical discovery that connects data heterogeneity and dimensional collapse.
Third, based on the motivation of mitigating dimensional collapse, we propose a novel method called F ED D ECORR , which yields consistent improvements while being implementation-friendly and computationally-efficient.
R ELATED W ORKS
Federated Learning.
McMahan et al.
(2017) proposed FedAvg, which adopts a simple averaging scheme to aggregate local models into the global model.
However, under data heterogeneity, FedAvg suffers from unstable and slow convergence, resulting in performance degradation.
To tackle this challenge, previous works either improve local training (Li et al., 2021b; 2020; Karimireddy et al., 2020; Acar et al., 2021; Al-Shedivat et al., 2020; Wang et al., 2021) or global aggregation (Wang et al., 2020b; Hsu et al., 2019; Luo et al., 2021; Wang et al., 2020a; Lin et al., 2020; Reddi et al., 2020; Wang et al., 2020a).
Most of these methods focus on the model parameter space, which may result in high computation or communication cost due to deep neural networks being overparameterized.
Li et al.
(2021b) focuses on model representations and uses a contrastive loss to maximize agreements between representations of local models and the global model.
However, one drawback of Li et al. (2021b) is that it requires additional forward passes during training, which almost doubles the training cost.
In this work, based on our study of how data heterogeneity affects model representations, we propose an effective yet highly efficient method to handle heterogeneous data.
Another research trend is in personalized federated learning (Arivazhagan et al., 2019; Li et al., 2021c; Fallah et al., 2020; T Dinh et al., 2020; Hanzely et al., 2020; Huang et al., 2021; Zhang et al., 2020), which aims to train personalized local models for each client.
In this work, however, we focus on the typical setting that aims to train one global model for all clients.
Dimensional Collapse.
Dimensional collapse of representations has been studied in metric learning (Roth et al., 2020), self-supervised learning (Jing et al., 2021), and class incremental learning (Shi et al., 2022).
In this work, we focus on federated learning and discover that stronger data heterogeneity causes a higher degree of dimensional collapse for locally trained models.
To the best of our knowledge, this work is the first to discover and analyze dimensional collapse of representations in federated learning.
Gradient Flow Dynamics.
Arora et al.
(2018; 2019) introduce the gradient flow dynamics framework to analyze the dynamics of multi-layer linear neural networks under the ℓ 2 -loss and find deeper neural networks biasing towards low-rank solution during optimization.
Following their works, Jing et al. (2021) finds two factors that cause dimensional collapse in self-supervised learning, namely strong data augmentation and implicit regularization from depth.
Differently, we focus on federated learning with the cross-entropy loss.
More importantly, our analysis focuses on dimensional collapse caused by data heterogeneity in federated learning instead of depth of neural networks.
Feature Decorrelation.
Feature decorrelation had been used for different purposes, such as preventing mode collapse in self-supervised learning (Bardes et al., 2021; Zbontar et al., 2021; Hua et al., 2021), boosting generalization (Cogswell et al., 2015; Huang et al., 2018; Xiong et al., 2016), and improving class incremental learning (Shi et al., 2022).
We instead apply feature decorrelation to counter the undesired dimensional collapse caused by data heterogeneity in federated learning.
D IMENSIONAL C OLLAPSE C AUSED BY D ATA H ETEROGENEITY
In this section, we first empirically visualize and compare representations of global models trained under different degrees of data heterogeneity in Sec. 3.1.
Next, to better understand the observations on global models, we analyze representations of local models in Sec. 3.2. Finally, to theoretically understand our observations, we analyze the gradient flow dynamics of local training in Sec. 3.3.
3. E MPIRICAL O BSERVATIONS ON THE G LOBAL M ODEL
We first empirically demonstrate that stronger data heterogeneity causes more severe dimensional collapse on the global model.
Specifically, we first separate the training samples of CIFAR100 into 10 splits, each corresponding to the local data of one client.
To simulate data heterogeneity among clients as in previous works (Yurochkin et al., 2019; Wang et al., 2020a; Li et al., 2021b), we sample a probability vector p c = ( p c, 1 , p c, 2 , . . . , p c,K ) ∼ Dir K ( α ) and allocate a p c,k proportion of instances of class c ∈
[ C ] = { 1 , 2 , . . , C } to client k ∈ [ K ] , where Dir K ( α ) is the Dirichlet distribution with K categories and α is the concentration parameter.
A smaller α implies stronger data heterogeneity ( α = ∞ corresponds to the homogeneous setting).
We let α ∈ { 0 . 01 , 0 . 05 , 0 .
25 , ∞} .
For each of the settings generated by different α ’s, we apply FedAvg (McMahan et al., 2017) to train a MobileNetV2 (Sandler et al., 2018) with CIFAR100 (observations on other federated learning methods, model architectures, or datasets are similar and are provided in Appendix D).
Next, for each of the four trained global models, we compute the covariance matrix Σ = 1 N (cid:80) Ni
=
( z i − [-¯ z )( z i − ¯ z ) ⊤ , where z i is therepresentation given the i -th test data point, N is number of test samples, and ¯ z = 1 N (cid:80) Ni =
-]l og
s i
ngu l a r v a l ue s [Figure] Finally, we apply the singular value decomposition (SVD) on each of the covariance matrices and visualize the top 100 singular values in Fig. 2(a).
If we define a small value τ as the threshold for a singular value to be significant (e.g., log τ = − 2 ), we observe that for the homogeneous setting, almost all the singular values are significant, i.e., they surpass τ .
However, as α decreases, the number of singular values exceeding τ monotonically decreases.
This implies that with stronger heterogeneity among local training data, the representation vectors produced by the trained global model tend to reside in a lower-dimensional space, corresponding to more severe dimensional collapse. 3. E MPIRICAL O BSERVATIONS ON L OCAL M ODELS
Since the global model is obtained by aggregating locally trained models on each client, we conjecture that the dimensional collapse observed on the global model stems from the dimensional collapse of local models.
To further validate our conjecture, we continue to study whether increasing data heterogeneity will also lead to more severe dimensional collapse on locally trained models.
Specifically, for different α ’s, we visualize the locally trained model of one client (visualizations on local models of other clients are similar and are provided in Appendix E). Following the same procedure as in Sec. 3.1, we plot the singular values of covariance matrices of representations produced by the local models. We observe from Fig. 2(b) that locally trained models demonstrate the same trend as the global models—namely, that the presence of stronger data heterogeneity causes more severe dimensional collapse.
These experiments corroborate that the global model inherit the adverse dimensional collapse phenomenon from the local models.
Based on the empirical observations in Sec. 3.1 and Sec. 3.2, we now develop a theoretical understanding to explain why heterogeneous training data causes dimensional collapse for the learned representations. Since we have established that the dimensional collapse of global model stems from local models, we focus on studying local models in this section.
Without loss of generality, we study local training of one arbitrary client.
Specifically, we first analyze the gradient flow dynamics of the model weights during the local training.
This analysis shows how heterogeneous local training data drives the model weights towards being low-rank, which leads to dimensional collapse for the representations.
3.3. S ETUPS AND N OTATIONS
We denote the number of training samples as N , the dimension of input data as d in , and total number of classes as C .
The i -th sample is denoted as X i ∈ R d in , and its corresponding one-hot encoded label is y i ∈ R C .
The collection of all N training samples is denoted as X =
[ X 1 , X 2 . . .
, X N ] ∈ R d in × N , and the N one-hot encoded training labels are denoted as y =
[ y 1 , y 2 , . . .
, y N ] ∈ R C × N .
For simplicity in exposition, we follow Arora et al. (2018; 2019) and Jing et al. and analyze linear neural networks (without nonlinear activation layers).
We consider an ( L + 1) -layer (where
L ≥ 1 ) linear neural network trained using the cross entropy loss under gradient flow (i.e., gradient descent with an infinitesimally small learning rate).
The weight matrix of the i -th layer ( i ∈ [ L +1] ) at the optimization time step t is denoted as W i ( t ) . The dynamics can be expressed as
[Equation]
where ℓ denotes the cross-entropy loss. In addition, at the optimization time step t and given the input data X i , we denote z
i ( t ) ∈ R d as the output representation vector ( d being the dimension of the representations) and γ i ( t ) ∈ R C as the output softmax probability vector.
We have
[Equation]
We define µ c = N c N , where N c is number of data samples belonging to class c .
We denote e c as the C -dimensional one-hot vector where only the c -th entry is 1 (and the others are 0 ).
In addition, let ¯ γ c ( t ) = 1 N c (cid:80) Ni =
1 γ i ( t ) 1 { y i =
e c } and ¯ X c = 1 N c (cid:80) Ni =
X
i 1 { y i =
e c } .
3.3. A NALYSIS ON G RADIENT F LOW D YNAMICS
Since our goal is to analyze model representations z i ( t ) , we focus on weight matrices that directly produce representations (i.e., the first L layers).
We denote the product of the weight matrices of the first L layers as Π
[Equation]
W 1 ( t ) and analyze the behavior of Π( t ) under the gradient flow dynamics.
In particular, we derive the following result for the singular values of Π( t ) .
Theorem 1 (Informal) .
Assuming that the mild conditions as stated in Appendix A.3 hold.
Let σ k ( t ) for k ∈ [ d ] be the k -th largest singular value of Π( t ) .
Then,
[Equation]
where u L +1 ,k ( t ) is the k -th left singular vector of W L +1 ( t ) , v k ( t ) is the k -th right singular vector of Π( t ) , M is a constant, and G ( t ) is defined as
[Equation]
where µ c , e c , ¯ γ c ( t ) , ¯ X c are defined after Eqn.
The proof of the precise version of Theorem 1 is provided in Appendix A. Based on Theorem 1, we are able to explain why greater data heterogeneity causes Π( t ) to be biased to become lower-rank.
Note that strong data heterogeneity causes local training data of one client being highly imbalanced in terms of the number of data samples per class (recall Fig. 1(a)).
This implies that µ c , which is the proportion of the class c data, will be close to 0 for some classes.
Next, based on the definition of G ( t ) in Eqn.
(4), more µ c ’s being close to 0 leads to G ( t ) being biased towards a low-rank matrix.
If this is so, the term ( u L +1 ,k ( t ))
⊤ G ( t ) v k ( t ) in Eqn.
(3) will only be significant (large in magnitude) for fewer values of k .
This is because u L +1 ,k ( t ) and v k ( t ) are both singular vectors, which are orthogonal among different k ’s. This further leads to ˙ σ k ( t ) on the left-hand side of Eqn.
(3), which is the evolving rate of σ k , being small for most of the k ’s throughout training.
These observations imply that only relatively few singular values of Π( t ) will increase significantly after training.
Furthermore, Π( t ) being biased towards being low-rank will directly lead to dimensional collapse for the representations.
To see this, we simply write the covariance matrix of the representations in terms of Π( t ) as
[Equation]
From Eqn.
(5), we observe that if Π( t ) evolves to being a lower-rank matrix, Σ( t ) will also tend to be lower-rank, which corresponds to the stronger dimensional collapse observed in Fig. 2(b).
l og
s i
ngu l a r v a l ue s
M ITIGATING D IMENSIONAL C OLLAPSE WITH F ED D ECORR
Motivated by the above observations and analyses on dimensional collapse caused by data heterogeneity in federated learning, we explore how to mitigate excessive dimensional collapse. Since dimensional collapse on the global model is inherited from local models, we propose to alleviate the problem during local training.
One natural way to achieve this is to add the following regularization term on the representations during training
[Equation]
where λ i is the i -th singular value of the covariance matrix of the representations.
Essentially,
L singular penalizes the variance among the singular values, thus discouraging the tail singular values from collapsing to 0 , mitigating dimensional collapse.
However, this regularization term is not practical as it requires calculating all the singular values, which is computationally expensive.
Therefore, to derive a computationally-cheap training objective, we first apply the z-score normalization on all the representation vectors
z i as follows:
ˆ z i =
( z i − ¯ z ) / (cid:112) Var( z ) .
This results in the covariance matrix of ˆ z i being equal to its correlation matrix (i.e., the matrix of correlation coefficients).
The following proposition suggests a more convenient cost function to regularize.
Proposition 1.
For a d -by- d correlation matrix K with singular values ( λ 1 , . . .
, λ d ) , we have:
[Equation]
The proof of Proposition 1 can be found in Appendix B.
This proposition suggests that regularizing the Frobenius norm of the correlation matrix ∥ K ∥ F achieves the same effect as minimizing L singular .
In contrast to the singular values, ∥ K ∥ F can be computed efficiently.
To leverage this proposition, we propose a novel method, F ED D ECORR , which regularizes the Frobenius norm of the correlation matrix of the representation vectors during local training on each client.
Formally, the proposed regularization term is defined as:
[Equation]
where w is the model parameters, K is the correlation matrix of the representations.
The overall objective of each local client is
[Equation]
where ℓ is the cross entropy loss, and β is the regularization coefficient of F ED D ECORR .
The pseudocode of our method is provided in Appendix G. To visualize the effectiveness of L FedDecorr in mitigating dimensional collapse, we now revisit the experiments of Fig. 2 and apply L FedDecorr under the heterogeneous setting where α ∈ { 0 . 01 , 0 . 05 } .
We plot our results in Fig.
3 for both local and global models. Figs.
3(a-b) show that for local models,
F ED D ECORR encourages the tail singular values to not collapse to 0 , thus effectively mitigating dimensional collapse.
Moreover, as illustrated in Figs. 3(c-d), this desirable effect introduced by
F ED D ECORR on local models can also be inherited by the global models.
[Table]
E XPERIMENTS 5. E XPERIMENTAL S ETUPS
[Table]
Datasets:
We adopt three popular benchmark datasets, namely CIFAR10, CIFAR100, and TinyImageNet.
CIFAR10 and CIFAR100 both have 50 , 000 training samples and 10 ,test samples, and the size of each image is 32 × 32 .
TinyImageNet containsclasses, with 100 , 000 training samples and 10 , 000 testing samples, and each image is 64 × 64 .
The method generating local data for each client was introduced in Sec. 3.1.
Implementation Details: Our code is based on the code of Li et al.
(2021b). For all experiments, we use MobileNetV2 (Sandler et al., 2018).
We run 100 communication rounds for all experiments on the CIFAR10/ datasets and 50 communication rounds on the TinyImageNet dataset.
We conduct local training for 10 epochs in each communication round using SGD optimizer with a learning rate of 0 .
01 , a SGD momentum of 0 . 9 , and a batch size of 64 .
The weight decay is set to 10 − 5 for CIFAR10 and 10 − 4 for CIFAR100 and TinyImageNet.
We apply the data augmentation of Cubuk et al. (2018) in all CIFAR100 and TinyImageNet experiments.
The β of F ED D ECORR (i.e., β in Eqn.
(9)) is tuned to be 0 .
The details of tuning hyper-parameters for other federated learning methods are described in Appendix F.
5. F ED D ECORR S IGNIFICANTLY I MPROVES B ASELINE M ETHODS
To validate the effectiveness of our method, we apply F ED D ECORR to four baselines, namely FedAvg (McMahan et al., 2017), FedAvgM (Hsu et al., 2019), FedProx (Li et al., 2020), and MOON (Li et al., 2021b).
We partition the three benchmark datasets (CIFAR10, CIFAR100, and TinyImageNet) into 10 clients with α ∈ { 0 . 05 , 0 . 1 , 0 . 5 , ∞} .
Since α = ∞ is the homogeneous setting where local models should be free from the pitfall of excessive dimensional collapse, we only expect F ED D ECORR to perform on par with the baselines in this setting.
We display the CIFAR10/100 results in Tab. 1 and the TinyImageNet results in Tab. 2.
We observe that for all of the heterogeneous settings on all datasets, the highest accuracies are achieved [Figure] by adding F ED D ECORR on top of a certain baseline method.
In particular, in the strongly heterogeneous settings where α ∈ { 0 . 05 , 0 . 1 } , adding F ED D ECORR yields significant improvements of around 2 % ∼ 9 % over baseline methods on all datasets.
On the other hand, for the less heterogeneous setting of α = 0 . 5 , the problem of dimensional collapse is less pronounced as discussed in Sec 3, leading to smaller improvements from F ED D ECORR .
Such decrease in improvements is a general trend and is also observed on FedProx, FedAvgM, and MOON.
In addition, surprisingly, in the homogeneous setting of α = ∞ , F ED D ECORR still produces around 2% of improvements on the TinyImageNet dataset.
We conjecture that this is because TinyImageNet is much more complicated than the CIFAR datasets, and some other factors besides heterogeneity of label may cause undesirable dimensional collapse in the federated learning setup. Therefore, federated learning on
TinyImageNet {+can benefit from F ED D ECORR even in the homogeneous setting. # clients Method α = 0 . 
 FedAvg 35.+}
+ F ED D ECORR 40. FedAvg 31. + F ED D ECORR 39. FedAvg 26.
+ F ED D ECORR 36.
FedAvg 25.
+ F ED D ECORR 34.
FedAvg 21.
+ F ED D ECORR 30.
To further demonstrate the advantages of F EDD ECORR , we apply it on FedAvg and plot how the test accuracy of the global model evolves throughout the federated learning in Fig.
In this figure, if we set a certain value of the testing accuracy as a threshold, we see that adding
F ED D ECORR significantly reduces the number of communication rounds needed to achieve the given threshold.
This further shows that F EDD ECORR not only improves the final performance, but also greatly boosts the communication efficiency in federated 5. A BLATION S TUDY ON THE N UMBER OF C LIENTS Next, we study whether the improvements brought by F ED D ECORR are preserved as number of clients increases. We partition the TinyImageNet dataset into 10 , 20 , 30 , 50 , andclients according to different α ’s, and then run
FedAvg with and without F ED D ECORR .
For the experiments with 10 , 20 and 30 clients, we run 50 communication rounds.
For the experiments with 50 and 100 clients, we randomly select 20% of the total clients to participate the federated learning in each round and run 100 communication rounds.
Results are shown in Tab. 3. From this table, we see that the performance improvements resulting from F ED D ECORR increase from around 3 % ∼ 5 % to around 7 % ∼ 10 % with the growth in the number of clients. Therefore, interestingly, we show through experiments that the improvements brought by F ED D ECORR can be even more pronounced under the more challenging settings with more clients. Moreover, our experimental results under random client participation show that the improvements from F ED D ECORR are robust to such uncertainties.
These experiments demonstrate the potential of F ED D ECORR to be applied to real world federated learning settings with massive numbers of clients and random client participation.
[Figure]
5. A BLATION S TUDY ON THE R EGULARIZATION C OEFFICIENT β
Next, we study F ED D ECORR ’s robustness to the β in Eqn.
(9) by varying it in the set { 0 . 01 , 0 . 05 , 0 . 1 , 0 . 2 , 0 . 3 } .
We partition the CIFAR10 and TinyImageNet datasets into 10 clients with α equals to 0 .
05 and 0 .
1 to simulate the heterogeneous setting.
Results are shown in Fig.
We observe that, in general, when β increases, the performance of F ED D ECORR first increases, then plateaus, and finally decreases slightly.
These results show that F ED D ECORR is relatively insensitive to the choice of β , which implies F ED D ECORR is an easy-to-tune federated learning method.
In addition, among all experimental setups, setting β to be 0 .
1 consistently produces (almost) the best results.
Therefore, we recommend β = 0 .
1 when having no prior information about the dataset.
5. A BLATION S TUDY ON THE N UMBER OF L OCAL E POCHS
E Method
CIFAR
{+TinyImageNet α = 0 . 
 FedAvg 50. + F ED D ECORR 53. FedAvg 59.+}
+ F ED D ECORR 61.
FedAvg 59.
+ F ED D ECORR 61.
FedAvg 58.
+ F ED D ECORR 60.
Lastly, we ablate on the number of local epochs per communication round.
We set the number of local epochs E to be in the set { 1 , 5 , 10 , 20 } .
We run experiments with and without F ED D ECORR , and we use the
CIFAR100 and TinyImageNet datasets with α being 0 .
05 and 0 .
1 for this ablation study.
Results are shown in Tab. 4, in which one observes that with increasing E , F ED A VG performance first increases and then decreases.
This is because when E is too small, the local training cannot converge properly in each communication round.
On the other hand, when E is too large, the model parameters of local clients might be driven to be too far from the global optimum.
Nevertheless,
F ED D ECORR consistently improves over the baselines across different choices of local epochs E .
5. A DDITIONAL E MPIRICAL A NALYSES
We present more empirical analyses in Appendix C. These include comparing F ED D ECORR with other baselines (Appendix C.4) and other decorrelation methods (Appendix C.2), experiments on other model architectures (Appendix C.3) and another type of data heterogeneity (Appendix C.5), and discussing the computational advantage of F ED D ECORR (Appendix C.1).
C ONCLUSION
In this work, we study representations of trained models under federated learning in which the data held by clients are heterogeneous.
Through extensive empirical observations and theoretical analyses, we show that stronger data heterogeneity results in more severe dimensional collapse for both global and local representations.
Motivated by this, we propose F ED D ECORR , a novel method to mitigate dimensional collapse during local training, thus improving federated learning under the heterogeneous data setting.
Extensive experiments on benchmark datasets show that F ED D ECORR yields consistent improvements over existing baseline methods.
A CKNOWLEDGEMENTS
The authors would like to thank anonymous reviewers for the constructive feedback.
Yujun Shi and Vincent Tan are supported by Singapore Ministry of Education Tier 1 grants (Grant Number:
A-0009042-01-00, A-8000189-01-00, A-8000980-00-00) and a Singapore National Research Foundation (NRF) Fellowship (Grant Number: A-0005077-01-00).
Jian Liang is supported by National
Natural Science Foundation of China (Grant No. 62276256) and Beijing Nova Program under Grant
Z211100002121108.
R EPRODUCIBILITY S TATEMENT
All source code has been released at https://github.com/Yujun-Shi/FedCLS.
Pseudo-code of F EDD ECORR is provided in Appendix G. We introduced all the implementation details of baselines and our method in Sec. 5.1.
In addition, the proofs of Theorem 1 and Proposition 1 are provided in
Appendix A and Appendix B, respectively.
All assumptions are stated and discussed in the proof.
A P ROOF OF T HEOREM 1 IN M AIN P APER A. N OTATIONS R EVISITED
Here, for the reader’s convenience, we summarize the notations used in both the main text and this appendix.
[Table]
N
[Table]
, X
[Table]
R d in × N .
[Table]
one
[Table]
the N
[Table]
, y ∈ R C × N .
[Table]
all N
[Table]
γ
[Table]
C
[Table]
N
[Table] t [Table]
step.
[Table] W [Table] [Table]
W l .
[Table] W [Table] [Table]
of W l .
[Table]
Π .
[Table]
[Table]
[Table]
C -dimensional one-hot vector where only the c -th entry is 1 .
[Table] A. T [Table] LEMMAS
Here, we elaborate two useful lemmas from Arora et al.
The first lemma is adopted from Arora et al.
Lemma 1.
Assuming the weight matrix W evolves under gradient descent dynamics with infinitesimally small learning rate, the k -th singular value of this matrix (denoted as σ k ) evolves as
[Equation]
[Equation]
where u k ( t ) and v k ( t ) are the k -th left and right singular vectors of W ( t ) , respectively. Proof. By performing an SVD on W ( t ) , we have W ( t ) = U ( t ) S ( t ) V ( t ) ⊤ .
Therefore, by the chain rule in differention, we have:
[Equation]
Next, for both sides of the above equation, we left multiply U ( t ) ⊤ and right multiply V ( t ) :
[Equation]
Since S ( t ) is a diagonal matrix, we consider the k -th diagonal entry of S ( t ) , namely σ k ( t ) :
[Equation]
Since u k ( t ) and v k ( t ) are unit vectors and are evolving in time with infinitesimal rate, we have ( u k ( t ))
⊤ ˙
u k ( t ) = 0 and ( v k ( t ))
⊤ ˙ v k ( t ) = 0 .
Next, Eqn.
(13) can be simplified as
[Equation]
[Equation]
The proof is thus complete.
The second lemma is adopted from Arora et
[Table] Lemma 2.
Given L consecutive linear layers in a neural network characterized by weight matrices
W 1 , W 2 , . .
, W L .
We denote Π =
W L W L − 1 . . .
W 1 .
We further denote W j ( t ) as weight matrix W j after the t -th gradient descent optimization step.
Correspondingly, the initialization of W j is W j (0) .
Assuming we have W j (0)( W j (0)) ⊤ = ( W j +1 (0)) ⊤ W j +1 (0) for any j ∈
[ L − 1] at initialization.
Then, under the gradient descent dynamics, Π( t ) satisfies
[Equation]
where [ · ] L − jL and [ · ] j − 1 L are fractional power operators defined over positive semi-definite matrices.
Proof.
Here, we first define some additional notation.
Given any square matrices (or possibly scalar)
A 1 , A 2 , . . .
, A m , we denote diag(
A 1 , A 2 , . . . ,
A m ) to be the block diagonal matrix
[Equation]
[Equation]
Here, we first consider dynamics of an arbitrary W j where j ∈ [ L − 1] .
By the chain rule, we have
[Equation]
[Equation]
Given Eqn. (16), we right multiply ˙ W j ( t ) by ( W j ( t )) ⊤ and we left multiply ˙ W j +1 ( t ) by ( W j +1 ( t ))
⊤ , which yields
[Equation]
Applying the same trick on W j ( t ) ⊤ and W j +1 ( t ) ⊤ yields
[Equation]
Adding Eqns. (17) and (18) on both sides yields [Equation] Next, by the chain rule for differentiation, Eqn.
(19) directly implies that
[Equation]
Since we have assumed that W j (0) W j (0) ⊤
=
W j +1 (0)
⊤ W j +1 (0) , we can conclude that
[Equation]
Next, we apply an SVD on W j ( t ) and W j +1 ( t ) in Eqn.
This yields
[Equation]
[Equation]
Based on Eqn.
(22) and given the uniqueness property of SVD, we know:
[Equation]
where √ ρ 1 , . . .
√
ρ m represent the m distinct singular values satisfying ρ 1 > ρ 2 > . . .
>
ρ m ≥ 0 , and I d r for any r ∈ [ m ] are identity matrix of size d r × d r .
Since Eqn.
(23) holds for any j , we know by induction that the set of values of ρ ’s is the same across all layers
j ∈ [ L ] . In addition, there exist orthogonal matrices O j,r ∈ R d r × d r for any r ∈ [ m ] such that [Equation]
[Equation]
Given Eqns.
(24), next, we study W j +1 ( t ) W j ( t
[Equation]
W ⊤ j
[Equation]
t
[Equation]
⊤ j +
[Equation]
t ) for any j ∈ [ N − 1] :
[Equation]
Similarly, it holds that
[Equation]
Next, by induction and Eqns.
[Equation]
by induction and Eqns.
(26), it holds that
[Equation]
From Eqns.
[Equation]
we know that for any j ∈ [ L − 1] ,
[Equation]
Similarly, from Eqn.
[Equation]
we know that for any 2 ≤ j ≤ L − 1 ,
[Equation]
With everything derived above
[Equation]
we now study the dynamics of Π( t ) as follows ˙Π
[Equation]
This completes the proof
[Equation]
A. A SSUMPTIONS
Assumption 1.
We assume that the initial values of the weight matrices satisfy W ⊤ i +1 (0) W i +1 (0) =
W i (0) W ⊤ i (0) for any i ∈
[ L − 1] .
Assumption 2.
We assume | u k ( t ) ⊤ v L +1 ,k ′ ( t )
| = 1 { k = k ′ } holds for all t , where u k ( t ) is the k -th left singular vector of Π( t ) and v L +1 ,k ′ ( t ) is the k ′ -th
right singular vector of W L +1 ( t ) .
Remark:
For Assumption 1, it can be achieved in practice by proper random initialization.
For
Assumption 2, Ji & Telgarsky (2018) proved that under some assumptions, gradient descent optimization will drive consecutive layers of linear networks to satisfy it. We also provide empirical evidence in Fig. 6 to corroborate that this assumption approximately holds. A. P ROOF OF T HEOREM 1 IN THE MAIN TEXT
Theorem 1 (formally stated).
Let σ k ( t ) for k ∈ [ d ] be the k -th largest singular value of Π( t ) . Then, under Assumptions 1 and 2, we have [Equation] 32) where u L +1 ,k ( t ) is the k -th left singular vector of W L +1 ( t ) , v k ( t ) is the k -th right singular vector of Π( t ) , M is a constant, and G ( t ) is defined as
[Equation]
Proof.
Recall that for ( L + 1) -layer linear neural networks, given the i -th training sample X i ∈ R d , we have [Equation] [Equation] and the loss is the standard cross-entropy loss defined as follows
[Equation]
[Equation]
By the chain rule, we can derive gradient of ℓ with respect to W L +1 and Π , which are respectively,
[Equation]
36) and
[Equation]
Next, under the gradient descent dynamics, the dynamics on W L +1 satisfies
[Equation]
while the dynamics on Π requires invoking Lemma
[Equation]
which allows us to write
[Equation]
Next
[Equation]
we invoke Lemma 1 on Eqn.
(39) and Eqn.
[Equation]
respectively
[Equation]
yielding:
[Equation]
and
[Equation]
Combining Eqns.
(40) and (41)
[Equation]
we have:
[Equation]
Next, apply integration on both sides
[Equation]
which yields
[Equation] where M a constant. By Eqn. (43), Eqn.
(40) can be rewritten as
[Equation]
[Equation]
Finally
[Equation]
notice that ( y − γ ( t ))
X ⊤ can be rewritten as
[Equation]
We further substitute Eqn.
(45) into Eqn.
and obtain ˙
[Equation]
where G ( t ) is defined as
[Equation]
This completes the proof.
B P ROOF OF P ROPOSITION 1 IN THE M AIN P APER
Proposition 1 (restated). For a d -by- d correlation matrix K with singular values ( λ 1 , . . . , λ d ) , we have: [Equation]
Proof.
Given a d -by- d correlation matrix K , since the diagonal entries of K are all 1 , we have [Equation] [Equation] This is because for any symmetric positive definite matrix, the sum of all singular values equals the trace of the matrix.
Next, for the left-hand side of Eqn.
(7), we have:
[Equation]
Next, for the right-hand side of Eqn.
(7), we have:
[Equation]
[Equation] SVD on K [Equation] Therefore, we have shown that the left-hand side of Eqn.
(7) equals its right-hand side. C A DDITIONAL E MPIRICAL A NALYSES C. C OMPUTATIONAL E FFICIENCY We demonstrate F ED D ECORR ’s advantage vis-`a-vis some of its competitors in terms of its computational efficiency.
We compare F ED D ECORR with some other methods that also apply additional regularization terms during local training such as FedProx and MOON.
We partition CIFAR10, CIFAR100 and TinyImageNet into 10 clients with α = 0 . 5 and report the total computation times required for one round of training for FedAvg, FedProx, MOON, and F ED D ECORR . Results are shown in Tab. 5.
All results are produced with a NVIDIA Tesla V100 GPU. We see that F EDD ECORR incurs a negligible computation overhead on top of the na¨ıve FedAvg, while FedProx and MOON introduce about 0 .
5 ∼ 1 × additional computation cost. The advantage of F ED D ECORR in terms of efficiency is mainly because it only involves calculating the Frobenius norm of a matrix which is extremely cheap. Indeed this regularization operates on the output representation vectors of the model, without requiring computing parameter-wise regularization like FedProx nor extra forward passes like MOON.
[Table]
C. C OMPARISON WITH O THER D ECORRELATION M ETHODS
Some decorrelation regularizations such as DeCov (Cogswell et al., 2015) and StructuredDeCov (Xiong et al., 2016) were proposed to improve the generalization capabilities in standard classification tasks.
Both these methods operate directly on the covariance matrix of the representations instead of the correlation matrix like our proposed method—F ED D ECORR .
To compare our
F ED D ECORR with the existing decorrelation methods, we follow the same procedure as in F EDD ECORR and apply DeCov and Structured-DeCov during local training.
Our experiments are based on TinyImageNet and FedAvg.
TinyImageNet is partitioned into 10 clients according to various α ’s.
Results are shown in Tab. 6.
Surprisingly, we see that unlike our F ED D ECORR which steadily improves the baseline, adding DeCov or Structured-DeCov both degrade the performance in federated learning.
We conjecture that this is because directly regularizing the covariance matrix is highly unstable, leading to undesired modification on the representations.
This experiment shows that our design of regularization of the correlation matrix instead of the covariance matrix is of paramount importance. [Table] DeCov St [Table]
Decov [Table] α = 0 . [Table]
[Table]
[Table]
α = 0 .39.
α = 0 .
[Table]
In this section, we demonstrate the effectiveness of our method across different model architectures. Here, besides the MobileNetV2 used in the main paper, we also experiment on ResNet and ResNet32. Note that ResNet18 is the wider ResNet whose representation dimension is 512 and
ResNet32 is the narrower ResNet whose representation dimension is 64 .
The coefficient of the FedDecorr objective is set to be 0 .
1 as suggested to be a good universal value of β in the paper.
The heterogeneity parameter α is set to be 0 . and we use the CIFAR10 dataset. Our results are shown in Tab. 7.
As can be seen, FedDecorr yields consistent improvements across different neural network architectures.
One interesting phenomenon is that the improvements brought about by FedDecorr are much larger on wider networks (e.g., MobileNetV2, ResNet18) than on narrower ones (e.g.
ResNet32). We conjecture this is because the dimension of the ambient space of wider networks are clearly higher than that of shallower networks. Therefore, relatively speaking, the dimensional collapse caused by data heterogeneity will be more severe for wider networks. [Table]
[Table] [Table] C. C OMPARISON WITH O THER F EDERATED L EARNING B ASELINES In this section, we compare F ED D ECORR with two other baselines, namely Scaffold (Karimireddy et al., 2020) and FedNova (Wang et al., 2020b).
We use the same experimental setups as in the main paper to implement these two baselines.
Results on CIFAR10/100 and TinyImageNet are shown in Tab. 8 and Tab. 9, respectively.
As shown in the tables, across various datasets and degrees of heterogeneity, adding F ED D ECORR on top of a baseline method can outperform the baselines when there is some heterogeneity across the agents, i.e., α < ∞ .
C. E XPERIMENTS ON A NOTHER T YPE OF H ETEROGENEITY
In this section, we run experiments under another type of data heterogeneity.
Specifically, we follow McMahan et al. and split the CIFAR10 dataset across different clients such that each client only has a fixed number of classes C (e.g., C = 2 indicates each client only has data of two classes).
We split the data across 10 clients and choose C to be 2 and 3 . Results are shown in Tab. 10.
As can FedAvgM FedProx MOON l og s i ngu l a r v a l ue s
ResNet
ResNet l og s i ngu l a r v a l ue s be observed, under this different scenario, ED D ECORR also yields noticeable and improvements.
[Equation]
In this section, we provide additional visualizations on global models with different federated learning methods, model architectures, and datasets.
Through our extensive experimental results, we demonstrate that dimensional collapse is a general problem under heterogeneous data in federated learning.
D. V ISUALIZATION ON G LOBAL M ODELS OF O THER F EDERATED L EARNING M ETHODS
In the main text, we have shown that global models produced by FedAvg (McMahan et al., 2017) suffer stronger dimensional collapse with increasing data heterogeneity.
To further show such dimensional collapse phenomenon is a general problem in federated learning, we visualized global models produced by other federated learning methods such as FedAvg with server momentum (Hsu et al., 2019),
FedProx (Li et al., 2020), and MOON (Li et al., 2021b).
Specifically, we follow the same procedure as in the main text and plot the singular values of covariance matrices of representations.
Results are shown in Fig. From the figure, one can see that all these three other methods also demonstrated the similar hazard of dimensional collapse as in FedAvg.
In the main text, we have shown the dimensional collapse on global models caused by data heterogeneity with MobileNetV2.
In this section, we perform the similar visualization based on other
CIFAR
TinyImageNet l og
s i
ngu l a r v a l ue s model architectures such as ResNet32 and ResNet18.
Note that ResNet32 is a narrower ResNet whose representation dimension is 64 and ResNet18 is a wider ResNet whose representation dimension is 64 . We visualize the top 50 singular values for ResNet32 and the top 100 singular values for
ResNet18.
Results are shown in Fig. From the figure, one can observe that heterogeneous data also lead to dimensional collapse on ResNet32 and ResNet18.
In the main text, we use the CIFAR100 dataset for our visualizations.
In this section, we perform similar visualizations with other datasets such as CIFAR10 and TinyImageNet. Results are shown in Fig. From the figure, one can also observe that dimensional collapse results from data heterogeneity.
E V ISUALIZATION ON O THER L OCAL C LIENTS
In the main text Fig. 2(b), under the four different degrees of data heterogeneity (i.e., α ∈ { 0 . 01 , 0 . 05 , 0 . 25 , ∞} ), we compare representations of local models of client 1 and empirically show how data heterogeneity affects representations produced by the local models.
In this section, to further corroborate our conclusion, we follow the same procedure and visualize singular values of the covariance matrix of representations produced by local models trained on the rest of the 9 clients under the same α ’s. Results are shown in Fig.
From the results, we can obtain the similar observations as in Fig. 2(b) of the main text, namely that stronger data heterogeneity causes more severe dimensional collapse for local models. F H YPERPARAMETERS OF O THER F EDERATED L EARNING M ETHODS
The regularization coefficient of
FedProx (Li et al., 2020) µ is tuned across { 10 − 4 , 10 − 3 , 10 − 2 , 10 − 1 } and is selected to be µ = 10 − 3 ; the regularization coefficient of
MOON (Li et al., 2021b) µ is tuned across { 0 . 1 , 1 . 0 , 5 . 0 , 10 . 0 } and is selected to be µ = 1 . 0 ; the server momentum of FedAvgM (Hsu et al., 2019) ρ is tuned across { 0 . 1 , 0 . 5 , 0 . 9 } and is selected to be ρ = 0 .
G P SEUDO - CODE OF F ED D ECORR
Here, we provide a pytorch-style pseudo-code for F ED D ECORR in Alg. 1.
All F ED D ECORR -specific components are highlight in blue.
As indicated in the pseudocode, the only additional operation of
F ED D ECORR is in adding a regularization term L FedDecorr ( w, X ) defined in Eqn.
This shows that F ED D ECORR is an extremely convenient plug-and-play federated learning method.
In this section, we first split CIFAR10 into 10 clients with α = 0 .
Then, we plot how FedDecorr loss evolve within 10 local epochs for all the 10 clients in Fig. All training configurations are the same as in the main paper.
From the results, one can observe that the optimization process of
FedDecorr loss is stable.
[Figure]
l og s i n gu l a r v a l ue s
[Figure] [Table]