Towards Understanding and Mitigating Dimensional Collapse in Heterogeneous Federated Learning
Authors: Yujun Shi
Jian Liang
Wenqing Zhang
Vincent Tan
Song Bai
Venue: ICLR 2023 poster (2023)
Keywords: federated Learning
representation Learning
data heterogeneity
dimensional collapse

Abstract: 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 FedDecorr, a novel method that can effectively mitigate dimensional collapse in federated learning. Specifically, FedDecorr applies a regularization term during local training that encourages different dimensions of representations to be uncorrelated. FedDecorr, which is implementation-friendly and computationally-efficient, yields consistent improvements over baselines on standard benchmark datasets. Code: https://github.com/bytedance/FedDecorr.

Published as a conference paper at ICLR 2023
TOWARDS UNDERSTANDING AND MITIGATING DIMENSIONAL COLLAPSE IN HETEROGENEOUS FEDERATED LEARNING
Yujun Shi! Jian Liang'
Wenqing Zhang
Vincent F. Tan' Song Bai²
National University of Singapore shi. yujunuu. edu
BByteDance Inc. Institute of Automation, CAS
vtan@nus www.iinussduu.ceo
..m.t...mm...cmmmmmmmm com
ABSTRACT
Federated learning aims to train models collaboratively across different clients without sharing data for privacy considerations. However, 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 lower-dimensional space instead of the ambient space. Moreover, we observe similar phenomenon on models locally trained on each client and deduce that the dimensional collapse on the global model is in- herited 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 FEDDECORR, a novel method that can effectively mitigate dimensional collapse in federated learning. Specifically, FEDDECORR applies a regulariza- tion term during local training that encourages different dimensions of represen- tations to be uncorrelated. FEDDECORR, which implementation-frieen compuuationally--ffiiienn, yields consistent improvements over baselines on stan- dard benchmark datasets. Code: http://gghhbbooooooyddde/eeeeeee
and
INTRODUCTION
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, McMMhan et al. (2017) propose Federated Learning decentralized train- ing 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. (a) for an example). Such dis- crepancies can result in drastic disagreements between the local optima of the clients and the desired global optimum, which may lead severe performance degradation of the global model. Previous works attempting to tackle this challenge mainly focus on the model parameters, either during lo- cal training (Li et al., (2020) Karimireddy et al.) (020)) or global aggregation (Wang al., 20206). However, these methods usually result in an excessive computation burden high communication costs (Li et al.) (021aa) because deep neural networks are typically heavily over prerpparameterized. In contrast, in this work, we focus the representation space of the model and study the impact of data heterogeneity.
To commence, 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 dif- ferent degrees of data heterogeneity. Since the singular values of the covariance matrix provide comprehensive characterization of the distribution of high-dimensional embeddings, we use it to
Work done when interning with Song Bai.

---

Published as a conference paper at ICLR 2023
Client
Ciint 2
Dela
dog bird
bird
(A) Hotoragenalty
Train Homo.
Train Hetero.
Herero, FedDlccor
Figure 1: (a) illustrates data heterogeneity in terms of number of samples per class. (b). (c). (d) show representations (normalized to the unit sphere) of global models trained under homogencous data, heterogeneous data, and heterogeneous data with FEDDECORR, respectively. Only (c) suffers
dimensional collapse. (b), (c). (d) are produced with ResNet20 on CIFARIO. Best viewed color.
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 oower-dimensional space (or manifold). graphical illustration of how heterogeneous training data affect output repre- sentations is shown in Fig. [[bcc. 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 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 establish the connection between dimensional
produced by local models. With this observation, collapse of the global model and local models. To further understand the dimensional collapse on local models, analyze the gradient flow dynamics of local training in Sec. 3.3 Interestingly,
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, consider mitigating dimensional collapse during local training Sec. In particular, we propose a novel federated learning method termed FEDDECORR. FEDDECORR adds a regularization term during local training encourage the Frobenius norm of the correlation matrix of representations to be small. We show theoretically and empirically that this proposed regularization term can effec- tively mitigate dimensional collapse (see Fig. (d) for example). Next, in Sec. 5)) through extensive experiments on standard benchmark datasets including CIFAR10. CIFAR100, and TinyImageNe we show that FEDDECORR consistently improves over baseline federated learning methods. In ad- dition, we find that FEDDECORR yields more dramatic improvements in more challenging federated learning setups such as stronger heterogeneity or more number of clients. Lastly, FEDDECORR 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, discover through experiments that stronger data heterogeneity in federated learning leads to greater dimensional collapse for global and local models. Second, develop 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 novel method called FEDDECORR, which yields consistent improvements while being pnpllmmentation-frienllly and computationalll--ffiiiiin
RELATED WORKS
Federated Learning. McMahan et al. (007)) proposed FedAvg. which adopts a simple averaging scheme to aggregate local models into the global model. However, under data heterogeneity, FedAvg
2

---

Published as a conference paper at ICLR 2023
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-Shedivaa et al. 2020;; Wang al. 2021 global aggregation (Wang
et al.
2020b
2019 Luo et al. 2021 Wang et al. 2020a Lin et al. 2020; Reddi et al. et al.
2020
Most of these methods focus the model parameter space, which may result in high computation or communication cost due to deep neural networks being over- parameterized. Li al. (202Ib) focuses on model representations and uses a contrastive loss to maximize agreements between representations of local models and the global model. However, one drawback Li et al. (202Ib) 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
Wang et al. 2020aa.
data. Another research trend is in personalized federated learning (Arivazhagan et al. 2019 [.i et al.,
2021c Fallah et al. 2020 T Dinh et al. 2020 Hanzely et al. |020 Huang et al. 20021] Zhang et al,,
2020). which aims train personalized local models for each chient. In this work, however,
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) al. 2022 In this work, focus on federated learning and discover that stronger data hetero- geneity 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. 20118 2019) introduce the gradient flow dynamics frame- work to analyze the dynamics of multi- multi-layer linear neural networks under the fy-loss and find deeper neural networks biasing towards low-rank solution during optimization. Following their works. Ding 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 feder- ated 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 pre- venting mode collapse in self-supervised learning (Bardes et al.) 2021 Zbontar et al. 2021 Hua
et al. )021), boosting generalization (Cogswell et al. 2015| | Huang et al. 2018 [188g al. (016),
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.
DIMENSIONAL COLLAPSE CAUSED BY DATA HETEROGENEITY
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 local training in Sec. 3.3
3.1 EMPIRICAL OBSERVATIONS ON THE GLOBAL MODEL
We first empirically demonstrate that stronger data heterogeneity causes more severe dimensional collapse the global model. Specifically, we first separate the training samples of CIFAR 100 into 10 splits, each corresponding to the local data of one client. To simulate data heterogeneity among
sam- Dir (a) and allocate a Pe.k proportion of instances of class e E [C] f1, C} to client & E [K]. where Dirg (a) is the Dirichlet dis- tribution with K categories and a is the concentration parameter. A smaller implies stronger data heterogeneity (a = oo corresponds to the homogeneous setting). We let E (0.01, 0.05, 0.25, oo).
clients as previous works (Yurochkin et al. 2019; Wang al. 2020a Li et al. (021b),
ple a probability vector Pe = (Pc.1. Pe.2, Pc.K)
For each of the settings generated by different a's, we apply FedAvg (McMahan et al.] (017) to train a MobileNetV2 (Sandler et al.) 2018) with CIFAR 100 (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 Σ N 2i=1
3

---

Published as a conference paper at ICLR 2023
(a) Global Model
(b) Local Model
(homo)
(homal
log singular values
2-0.05
0-0.01
dl 01
0-
-2
20
40
60
80
100
4 20
k
k
Figure 2: Data beterogeneity causes dimensional collapse on (a) global models and (b) local mod- els. We plot the singular values the covariance matrix of representations in descending order. The r-axii (k) is the index of singular values and the y-axis is the logarithm of the singular values.
2)( - T)T of the representations over the test data points in CIFARIO0. Here z; is the i-th test data point and z N/2 Ei= is their average.
Finally, we apply the singular value decomposition (SVD) on each of the covariance matrices and visualize the top 100 singular values in Fig. a If we define small value as the threshold for singular value be significant (e.g.. log most all the singular values are significant, i.e., they surpass However, as decreases, the number of singular values exceeding 7 monotonically decreases. This implies that with stronger heterogene- ity among local training data, the representation vectors produced by the trained global model tend to reside a lower-dimensional space, corresponding to more severe dimensional collapse.
2), we observe that for the homogeneous setting, al-
3.2 EMPIRICAL OBSERVATIONS ON LOCAL MODELS
Since the global model is obtained by aggregating locally trained models on each client, we conjec- ture 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 a's, we visualize the locally trained model of one client (visualizations on local models of other clients are similar and provided in Appendix ED. Following the same procedure as in Sec. [.I] we plot the singular values of covariance matrices representations pro- duced by the local models. We observe from Fig. (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.
3.3 A THEBRETICAL EXPLANATION FOR DIMENSIONAL COLLAPSE
Based on the empirical observations in Sec. 3.1 and Sec. 3.2, we now develop a theoretical under-
standing to explain why heterogeneous training data causes dimensional collapse for the learned representations.
Since 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. SETUPS AND NOTATIONS
We denote the number of training samples as N, the dimension of input data as din and total number of classes C. The i-th sample is denoted X, E Rd,,, and its corresponding one-hot encoded label is y; E R The collection of all N training samples is denoted as X = [X,, Xx...,Xy e Rdxx and the N one-hot encoded training labels are denoted y ..,y...,,y RCN
4

---

Published as a conference paper at ICLR 2023
For simplicity in exposition, follow Arora et al. (2018 2019) and Jing al. (2021) and analyze linear neural networks (without nonlinear activation layers). We consider an (L + layer (where L 2 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 E [1++] 1) the optimization time step is denoted as W,(t) The dynamies can be expressed
Wi(t) ow; ((W) WLL11())))
(1)
where l denotes the cross-entropy loss.
In addition, at the optimization time step t and given the input data X,, we denote 2,(t) E R as the
output representation vector (d being the dimension of the representations) and 7:(t) E RC as the output softmax probability vector. We have
(t)
(t
softmax(W (i)))).)
We define Me N where No is number of data samples belonging to class We denote e, as the C-dimensional one-hot vector where only the e-th entry is 1 (and the others are 0). In addition, let Ta(1)= N. Σi= Ty; =..) and X N. Ei X,1 (y: = =ee.
No.
3.3.2 ANALYSIS ON GRADIENT FLOW DYNAMICS
Since goal is to analyze model representations z, (f), we focus 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 II(t) (L(t)WL-1(() gradient flow dynamics. In particular, we derive the following result for the singular values of [[(t).
W1 (t) and analyze the behavior of II(t) under the
Theorem (Informal). Assuming that the mild conditions as stated in Appendix A.3 for k E [d] be the k-th largest singular value of (I()). Then,
hold. Let (x((t
(xt()= NL (t)) E (Ox( + M (uL+1A(()) G(t)va(t),
(3)
where (t) is the k-th left singular vector of W WL+1(t), vi(() the k-th right singular vector IIt), M is constant, and G(t) is defined as
G(I) Enole - (o(()--6(t))
where Me e. Pt(,), X are defined after Eqn.
The proof the precise version of Theorem is provided in Appendix A
Based on Theorem I we are able to explain why greater data heterogeneity causes II((t) to be biased 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, [T]). This implies that Per which is the proportion the class e data, will be close to for some classes.
Next, based on the definition of G(t) in Eqn. (4) more Me 's being close to 0 leads to G(I) being biased towards low-rank matrix. If this is so, the term (HL++1, (L)) ((t)vv((t) in Eqn. (3) will only be significant (large in magnitude) for fewer values of k. This is because UL+1& (t) and (t) are both singular vectors, which are orthogonal among different k's. This further leads to o on the left-hand side of Eqn. 3 which is the evolving rate of σke being small for most of the k's throughout training. These observations imply that only relatively few singular values of II(t) will increase significantly after training.
Furthermore, IT(L) being biased towards being low-rank will directly lead to dimensional collapse for the representations. To this, simply write the covariance matrix of the representations in terms of (T( as
1 N (tt = N
(t) = II(t) 1 N
- -x) littT. (5)
E=1
From Eqn. (5). we observe that if II(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).
5

---

Published as a conference paper at ICLR 2023
(+) local model (ox 0.00)
[b) local model (o 0.06)
(6) global model (ar 0.01)
(4) global model 0.05)
singiler voluee
60
50
100
0
20
40
50
30
100
40
00 05
100
20 40 63 80 100
W/o FedDecorr w/ FedDecorf
Figure 3: FEDDECORR effectively mitigates dimensional collapse for (a-b) local models and (c-d) global models. For each heterogeneity parameter a E (0.01, 0.05), we apply FEDDECORR and plot the singular values of the representation covariance matrix. The æ-axis (k) is the index of singular values. With FEDDECORR, the tail singular values are prevented from dropping to too rapidly.
MITIGATING DIMENSIONAL COLLAPSE WITH FEDDECORR
Motivated by the above observations and analyses on dimensional collapse caused by data hetero- geneity in federated learning, we explore how to mitigate excessive dimensional collapse.
Since dimensional collapse on the global model inherited from local models, we propose al- leviate the problem during local training. One natural way to achieve this is to add the following regularization term on the representations during training
Lsingular (w. X)
Σ - j=1
Σ x
(6)
where A is the i-th singular value of the covariance matrix of the representations. Essentially, Lsingular penalizes the variance among the singular values, thus discouraging the tail singular values from collapsing to 0, mitigating dimensional collapse. However, this regularization practical as it requires calculating all the singular values, which is computationally expensive.
is not
Therefore, to derive a computationalll-ccenn training objective, we first apply the z-score normal- ization on all the representation vectors z, as follows: 2; = (z, z)/ /)/v/arr Var(z). This results in the covariance matrix of 2; being equal to its correlation matrix (i.e., the matrix of correlation coeffi- cients). The following proposition suggests a more convenient cost function to regularize.
Proposition 1. For a d-by-d correlation matrix K with singular values (AI,..., (A....,Ad)) Ad), we have:
d
2
Σ
A
A j=1
= |KIl
This proposition suggests that regularizing the Frobenius norm of the correlation matrix || achieves the same effect as minimizing Lsingular- In contrast to the singular values, ||||F can computed efficiently.
The proof Proposition Cean be found in Appendix B
To leverage this proposition, we propose novel method, FEDDECORR, 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:
LeadDecorr (w, X) 1
|
(8)
where w is the model parameters, K is the correlation matrix of the representations. The overall objective of each local client is
min ((w, X,y) + B (w, X),
where the cross entropy loss, and / is the regularization coefficient of FEDDECORR. The pseudocode of our method is provided in Appendix G
To visualize the effectiveness of LFedDecorr in mitigating dimensional collapse, we now revisit the experiments of Fig. and apply LeedDecom under the heterogencous setting where {0.01, 0.05). pe our results in Fig. B both local and global models. Blgs.Blabbi show that for local models, FEDDECORR encourages the tail singular values to not collapse to 0, thus effectively mitigating dimensional collapse. Moreover, as illustrated in Figs. B(--d), this desirable effect introduced by FRDDECORR local models can also be inherited by the global models.

---

Published as a conference paper at ICLR 2023
CIFARIO
CIFART00
Method
a = 0.05
0.1
0.5
8
0.05
0.1
0.5
0C
FedAvg
64.
2..
22.
29.87
16.462216
71.69=015
24.54
FEDDECORR 73.06.0..1 80.60
89.8
19.84.1..0
92.19.1..0
61.53
::.11 67.12 10.09
71.91
10.04
73.87
FedProx
89.57 FEDDECORR 71.38 =0.1 81.74
64.11 -0.4 76.
10
00.4
92.38
60.02+006
66.41
=0:: 71.78
74.34
92. 1:
6 67.00
71.64
74.15=mm
FedAvgM
71.34-0711
77.51 10.58 88.39
10.17 t0.17 91.35
10.15
59.64
66.36
71.17 74.20=a.6
FEDDECORR 73.60 00.2 79.21 +0.15 88. 70 +..26 91.33
+0.13 61.48
4o.7 66.60 71.26
73.86=0.25
MOON
=8.79=0.9 78. 00.60 90.08 10010 92.62
56.79
65.48
71.81 #0.14 74.30
FEDDECORR +..4 81.63 t..5 90.61 t..0 92.63 4-0119 59.43 4:344 66. +220 71.68 +0.05 +3.70
Table I: CIFAR10/100 Experiments. We experiments under various degrees of heterogeneity (a E {0.05, 0.1, 0.5, )o}) and report the test accuracy (%). All results are (re) (roduuddd by us and averaged over 3 runs (mean std). Bold font highlights the highest accuracy in each column.
5 EXPERIMENTS
5.1 EXPERIMENTAL SETUPS
TinylmageNet
Method
Datasets: We adopt three popu- lar benchmark datasets, namely CI- FARIO, CIFARI00, and TinylmageNet. CIFAR1O and CIFAR100 both have 50, 000 training samples and 10, 000 test samples, and the size of each image is 32 x 32. TinyImageNet contains 200 classes, with 100, 000 training samples and 10, 000 testing samples, and each image is 64x64. The method generating local data for each client was introduced in Sec. 3.
a 0.05 0.1
0.5
FedAvg
35.02+t..6 39.30 46.92 49.
+ FEDDECORR 40.29 43.864 50.01 + 52.63 +0.20
39.66404s 47. 16 4 49. 06.36 FedProx 35.20
+ FEDDECORK 40.63 44.19 50.26 52.37 +0.30
FedAvgM 34.81 39. 11 47
1004 49.67
FEDDECORR 39.97 43.95
50. +4. 52.05
MOON 35.23 40.53
50.48
FEDDECORR 40.40 44.20 50.81 10.5:
1 53.01 1:045
Table 2: Tiny Imageett Experiments. We run with a E (0.05, 0.1, 0.5, oo)) and report the test accuracy (%). All results are (re)produced by us and are averaged over runs (mean + std is reported). Bold font highlights the highest accuracy in each column.
Implementation Details: Our code 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 CIFARIO/100 datasets and 50 communication rounds the TinylmageNet 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 5 for CIFAR 10 and 10-4 for CIFAR 100 and Tiny TinllmaeNNet. We apply the data augmentation of Cubuk et al. 2018 in all CIFAR100 and EnylmageNet experiments. The 3 of FEDDECORR (i.e., Equ. (9 is tuned to be 0.1. The details of tuning hyper- paremmerrssssss for other federated learning methods are described
in Appendix F
5.2 FEDDECORR SIGNIFICANTLY IMPROVES BASELINE METHODS
To validate the effectiveness of our method, we apply FEDDECORR to four baselines, namely Fe- dAvg (McMahan et al. 2017), FedAvgM (Hsu et al., 2019), FedProx (Li et al., 2020), and MOON (ii et al.) 2021b). We partition the three benchmark datasets (CIFARIO, 10, CIFARI00, and Tiny- ImageNet) into 10 clients with where local models should be free from the pitfall of excessive dimensional collapse, we only expect FEDDECORR to perform on par with the baselines in this setting.
(0.05, 0.1, 20). Since
is the homogeneous setting
We ob- serve that for all of the heterogeneous settings all datasets, the highest accuracies are achieved
We display the CIFAR10/100 results in Tab. M and the TinylmageNet results in Tab. 2

---

Published as a conference paper at ICLR 2023
(a)
0.05
(b) 0.1
(c) 0.5
Accuracy
30
30
20-
20
10
10
10
01
comm. round
comm. round
comm. round
w/o FedDecorr
w/ FedDecorr
Figure 4: Test accuracy (% at each communication round. Results are averaged over 3 runs. Shaded areas denote one standard deviation above and below the mean.
by adding FEDDECORR on top of a certain baseline method. In particular, the strongly hetero- geneous settings where E {0.05, 0.1 adding FEDDECORR yields significant improvements of around 2% ~ 9% over baseline methods on all datasets. On the other hand, for the less heteroge- neous setting of a 0.5, the problem of dimensional collapse is less pronounced See 3 leading to smaller improvements from FEDDECORR. 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 a = oc, FEDDECORR still produces around 2% of improvements on the TinylmageNet dataset. We conjecture that this is because TinyImageNet much more com- plicated 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 TinylmageNet can benefit from FEDDECORR even in the homogeneous setting.
discussed in
To further demonstrate the advantages of FED- DECORR, 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 test- ing accuracy as threshold, we see that adding FEDDECoRR significantly reduces the number of communication rounds needed to achieve the given threshold. This further shows that FED- DECORR not only improves the final perfor- mance, but also greatly boosts the communica- tion efficiency in federated learning.
# clients
Method
a: 0.05
0.1
0.5
FedAvg
35.02
39.30
46.92
10
FEDDECORR
40.29
43.86
50.01
FedAvg
31.21
35.30
43.64
20
FEDDECORR
39.41
41.27
46.17
FedAvg 26.20
30.88
37.22
30
FEDDECORR
36.50
39.02
44.38
FedAvg 25.70
28.88
34.89
50
FEDDECORR
34.50
36.67
42.34
FedAvg
21.53
24.69
30.21
100
5.3 ABLATION
FEDDECORR
30.55
33.85
38.65
STUDY ON THE NUMBER OF CLIENTS
Table 3:
Ablation study the number of clients. Based on TinylmageNet, we run exper- iments with different number of clients and differ- ent amounts of data heterogeneity.
Next, we study whether the improvements brought by FEDDECORR are preserved as num- ber of clients increases. We partition the Tiny- ImageNet dataset into 10, 20, 30, 50, and 100 clients according to different a's, and then run FedAvg with and without FEDDECORR. 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 FEDDECORR increase from around 3% ~ 5% to around 7% ~ 10% with the growth in the number of clients. Therefore, in- terestingly, we show through experiments that the improvements brought by FEDDECORR can be even more pronounced under the more challenging settings with more clients. Moreover, our exper- imental results under random client participation show that the improvements from FEDDECORR robust to such uncertainties. These experiments demonstrate the potential of FEDDECORR to be applied to real world federated learning settings with massive numbers of clients and random client participation.
8

---

Published as a conference paper at ICLR 2023
CIP CIPAR10 (a 0.06)
CIFARI1 (a = 0.1)
TinylnngeNe et (a = 0.06)
TinyImageNel (ar 0.1)
72-
Accuracy
68
66-
77
36-
0.
0.1
012
073
0.0
61
022
0.3
6.0
03
0.2
0.3
0.0
011
02
03
8
FedDecorr
w/ FedDecorr
Figure 5: Ablation study on 3. We apply FEDDECORR with different choices of [3 on FedAvg.
5.4 ABLATION STUDY ON THE REGULARIZATION COFFFICIENT B
Next, we study FEDDECORR's robustness to the 3 in Eqn. 9 by varying the set {0.01, 0.05, 0.1, 0.2, 0.3). We partition the CIFARIO and TinylmageNet datasets into 10 clients with a equals to 0.05 and 0.1 to simulate the heterogeneous setting. Results are shown Fig. 5 We observe that, in general, when increases, the performance of FEDDECORR first increases, then plateaus, and finally decreases slightly. These results show that FEDDECORR is relatively insensi- tive to the choice of B3 which implies FEDDECORR is easy-to-tune federated learning method. In addition, among all experimental setups, setting 3 to be 0.1 consistently produces (almost) the best results. Therefore, we recommend 0.1 when having no prior information about the dataset.
5.5 ABLATION STUDY ON THE NUMBER OF LOCAL EPOCHS
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 51,5, 10, 20). We run experiments with and without FEDDECORR, and use the
CIFAR100
TinylmageNet
Method
0.05
0.1
0.05
0.1
FedAvg
50.67
55.98
32.31
34.88
1 + FEDDECORR
53.18
57.02
36.49
38.99
CIFAR100 and Imayent datasets with being 0.05 and 0.1 for this ablation study. Results are shown in Tab, 4 in which one ob-
FedAvg 59.57
65.02
36.02
40.75
5 FEDDECORR
61.42
65.98
41.68
44.77
serves that with increasing E, FEDAVG per- formance first increases and then decreases. This is because when E is too small, the lo- cal 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, FEDDECORR consistently improves over the baselines across different choices of local epochs E.
FedAvg 59.87
66.46
35.02
39.30
10 FEDDECORR
61.53
67.12
40.29
43.86
FedAvg 58.50
66.37
31.23
37.23
20
+ FEDDECORR
60.65
66.86
35.44
42.04
Table 4: Ablation study on local epochs. Experi- ments with different number of local epochs E.
5.6 ADDITIONAL EMPIRICAL ANALYSES
We present empirical analyses in Appendix C These include comparing FEDDECORR with other baselines (Appendix CC.4 and other decorreiation 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 FEDDECORR (Appendix C.I).
6 CONCLUSION
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 FEDDECORR, a novel method mitigate dimensional collapse during local training, thus improving federated learning under the heterogeneous data setting. Extensive experiments on benchmark datasets show that FEDDECORR yields consistent improvements over existing baseline methods.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
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 1-8000189-01-000 19-001-00,, 1-8000980-00--00) and a Singapore National Research Foun- dation (NRF) Fellowship (Grant Number: A-0005077-01-40 Jian Liang supported by National Natural Science Foundation of China (Grant No. 62276256) and Beijing Nova Program under Grant 221110000212100.
REPRODUCIBILITY STATEMENT
source code has been released at https://github.ccm https://gichu...om/byeeeeeeeeeee conlorraaaaeee///eeeeoooooo Pseudo-code FED- DECORR is provided in Appendix G We introduced all the implementation details of baselines and our method in Sec. 5.T In addition, the proofs of Theorem [] and Proposition [ are provided in Appendix A and Appendix
B respectively. All assumptions stated and discussed in the proof.
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas Navarro, Matthew Mattina, Paul N Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. arXiv preprint DrXiv:2/11.04000 2021.
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. arXiv preprint 2.XV020010.0523, 2020.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary, Fed- erated learning with personalization layers. arXiv preprint 1912v: 2019.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameteriiaaton.. In International Conference on Machine Learning, pp- 244-253. PMLR, 2018.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Variane-iivaranc--ovwaaaee regaaaaeeccocttaiaaaeegrrarriaai for self-supervised learning. arXiv preprint 2XXx:22005009066 2021.
Michael Cogswell, Faruk Ahmed, Ross Girshick, Larry Zitnick, and Dhruv Batra. Redueing overfit- ting in deep networks by decorrelating representations. arXiv preprint arXiv: 1511155111,,,..111 2015.
Ekin D Cubuk, Barret Zoph. Dandelion Mane, Vijay Vasudevan, and Quoc V Le. Autoaugment: Learning augmentation policies from data. arXiv preprint arXiv: 09501, 2018.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with the- oretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems, 33:3557-3566 2020.
Filip Hanzely, Slavomír Hanzely, Samuel Horváth, and Peter Richtárik. Lower bounds and opti- mal algorithms for personalized federated learning. Advances in Neural Information Processing Systems, 33:2304-2315, 2020.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown, Measuring the effects of non- non-identical data distribution for federated visual classification. arXiv preprint (00Xv::9090063355 2019.
Tianyu Hua, Wenxiao Wang, Zihui Xue, Sucheng Ren, Yue Wang, and Hang Zhao. On feature decorrelation in self-supervised learning. In Proceedings of the IEEE/CVF International Confer- ence on Computer Vision, pp. 9598-9608, 2021.
Lei Huang, Dawei Yang, Bo Lang, and Jia Deng. Decorrelated batch normalization. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 791-800, 2018
10

---

Published as a conference paper at ICLR 2023
Yutao Huang, Lingyang Chu, Zirui Zhou, Lanjun Wang, Jiangchuan Liu, Jian Pei, and Yong Zhang. Personalized cross-silo federated learning on non-iid data. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 7865-7873, 2021.
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. arXiv preprint arXiv: 180.0028 2018.
Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. arXiv preprint arXiv:2110.09548, 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference Machine Learning, pp. 5132-5143, PMLR, 2020.
Qinbin Li, Yiqun Diao, Quan Chen, and Bingsheng He. Federated learning on non-iid data silos: An experimental study, arXiv preprint arXiv:2102.0200 2021a.
Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10713-10722, 2021b.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Sys- tems, 2:429-450, 2020.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. International Conference on Machine Learning, pp. 6357- 6368. PMLR, 2021c
Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351-2363. 2020.
Mi Luo, Fei Chen, Dapeng Hu, Yifan Zhang. Jian Liang, and Jiashi Feng. No fear of heterogeneity: Classifier calibration for federated learning with non-iid data. Advances in Neural Information Processing Systems, 34, 2021.
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efffceeent learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1273-1282. PMLR, 2017.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295. 2020.
Karsten Roth, Timo Milbich, Samarth Sinha, Prateck Gupta, Bjorn Ommer, and Joseph Paul Co- hen. Revisiting training strategies and generalization performance in deep metric learning. In International Conference on Machine Learning, pp. 8242-8252. PMLR, 2020.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen. Mo- bilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 4510-4520, 2018.
Yujun Shi, Kuangqi Zhou, Jian Liang. Zihang Jiang, Jiashi Feng, Philip HS Torr, Song Bai, and Vincent YF Tan. Mimicking the oracle: an initial phase decorrelation approach for class incre- mental learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 16722-16731, 2022.
Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau en- velopes. Advances in Neural Information Processing Systems, 33:21394-211 2020.
Hongyi Wang, Mikhail Yurochkin. Yuekai Sun. Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. arXiv preprint arXiv:2002.06444) 2020a.
11

---

Published as a conference paper at ICLR 2023
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-7623, 2020b.
Lixu Wang, Shichao Xu, Xiao Wang, and Qi Zhu. Addressing class imbalance in federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 10165-10173, 2021.
Wei Xiong, Bo Du, Lefei Zhang, Ruimin Hu, and Dacheng Tao. Regularizing deep convolutional neural networks with structured decorrelation constraint. In 2016 IEEE 16th International Con- ference on Data Mining (ICDM), pp. 519-528. IEEE, 2016.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In Interna- tional Conference on Machine Learning pp. 7252-7261. PMLR, 2019.
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In International Conference on Machine Learning, pp. 12310- 12320. PMLR, 2021.
Michael Zhang, Karan Sapra, Sanja Fidler, Serena Yeung, and Jose M Alvarez. Personalized feder- ated learning with first order model optimization. arXiv preprint arXiv:2012.08565, 2020.
12

---

Published as a conference paper at ICLR 2023
A PROOF OF THEOREM IN MAIN PAPER
A.1 NOTATIONS REVISITED
Here, for the reader's convenience, summarize the notations used in both the main text and this appendix.
Notation
Explanation
Number of training data points.
C
Total number of classes.
X
The collection of the N training samples, X E RR
The collection of one hot labels of the training samples, y E RRCN The collection model output softmax vectors given all N input data, y RCXN
y
Y
(.(()
The i-th layer weight matrix at the (-th optimization step.
[I())
The product of the weight matrices of the first layers: II(t) = Wz(t) ... Wi(t).
oik
The k-th singular value of Wi.
u,k
The k-th left singular vector of W
Vi,k
The k-th right singular vector of W..
Ok
The k-th singular value IT. The k-th left singular vector of II. The k-th right singular vector of II.
Ub
Vi
N
Number of samples of class c. The proportion of class e samples w.r.t. the whole training samples: Me N The C-dimensional one-hot vector where only the e-th entry is 1.
He
ec
X
Mean vector of the training examples in class c: X = 1 Σ X,1 =
9e
Mean output softmax vector given samples in class c: To N. EE 7/ 7 = cc
A.2 TWO LEMMAS
Here, we elaborate two useful lemmas from Arora et al. 2019 2018).
The first lemma adopted from Arora et (2019:
Lemma 1. Assuming the weight matrix W evolves under gradient descent dynamics with infinites- imally small learning rate, the k-th singular value of this matrix (denoted as σyk evolves as
(yg(t) = (t)( (t)v
(10)
where un (1) and vx(1) are the k-th left and right singular vectors of W W(), respectively.
Proof. By performing an SVD on W(t), we have W(t) = (()) Therefore, by the chain rule differention, we have:
( (() U())((V()) +U()S(O)))T
(11)
Next, for both sides of the above equation, we left multiply U(I) and right multiply V(():
(())T(())V() (())(( + 5(t) (())
(12)
Since S(t) is a diagonal matrix, we consider the k-th diagonal entry of S(t), namely (t)::: (t) ( + (x(()) + ox)
(t)(t)) W(t)v, (t) ūú
(t). v.,)),
(13)
Since ug (t) and vt)(t) unit vectors and evolving in time with infinitesimal rate, have (f) ün (t) = 0 and Tv = 0. Next, Eqn. 13 can be simplified as
(t)(t) (u, (1)) (1))vx())
(14)
The proof is thus complete.
The second lemma is adopted from Arora et al. 2018
13

---

Published as a conference paper at ICLR 2023
Lemma 2. Given L consecutive linear layers in a neural network characterized by weight matrices W1. W2...., WL. We denote II = W----.... W1. We further denote (t) as weight matrix W after the t-th gradient descent optimization step. Correspondingly, the initialization of W is W (0).
Assuming have W, (W) (0))" (W (0)) (W))"W(++ W;+1(0) for any j E IL at initialization.
Then, under the gradient descent dynamics, [t(t) satisfies
Ii(t)= = - Σ [ ' d((I(()) jm1
[I(tt [[tt]] '
(15)
where [1] " and [1 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) A1, Ag...., Ans we denote diag A2........ Am to be the block diagonal matrix
A
0
0
2
0
diag( A1, A Am)
0
0
0
Am
Here, we first consider dynamics of arbitrary W2 where j E [L 1]. By the chain rule, we have (t) (l),.( (1)
aW,(f)
(16)
(1) ... T) (E(TT(())
On (Wi(t) ... Tt)
Given Equ. (66). we right multiply (t))) by (W)( and we left multiply W)+1(t) by
(W)+ which yields
W5(OWW(()) = (W)+
(17)
Applying the same trick on (t) and W yields
(,())(W(()) =
(18)
Adding Eqns. and 18) on both sides yields
(O) + WW)))(((0) = W(+1(() T (() + ( Next, by the chain rule for differentiation, Eqn. (19) directly implies that
(19)
d(w, (t), d(W)+1(t) ( (t))(t) dt
(20)
dt
Since we have assumed that W, (0) W,+ WW we can conclude that
(t) (t)
(21)
Next, we apply an SVD (t) and ((+z(() in Eqn. This yields
(1))) (OU, (1) (0S)
(22)
Based on Eqn. (22 and given the uniqueness property of SVD, we know:
(t) = (t)((()( = diag(pila, P2la..... Pm
(23)
Pm represent the distinct singular values satisfying #1 > p2 ... Pm > 0. and Ia. for any E [m] are identity matrix of size d, x dr. Since Eqn. (33 holds for any j. we know by induction that the set of values of p's is the same across all layers j [L]. In addition, there exist orthogonal matrices j E Rdd xd, for any r E [m] such that
where /PI.... Pm
Uy()) Vs+1 (1)iii))) Oj,y ,0,...)..,,,,)..
(24)
14

---

Published as a conference paper at ICLR 2023
Given Eqns. 24
next, we study (H) (+)) for any j IN 1}:
= UUH1))4tding(0,,000p00.. (0m).......,,)).5....... 0ss) diag(O J....... ,, SS 11
jmm)) j+1 Ujft
(S, commutes with diag((),, (plugging-in (4))
= S,H
Uj+idiag (eddag(ggl ( 23 las ... Pm P U
(25)
Similarly, holds that
W (t)WT 3+1 (t) = Ta, . Idm. V3
(26)
Next, by induction and Eqns. 25
W,. (t)
(t) ... (t)
Uzdiag(p] - Id, P2
L-s+1 I " --J+1 Idm YUT,
(27)
by induction and Eqns.
26 it holds that
W (TT(t) w (t) wt( = Vidiag(p
(28)
From Eqns.
27 we know that for any j E [L-1]}
[(t))[(tTT
WZ(t). ... Wi (t) ... (L)
Uzdiag(pt Ia, P5 Idy..... p Ian JUT
[Uddiag(p] - Id, P2 /-1 Idz................. P Iam JUZ] 1
(29)
WL W W)+ (I(()
Similarly, from Eqn. 28 we know that for any 1SiL--A
(I()) (It))= = Wi(t) ... WL(t). WL(t) W.((t)
Vidiag(pt La. p I Ids...., ...,, p tn Jam) VUT
(30)
[Vidiag( - pb - Iais Pot Iam VV
A
= [WT
With everything derived above, we now study the dynamics of II(t) as follows
ri(t)= . Ws+1 (W) 1 (w) Wi(t) (differential chain rule)
)E(I((t)
01 [WT ( (0W,-- wi( (plugging-in (16)
[( + (((I(())
[I((t) [I(t) +
(plugging-in in
29 and 30
011
(31)
This completes the proof.
15

---

Published as a conference paper at ICLR 2023
epoch: 3
epoch: 5
epoch:
epoch: 9
Figure 6: Alignment effects between the singular spaces WL+1(t and [T()). We train a 3-layer linear neural network on the MNIST dataset and visualize the models at 3, 5, 7, 9 training epochs, respectively. In each figure, the W-th row and k-th column pixel the value of (u, V (t)l- Darker colors denote values close to 1 while lighter colors denote values close to O. From the figures, we empirically observe that Lu (L) VI+11. (1) 1(k = k'} approximately holds.
A.3 ASSUMPTIONS
Assumption 1. We assume that the initial values of the weight matrices satisfy (0) W/O,), (0) for any i [L [L-1] 1].
Assumption 2. We assume [t VL+1L(()) = Lik- = k'} holds for all t, where ug(t) is the k-th left singular vector of I1()) and VL+1E the k'-th right singular vector of Wi+1
Remark: For Assumption it can be achieved practice by proper random initialization. For Assumption 2 Ji & Tellgrsky (2018) proved that under some assumptions, gradient descent opti-
mization 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.4 PROOF OF THEBREE 1 IN MAIN TEXT
Theorem 1 (formally stated). Let (t) for le E [d] be the k-th largest singular value of II(t).
Then, under Assumptions we have
(x NL NL ( ( + M
(32)
where (t) is the k-th left singular vector of WL+1(t), vx(() is the k-th right singular vector IIt),, M is constant, and G(t) is defined as
G(I) Σ (e (o(1))
(33)
Proof. Recall that for (L layer linear neural networks, given the i-th training sample X; Rd.
we have
(1) softmax ( (4)x,,, softmax (W (W))
(34)
and the loss is the standard cross-entropy loss defined as follows
((TTL), WT W (1) =
yI (().
(35)
By the chain rule, we can derive gradient of & with respect to Wi++ and II, which are respectively.
aE(IT(L), WL+1(1) (y -+))))) (I()) OW
(36)
and
aE(II(t), WL+1(t) = -W2+1(1) - T
(37)
Ø11
16

---

Published as a conference paper at ICLR 2023
Next, under the gradient descent dynamics, the dynamics WL+1 satisfies (1++(() = (E)((()) Wr) (y (1)) (())x"'(()"
(38)
aw
while the dynamics on requires invoking Lemmal
2 which allows us to write
ii(t) = Σ [( (e())(t))
011 [I(() [(t)]
(39)
(I((t)
Next, we invoke Lemma on Eqn. 33) and Eqn.
38
respectively, yielding:
(t) (u)(t)) (I()))
" ( (x)( (I((t) = vx(()
L(on(t))² ( (y -()) va((t) (v((t)
(SVD on IT()))
= [(on(t))² [OLtINuu(t) (122L() (x)(() Vvx(t)
vx(()
(SVD on (z))
[(on(t))²- 1 v.((t) (Assumption 2
(40)
and
bL+1,((t) = uL+IN(() -()x VL+IN(()
E ComuCHIA()) ( - -(())) (w (o)
(41)
(RULLIL (y - ((1))xTvv(( (Assumption
Combining Eqns. (40) and (41 we have:
=
(42)
Next, apply integration on both sides, which yields
(OL+IA(t))² = (t)( + M,
(43)
where M a constant.
By Equ. 44
Eqn.
40 can be rewritten
dz = ((ax(t))² - & (ooa(t) + M (x vt()).
(44)
Finally, notice that (y x can be rewritten as
N
(y -
(t)xT
N Σ Hale (T
(45)
1=1
1=1
We further substitute Eqn.
45
into Eqn.
and obtain
(t) = (T() (ax(t)) + +M G())v(())
(46)
where G(t) is defined as
G(t)
Σ - - 8t)
(47)
c=l
This completes the proof.
17

---

Published as a conference paper at ICLR 2023
B PROOF OF PROPOSITION 1 IN THE MAIN PAPER
Proposition (restated), For a d- d-by-d correlation matrix K with singular values (A.....;Aa we have:
d
d
2 \1 = |K|E d.
A;
(48)
d
j=1
Proof.
Given d-by-d correlation matrix K, since the diagonal entries of K all have
tr(K) d.
(49)
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.
have:
(1: 1)²
(Plug-in Eqn. (49)
Σ -
A + d
(50)
1=1
i=1
Σ X2 d
(Plug-in Eqn. (49
Next, for the right-hand side of Eqn. 77. we have:
| d = tr(K K) d = T
(Apply SVD on K)
= =rr(USSTU) -d
(51)
= x3 d.
Therefore, we have shown that the left-hand side of Eqn.
equals its right-hand side.
C ADDITIONAL EMPIRICAL ANALYSES
C.1 COMPUTATIONAL EFFICIENCY
We demonstrate FEDDECORR'S advantage vis-à-vis some of its competitors in terms of its compu- tational efficiency. We compare FEDDECORR with some other methods that also apply additional regularization terms during local training such as FedProx and MOON. We partition CIFAR10. CI- FARI00 and TinyImageNe required for one round of training for FedAvg, FedProx, MOON, and FEDDECORR Results are shown in Tab. 5) All results are produced with NVIDIA Tesla V100 GPU. We see that FED- DECORR incurs negligible computation overhead on top the naive FedAvg, while FedProx and MOON introduce about 0.5 ~ 1x additional computation cost. The advantage of FEDDECORR in terms of efficiency is mainly because only involves calculating the Frobenius norm of a matrix which is extremely cheap. Indeed this regularization operates the output representation vectors of the model. without requiring computing parameter-wise regularization like FedProx nor extra forward passes like MOON.
into 10 clients with a = 0.5 and report the total computation times
18

---

Published as a conference paper at ICLR 2023
CIFARIO CIFART00 TinylmageNet
FedAvg
6.7
6.9
25.4
FedProx
12.1
12.3
33.2
MOON
12.2
12.7
38.1
FEDDECORR
6.9
7.1
25.7
Table 5: Comparison of computation times. We report the total computation times (in minutes) for one round of training on the three datasets for FedAvg, FedProx, MOON, and FEDDECORR. Here, FEDDECORR stands for applying FEDDECORR to FedAvg.
C.2 COMPARISON WITH OTHER DECORRELATION METHODS
Some decorrelation regularizations such DeCov (Cogswell et al. 2015) and Structured- DeCov (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 represen- tations instead of the correlation matrix like our proposed method FEDD--EEDDRRRRRR To compare our FEDDECORR with the existing decorrelation methods, we follow the same procedure in FED- DECORR and apply DeCov and Structured-DeCo during local training. Our experiments are based on TinylmageNer and FedAvg. TinyImageNet partitioned into clients according to various a's.. Results shown in Tab. 6 Surprisingly, see that unlike FEDDECORR which steadily im- proves the baseline, adding DeCov or Structured-Dect 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.
both degrade the performance in federated
FedAvg
DeCov St.-Decov
FEDDECORR
a = 0.05
35.02
32.88
32.04
40.29
0.1
39.30
37.29
37.74
43.86
a = 0.5
46.92
46.29
45.85
50.01
Table Comparison with other decorrelation methods. Based on FedAvg and the TinylmageNet dataset, we use different decorrelation regularizers in local training.
C.3 EXPERIMENTS ON OTHER MODEL ARCHITECCTURSS
In this section, demonstrate the effectiveness of method across different model architec- tures. Here, besides the MobileNerV2 used in the main paper, we also experiment on ResNet18 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 Fed- Decorr 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.05 and we use the CIFARIO 10 dataset. Our results are shown in Tab. 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.
MobileNetV2 ResNet18 ResNet32
FedAvg FEDDECORR
64.85
71.51
65.76
73.06
76.54
67.21
Table 7: Effectiveness of FEDDECORR other model architectures.
19

---

Published as a conference paper at ICLR 2023
CIFARIO
CIFARTO0
Method
a = 0.05
0.1
0.5
8
0.05
0.1
0.5
0C
Scaffold
51.99 34
87
29
54.51
68.37
70.97
FedNova
63.07 7
90.23 10.41 92.39 1:018
60.22.:0.3
66.43 10.26
71.79.t0.1
74.47
FedAvg
64.85_2.01
76.28
89. 14 92.39
20
66.4620.1 16 71.69=a.15
74.54
FEDDECORR 73.06=aa11
80.60
89. anas 92.19
61.53
67.12
71.91
=0.02
73.87
FedProx
64.11
89. 39.57 00.10
92.38
40
66.41 =0.27
71.78=0.19 74.34
FEDDECORR 71.38
81.74 89, + 92. +.000
+o.3
67.00
71.64
74.15=0m0
FedAvgM
71.34
77.51 +0.58 88.39 +0.17 91.35 +..1s
59
66.36
71.17 74.20
FEDDECORR 73.60 -ox2 79.21 10.15 88. + 91.33 +-.13 61.48 4: 66.60 71.26
73.86+0.25
MOON
6..79_m.6 78. 0..66 90.08 00.08:22 = 92.62
56.79
20
71.81 10 74.30=a.12
FEDDECORK 73.46 81.63 0.55 90.61 t.0.5 92.63++..1 49.4 66. 12 #0.20 71.68 +0.0 +3.7
Table 8: CIFAR10/100 Experiments. We run experiments under various degrees of heterogeneity (a E (0.05, 0.1, 0.5, oo)) and report the test accuracy (%). All results are (re)produced by and are averaged over 3 runs (mean + std). Bold font highlights the highest accuracy in each column.
We add results of Scaffold and FedNova comparing to Tab. I in the main paper.
Tiny
Method
a = 0.05
0.1
0.5
OC
Scaffold
35.16 78 44.24 44.8
FedNova
35.28 39.73 47.05 49.57
FedAvg
46.92 49.3 +919 35.02 39.300
FEDDECORR 40.29
50.01 #0.7 52.63 +026
FedProx
35.20 0..0 39.66
47.117..60000 16 49.76 4026
FEDDECORR 40.63 0.05 44.19 0.114 50.26 52.37 4
FedAvgM
34.81 34.81=....9.7722... 47.
1004 49.67
+ FEDDECORR
39.97 43.95 0.26 54.14........ 10.11 52.05
MOON 35.23 40.53
47.25
50.48
FEDDECORR 40.40 44.20 0022
2.22 50.81 10.51 53.01
Table 9: Tiny ImageNet Experiments. We run with E {0.05, 0.1, 0.5, oo} and report the test accuracies (%). All results are (re)produced by us and are averaged over 3 runs (mean + std is reported). Bold font highlights the highest accuracy in each column. We add results of Scaffold and FedNova comparing to Tab. 2 in the main paper.
C.4 COMPARISON WITH OTHER FEDERATED LEARNING BASELINES
In this section, we compare FEDDECORR with two other baselines, namely Scaffold (Karimireddy al. 2020) and FedNova (Wang al. 2020b). We use the same experimental setups as in the main paper to implement these two baselines. Results on CIFAR10/10 and TinyImageNet are shown in Tab. g and Tab. D respectively. As shown in the tables, across various datasets and degrees of heterogeneity, adding FEDDECORR on top of a baseline method can outperform the baselines when there is some heterogeneity across the agents, i.e., a < 00.
C.5 EXPERIMENTS ON ANOTHER TYPE OF HETEROGENEITY
In this section, we run experiments under another type of data heterogeneity. Specifically, we follow McMahan et al. (017)) and split the CIFARI0 dataset across different clients such that each client only has a fixed number of classes C' (c.g.. = 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 shown in Tab. TO As can
20

---

In the main text, we have shown that global models produced by FedAvg (MeMahan et al., 2017) suffer stronger dimensional collapse with increasing data heterogeneity. To further show such di- mensional 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 follow the Published as a conference paper at ICLR 2023
FedAvgM
FedProx
MOON
log singular values
(Pann)
& Figure 7: Data heterogeneity causes similar dimensional collapse on other federated learning meth-
k
k
ods such as FedAvgM (Hsu et al. 2019), FedProx (Li ct al. 2020), and MOON (Li al. 2021b).
The x-axis (k) is the index singular values.
ResNet32
ResNet18
log singular valuess
themee)
(hams)
25
a=025
30 40
50
20
40
60
100
Figure 8: Data heterogeneity causes similar dimensional collapse on other model architectures dur- ing federated learning. The x-axis (k) is the index of singular values.
be observed, under this different heterogeneity scenario, FEDDECORR also yields noticeable and consistent improvements.
CC= = 2 C = 3
FedAvg
45.61
67.53
+ FEDDECORR
47.63
74.51
Table 10: FEDDECORR yields noticeable and consistent improvements under another type of data heterogeneity.
D ADDITIONAL VISUALIZATIONS ON GLOBAL MODELS
In this section, we provide additional visualizations on global models with different federated learn- ing methods, model architectures, and datasets. Through extensive experimental results, demonstrate that dimensional collapse is a general problem under heterogeneous data in federated learning.
D.1 VISUALIZATION ON GLOBAL MODELS OF OTHER FEDERATED LEARNING METHODS
et al. 2019), FedProx (Li et al.] (020), and MOON (Li et al. 2021b). Specifically, same procedure as the main text and plot the singular values of covariance matrices of represen- that all these three other methods also demonstrated the similar hazard of dimensional collapse as in FedAvg.
tations. Results are shown in Fig. 7 From the figure, one
D.2 VISUALIZATION ON GLOBAL MODELS OF OTHER MODEL ARCHITECTURES
In the main text, we have shown the dimensional collapse on global models caused by data het- erogeneity with Neb2.e In this section, perform the similar visualization based other
21

---

Published as a conference paper at ICLR 2023
CIFAR10
TinylmageNet
log singular singular values
80 300
20
60
Figure 9: Data heterogeneity causes similar dimensional collapse on other datasets during federated learning. The x-axis (l) is the index of singular values.
model architectures such ResNet32 and ResNet18. Note that ResNet32 is narrower ResNet whose representation dimension 64 and ResNet18 is a wider ResNet whose representation dimen- sion is 64. We visualize the top 50 singular values for ResNet32 and the top 100 singular values for ResNet18. Results are shown in Fig. 8 From the figure, one can observe that heterogeneous data also lead to dimensional collapse ResNet32 and ResNet18.
D.3 VISUALIZATION ON GLOBAL MODELS OF OTHER DATASETS
In the main text, we use the CIFAR100 dataset for our visualizations. In this section, we perform similar visualizations with other datasets such CIFARIO and TinylmageNet. Results are shown in Fig. 9 From the figure, one can also observe that dimensional collapse results from data hetero- geneity.
E VISUALIZATION ON OTHER LOCAL CLIENTS
In the main text Fig, bb.,, under the four different degrees of data heterogeneity (i.c., {0.01, 0.05, 0.25, Da }). we compare representations of local models of client and empirically show how data heterogeneity affects representations produced by the local models. In this section, fur- ther corroborate our conclusion, we follow the same procedure and visualize singular values of the covariance matrix of representations produced
local models trained the rest of the 9 clients under the same a s. Results are shown in Fig. I0 From the results, we can obtain the similar obser- vations in Fig. (b) of the main text, namely that stronger data heterogeneity causes dimensional collapse for local models.
F HYPERPARAMETERS OF OTHER FEDERATED LEARNING METHODS
The regularization coefficient {10-4, 10 10 10-1} and is selected to be μ 10 the regularization coefficient of MOON (Li et al. 202 lb) ju is tuned across f0.1, 1.0, 5.0, 10.0) and is selected to be = 1.0; the server momentum of FedAvgM (Hsu et al. 2019) p is tuned across (0.1, 0.5, 0.9) and is selected to be =
FedProx
et al. 2020 1 tuned across
G PSEUDO-CODE OF FEDDECORR
Here, we provide a pytorch-style pseudo-code for FEDDECORR Alg [ All FEDDECORR- components are highlight in blue. As indicated in the pseudocode, the only additional operation of FEDDECORR is in adding a regularization term LeedDecorr (w, X) defined in Eqn. (8). This shows that FEDDECORR is an extremely convenient plug-and and-play federated learning method.
H STABILITY OF FEDDECORR REGULARIZATION LOSS
In this section, we first split CIFARIO into 10 clients with a = 0.5. Then, we plot how FedDecorr loss evolve within 10 local epochs for all the 10 clients in Fig. I All training configurations
the same as in the main paper. From the results, one can observe that the optimization process of FedDecorr loss is stable.
22

---

Published as a conference paper at ICLR 2023
Client 2
Client
Client 4
e-0.25
-2
302
Client
Client
Client
log singular values
(romo)
-- [toma]
(homst
20
100
20
60
100
20
60
100
Client
Client
Client 10
(hemo)
[tama)
g-t.25
k
Figure 10: Heterogeneous local training data cause dimensional collapse. For each of the clients, given the four models trained under different degrees of heterogeneity, we plot the singular values of covariance matrix of representations in descending orders (the results of client are shown in main text Fig. b)). Representations are computed over the CIFAR100 test set. The r-axis (k) the index of singular values and the y-axis is the logarithm of the singular values.
Cleat
Clent2
Clent 3
Cleet
Clert
22
034
0.29
02
6.55
15
FedDecorr
414
3300
Clest
Clent
Clent 8
Clent
Clant 10
0.25
033
Local Epochs
Figure 11: How FedDecorr loss evolve within 10 local epochs.
23

---

Published as a conference paper at ICLR 2023
Algorithm PyTorch-styye Pseudocode for FEDDECORR. (Blue highlights FEDDECORR-specifee code)
FedDecorrLoss (2)
N: batch size
d: representation dimension
batch representation, with shape (N,
N. z.shape
normalization
(z z.mean z.std(0)
estimate correlation matrix
corr.mat 1/N4torch.matmmnn matmul (z.t(), z)
calculate FedDecorr loss return loss fed_deorr
(corr_mat, ccor__aat.pox(2) .maa ()
LocalTraining (train loader, for e in range (local (local_epochar epochs) for data, targets in loader:
epochs, beta):
propagation.
# given the batch of data, compute batch representations z and loss
lo88,
loss bata+FadDecorr(s)
back propagation update local model parameters
def (lobalAggregationn)
receiving models from each clients
aggregating local models with certain schemes
sending aggregated models back to clients
def main ()
round: number of communication rounds. loader training data. n local.l epochs: number of local trainig epochs on each client.
train loader:
beta: coefficient of the FedDecorr regularization. range comm_rooud)
Localfraining (train loader, l_cocal opochs, beta) Globalkggrogation)))
24


Reviews:
[{"review_id": "wse7WPYY5_-", "reviewer": "Reviewer_deuX", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper makes an interesting observation in learning from heterogeneous data sources by empirically illustrating that representations tend to reside in a lower-dimensional space.  To overcome this issue, authors propose FEDDECORR algorithm that encourages the Frobenius norm of the correlation matrix of representations at each client (penultimate layer of a neural network) to be small.\n", "summary_of_the_review": "The paper makes an interesting observation and algorithmic solution, but the theoretical results and empirical results need to be strengthened to resolve  issues discussed above.\n\n", "strength_and_weaknesses": "The paper makes an interesting observation and has some merits, but \n\nOn the empirical side, \n\n1. There is a discrepancy between experiments the observation made from and proposed algorithmic solution: the way the heterogeneity is introduced among clients to infer empirical observation include discrepancy among both covariates and labels, while the proposed algorithm only tackles representations by regularizing  covariance matrix of the representations. This left me wondering about the contribution of each source of heterogeneity on dimensional collapse. \n\n2. Also, in the results reported in Table 4, it seems as the number of local updates increases, the proposed algorithm does not show significant improvement over baseline (FedAvg). Since in federated settings the number of local updates could be much larger than 20, it seems the proposed regularization becomes less beneficial (it is conjectured by authors that the dimensional collapse of the global model stems from local models). \n\n3. Finally, it would be nice to see how the observed  phenomena manifest itself in other algorithms where a drift/diversity mitigation schema is employed. This could potentially better illustrate the role of heterogeneity in dimensional collapse and the effectiveness of the proposed method. \n\nOn the theoretical side, the analysis is limited to linear neural networks, does not take into account the local updating and aggregation schema, and more importantly does not clearly illustrate how covariance regularization can improve accuracy as it analyzes the gradient flow of optimization which does not necessarily translate to a good generalization in the presence of data heterogeneity.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "rIF15cQ3Mf", "reviewer": "Reviewer_h8K7", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies dimension collapse in federated learning with label distribution shift, and methods on more to mitigate it. With the simple observation that the singular values of features are decaying fast for heterogeneous cases, the authors propose to minimize the corresponding variance, or equivalently, the Frobenius norm of the correlation matrix. This F-norm regularization has been tested on vision datasets with several baseline comparisons.\n\n--post rebuttal--\n\nThanks to the authors for showing extra experiments on new data splits and datasets as well as comparison with new baselines. From the discussion, this paper is primarily experimental, and there is not much theoretical analysis on the explanation of $\\beta$, or which algorithm  is the best when combined with DeCorr. These are some future directions. ", "summary_of_the_review": "In summary I would recommend weak accept as this paper identifies the problem of dimension collapse in FL and proposed a simple yet effective method to solve it. There are still some remaining issues with this method, such as baseline comparison, and deeper understanding beyond linear dynamics. The authors may start by proving faster convergence rate or better generalization error for FedDecorr, which could further improve the theoretical understanding. Also, it remains unclear to me why for TinyImageNet with $\\alpha = \\infty$ FedDecorr still has moderate improvement. Moreover, without FedDecorr, all the baseline algorithms do not beat FedAvg on CIFAR-100 and TinyImageNet by a large margin (1%), which raises questions for the correctness of the implementation.\n\n", "strength_and_weaknesses": "Strengths:\n1) Data heterogeneity is common in FL and this paper studies this problem from the perspective of singular values of features. This perspective is interesting and has not been studied before. The paper is well-motivated.\n2) The proposed method is easy to understand, by equalizing the variance of singular values of the covariance matrix, with the Frobenius norm of the correlation matrix, the authors propose to use the F-norm as a regularization method. This is easy to implement and can be adapted to most existing FL methods, such as FedAvg, FedProx, MOON, etc.\n3) This method has been tested on a variety of vision datasets, with different levels of data heterogeneity.\n\nWeaknesses:\nAlthough this paper proposes an interesting idea for mitigating the label shift problem, I still feel it needs some improvement.\n\n1) Sec 3.3 seems unnecessary and is detached from the main paper. The linear network assumption seems too restrictive and does not provide us with new insights. One can simply obtain eq. (6) from the observation in Figure 2, and there is no need to use gradient flow dynamics.\n2) Only Dirichlet partition has been considered. More testing needs to be done on other splits. Also, realistic partitions should be tested, such as FEMNIST, Shakespeare and StackOverflow.\n3) The current method seem more like heuristics and it needs more understanding. For example, why is $\\beta = 0.1$ more desirable? Why increasing $\\beta$ can sometimes decrease the performance? \n4) About baseline comparison, the authors should add more competitive baselines such as FedAdam [1] and FedLC [2] that aims to address the same data heterogeneity problem. \n\n\n[1] Reddi, Sashank J., et al. \"Adaptive Federated Optimization.\" International Conference on Learning Representations. 2020.\n[2] Zhang, Jie, et al. \"Federated Learning with Label Distribution Skew via Logits Calibration.\" International Conference on Machine Learning. PMLR, 2022.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "F98RvlvgcqX", "reviewer": "Reviewer_QKrC", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper identifies a key bottleneck in highly heterogeneous federated systems, namely dimensional collapse. The authors propose a simple but novel technique for mitigating dimensional collapse by adding a regularization term during local training. The paper shows both theoretical and empirical analysis to support their results.", "summary_of_the_review": "The paper points out a key problem in highly heterogeneous federated systems and proposes a novel approach to solving it. The study is a good direction for future research as well. The authors back up their arguments with both rigorous theoretical and experimental results on image classification tasks. I do have reservations regarding the stability of the algorithm and its effectiveness in other domains. Nevertheless, the results are more or less rigorous and show it is an effective technique.", "strength_and_weaknesses": "The paper identifies a key issue in data heterogeneous federated learning. The overview of dimensional collapse in FL is a significant stepping stone for future research. The proposed algorithm is simple yet effective. The paper also shows theoretical and experimental results which makes their arguments strong. I do have some suggestions.\n\n1. I believe it would be a good to see how the algorithm behaves when we vary the client sampling rate in addition to the number of clients.\n2. Add a convergence analysis for the algorithm if tractable.\n3. An ablation study studying how the added regression evolves over time is important.\n4. Add more tasks besides image classification.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "86AoU_7FPG", "reviewer": "Reviewer_mWEv", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper considers the federated learning problem with data heterogeneity. From experiments, the authors find an interesting phenomenon, dimensional collapse. The global model suffers from dimensional collapse in the feature space such that the covariance matrix of the representations of the global model becomes approximately low-rank. The model collapse also happens in local models. The authors claim that the model collapse is one of the main reasons that degrade the FL performance. To remedy this problem, they propose FedDecorr, which applies a Frobenius norm regularization term during local training, making the distributions of singular values more uniform. FedDecorr outperforms baselines on standard benchmark datasets.\n", "summary_of_the_review": "This paper proposes a simple yet efficient algorithm for FL with a good justification.", "strength_and_weaknesses": "Strengths\n\n- This paper is easy to follow.\n\n- The proposed algorithm is well-motivated and performed very well.\n\n- The dimensional collapse is an interesting observation.\n\nWeakness\n\n- More recent FL algorithms should be discussed and compared. \n\n- More experiments with other data sets and neural net architectures might be required.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster