FedLWS: Federated Learning with Adaptive Layer-wise Weight Shrinking
Authors: Changlong Shi
Jinmeng Li
He Zhao
Dan dan Guo
Yi Chang
Venue: ICLR 2025 Poster (2025)
Keywords: Federated Learning
Model aggregation
Deep neural networks
Machine learning
Primary area: other topics in machine learning (i.e., none of the above)

Abstract: In Federated Learning (FL), weighted aggregation of local models is conducted to generate a new global model, and the aggregation weights are typically normalized to 1. A recent study identifies the global weight shrinking effect in FL, indicating an enhancement in the global model’s generalization when the sum of weights (i.e., the shrinking factor) is smaller than 1, where how to learn the shrinking factor becomes crucial. However, principled approaches to this solution have not been carefully studied from the adequate consideration of privacy concerns and layer-wise distinctions. To this end, we propose a novel model aggregation strategy, Federated Learning with Adaptive Layer-wise Weight Shrinking (FedLWS), which adaptively designs the shrinking factor in a layer-wise manner and avoids optimizing the shrinking factors on a proxy dataset. We initially explored the factors affecting the shrinking factor during the training process. Then we calculate the layer-wise shrinking factors by considering the distinctions among each layer of the global model. FedLWS can be easily incorporated with various existing methods due to its flexibility. Extensive experiments under diverse scenarios demonstrate the superiority of our method over several state-of-the-art approaches, providing a promising tool for enhancing the global model in FL.

Published as a conference paper at ICLR 2025
FEDLWS: FEDERATED LEARNING WITH ADAPTIVE LAYER-WIS WISE WEIGHT SHRINKING
Changlong Jinmeng Li', He Zhao?, Dandan Guo',, Yi Chang 34+ 334. School of Artificial Intelligence, Jilin University! CSIRO's Data61² International Center of Future Science, Jilin University Engineering Research Center of Knowledge-Drine Driven Human-Maching Intelligence, MOE, China* (shic122, 1i jms921) @mails jlu. edu. en, he zhao@data611. csiro.au au, (guodandan, yichang) Bjlu. edu.
ABSTRACT
In Federated Learning (FL), weighted aggregation of local models conducted generate a new global model, and the aggregation weights are typically normalized 1. A recent study identifies the global weight shrinking effect FL, indicating an enhancement in the global model's ' generalization when the sum of weights (i.e., the shrinking factor) smaller than 1, where how to learn the shrinking factor becomes crucial. However, principled approaches to this solution have not been carefully studied from the adequate consideration of privacy concerns and layer-wise distinctions. To this end, we propose novel model aggregation strategy, Federated Learning with Adaptive Layer-wise Weight Shrinking (FedLWS), which adaptively designs the shrinking factor in layer-wise manner and avoids opti- mizing the shrinking factors on a proxy dataset. We initially explored the factors affecting the shrinking factor during the training process. Then calculate the layer-wise shrinking factors by considering the distinctions among each layer of the global model. FedLWS can be easily incorporated with various existing methods due to its flexibility. Extensive experiments under diverse scenarios demonstrate the superiority of our method over several state-of-the- approaches, providing a promising tool for enhancing the global model in FL.
INTRODUCTION
Federated Learning (FL) an innovative paradigm in machine learning has attracted substantial attention in recent years (Li et al. 22200a] Zhang et al. 22021] Kairouz et al. 2021 Qi et al. (203)). This distributed optimization approach allows model updates be computed locally and aggregated without exposing raw data, thereby effectively addressing the challenges arising from distributed data sources while simultaneously preserving privacy and security.
FL., weighted aggregation of local models is conducted to generate the global model, where how to design the aggregation scheme is
critical problem (Kairouz et al. 2021 Qi et al. 2023
Ye et al. 2023).
Most previous works (Li et al. 2020b McMahan al. 2017) conduct model aggregation simply based on the local dataset relative size, which however could sub-optimal empirically due to the data heterogeneity. Consequently, many methods (Ye et al.) 2223 Hsu et al. 2019 Waag al.) 2020b) focus on adjusting the model aggregation scheme to enhance the global model's performance, with the majority of these methods normalizing the aggregation weights (i.e., the sum of aggregation weights is 1, represented as 7 = 1, where y is the sum of weights). Recently, Li et al. (2023a) revisited the weighted aggregation process and gained new insights into the training dynamics of FL, identifying the Global Weight Shrinking effect (GWS, analogous to weight decay) in FL, when is smaller than 1, which can further enhance the global model's 's generalization. The study demonstrates that during the training process, the shrinking factor y plays a crucial role in maintaining the balance between the regularization term and the optimization term. [i al. (2023a) proposed FedLAW, which optimizes the value of y using gradient descent on proxy dataset that has the same distribution as the global dataset. Despite its notable performance improvements, two key issues limit its practical applicability and adaptability.
"Corresponding authors.

---

Published as a conference paper at ICLR 2025
Firstly, the use of proxy datasets may raise privacy-related concems, which are particularly erucial in the context of FL. Additionally, obtaining a dataset with a distribution identical to the global dataset is challenging in practice, and discrepancies between the proxy and global data distributions can negatively impact the method's effectiveness. Secondly, divergence across different layers of deep neural networks has been demonstrated in numerous previous studies (Ma et al. [2022] Rehmman et al..
2023 Lee et al. 2023). Hence, simply applying a single shrinking factor across the entire model may not effectively harness the benefits of global weight shrinking."
To address the aforementioned issues, we initially conducted a series of experiments to explore factors influencing the shrinking factor's value during Federated Learning training. Building on these empirical observations from our experiments, we subsequently propose our method, Federated Learning with Adaptive Layer-wise Weight Shrinking (FedLWS). By analyzing the relationship between the regularization and optimization terms, and examining how their ratio correlates with the variance of local gradients, we derive an expression for the shrinking factor ry. This factor is computed directly using the gradients and parameters of the global model, which are readily accessible in Federated Learning. Consequently, FedLWS eliminates the need to optimize the shrinking factor Y on proxy dataset, as required by previous work (Li et al. 2023a). This not only addresses privacy- related concerns but also enhances its feasibility and applicability in practical, real-world deployments. Furthermore, by leveraging the gradients and parameters of each layer in the global model, it is easy to calculate the layer-wise for each layer. This allows for an improvement in model generalization by considering layer-wise differences of shrinking factor. Moreover, our approach is conducted after the server-side model aggregation. That is to say, it is orthogonal to many existing Federated Learning methods, making it easily integrated with them to further enhance model performance. To validate the effectiveness of our proposed FedLWS, we conduct extensive experiments under diverse scenarios. We observe that FedLWS can significantly improve the performance of existing FL algorithms.
The contributions of this work are summarized as follows:
We empirically show that the ratio of regularization optimization terms in FL model aggregation is positively correlated with local gradient variance. This allows us to directly calculate the shrinking factor, eliminating the need for fine-tuning on a proxy dataset.
We experimentally demonstrate that the model-wise shrinking factor is suboptimal to im- prove the generalization of the global model. It is essential to consider the discrepancy between different layers.
We propose FedLWS, simple method that generates the global model with better general- ization through layer-wise weight shrinking. Fe requires neither additional data nor transmission of the original data, thus raising privacy concern.
We conduct extensive experiments under diverse scenarios to demonstrate that FedLWS brings considerable accuracy gains over the state-of-the-art FL. approaches.
2 RELATED WORKS
Federated learning (FL) is a rapidly advancing research field with many remaining open problems to address. FedAvg (MeMahan al. 2017) has been the standard algorithm of FL. It trains local models separately and conducts model aggregation based on the data size. However, the distribution heterogeneity of local datasets may significantly degrade FL's performance. this paper, we aim to attain a more effective and better generalized global model through collaborative training on both the client and server sides. Many previous works focus on this can be mainly divided into two directions.
2.1 CLIENT-SIDE DRIFT ADJUSTMENT
Due to the data heterogeneity, local models trained on the clients may exhibit different degrees of bias, thereby affecting the performance of the global model. Many methods aim to reduce this bias by adjusting the training process of local models. FedProx (Li et al., 2020b) utilizes the (g-distance between the global model and the local model as regularization term during the training of the local model. FedDyn (Acar et al. 2021) proposes a dynamic regularizer for each client to align the global and local solutions. MOON (Li et al. (021)) aligns the features of global and local models through contrastive learning. FedDC (Gao et al. (222) and SCAFFOLD (Karimireddy et al. 2020) adjust the drift in the local model by introducing control variates. FedETF (Li et al.! (023b) employs fixed
2

---

Published as a conference paper at ICLR 2025
ETF classifier during training to learn unified feature representations. These methods concentrate on adjusting local models, merely assigning aggregation weights based on the size of the local dataset. In contrast, our approach focuses the server-side aggregation process, which can be easily combined with these methods to enhance model performance further due to its flexibility.
2.2
SERVER-SIDE AGGREGATION SCHEME
Several other previous works adjust the server-side model by improving the model aggregation stage or incorporating a fine-tuning step to obtain a better global model. FedAvgM (Hsu al., 2019) adopts momentum updates on the server-side to stabilize the training process. CCVR Luo et al.
2021) adjusts the global model's ' classifier using virtual representations. FedDF (Lin et al. 2020) and
FedBE (Chen & Chao 202T fine-tune the global model on the server. FedDisco (Ye et al. 2023) leverages both dataset size and the discrepancy between local and global category distributions to determine more distinguishing aggregation weights. Several prior studies have investigated layer- model aggregation. Fedl AMA (Lee et al. 2023) adjusts the aggregation frequency for each layer to reduce communication costs while accounting for inter-layer differences. pFedLA (Ma al. 2022)
focuses on personalized FL, it designs hyper-network to predict the layer-wise aggregation weights for each client. L-DAWA (Rehman et al. 2023) employs cosine similarity between local models and the global model as aggregation weights for model aggregation, simultaneously taking into account variations across different layers of the model. In comparison to these methods, our method is neither focused on aggregation frequency adjustment nor layer-wise aggregation weight adjustment. Instead, we propose an adaptive layer- layer-wise weight shrinking step after model aggregation to mitigate aggregation bias, which is both computationally efficient and modular, enabling seamless integration with various FL frameworks and basclines. Recently, FedLAW (Li et al., 2023a) identifies the global weight shrinking phenomenon and then learns the optimal shrinking factor Y and the aggregation weights À at the server with a proxy dataset, which is assumed to have the same distribution as the test dataset. Its success is inseparable from the high degree of consistency between the proxy data and the test data. Considering data privacy is a significant concern in FL, obtaining the proxy dataset with distribution identical the test dataset in practice is challenging. limiting its application in real-world. In addition, FedL.W ignores the variations across different layers of model for model aggregation. In contrast FedLAW, our FedLWS calculate the shrinking factors directly through the casily available gradient and parameters of the global model, which takes into account the layer-wise differences, avoiding demanding proxy dataset and optimization. Moving beyond FedLAA, ours can be easily integrated with most of related model aggregation methods for decoupling shrinking factor and aggregation weights.
3 BACKGROUND
Federated Learning consists of K clients and central server, where each client has its own private local dataset Dx- FL aims to enable clients to collaboratively learn a global model for the server without data sharing. In communication round f, the parameters of the global model and the client k's model are denoted as 9 and w respectively. The workflow of the basic FL method, FedAvg (McMahan al. 2017). in communication round can be described follows:
Step 1: Server broadcasts the parameters of global model w/ to each client;
Step 2: Each client & performs E. epochs of local model training on private dataset D3 to obtain local model w
Step 3: Clients upload the local models to the server;
Step 4: Server merges the local models to get new global model: w E daw
where d is the aggregation weight of the client & and FedAvg sets X Σr D.I'
D.
Denote g as the local gradient for the model of the k-th client during communication round t, where the local model wi. equals to w gk. in Step 2. Therefore, the update process for the global model can be expressed as follows:
K 7g
w ++1
w
Agg s.t. A 2. |All1 = 1,
(I)
k=
3

---

Published as a conference paper at ICLR 2025
where X (A1.... ..... Ax] is the aggregation weights and f2 is the global learning rate. FedLAW (Li eet al. 2023a) identifies the global weight shrinking phenomenon during the model aggregation in FL, and introduces global shrinking factor 0 < Y < 1, which improves generalization. The model aggregation process is then reformulated as follows:
K
Σ Aaw S.L. 1 7 0, 2 0, ||
(2)
where the 7 is the shrinking factor in the communication round Therefore, the right of Eq. can be rewritten as:
k=
w = "w' nggg.) = w'y - -7' 71-7")w's
-E where = =k=1 Aug, is the global gradient. We refer to the 7 r'yggi as the global averaged gradient (optimization term) and - >'''ww, is the pseudo gradient of global weight shrinking (regularization term). Therefore, the value of 7t determines the strength of regularization, with smaller values of y corresponding to stronger regularization (weight shrinking) and smaller gradient updates. It can be seen that the global weight shriking is analogous to weight decay, however they are are distinet. More discussions about their differences can be found in Section 5..3 and Appendix A.2
FedLAW (Li et al., 2023a) learns the optimal shrinking factor y and the aggregation weights À proxy dataset, which is assumed to have the same distribution as the test global dataset. Nevertheless, the utilization of a proxy dataset may raise privacy-relate concerns, and obtaining a dataset with a distribution identical to the test data in real-world applications can be challenging. Moreover,
numerous previous studies (Ma al. 2022 Luo al. 202T Rehman et al. 2023) have illustrated that the different layers of deep neural networks exhibit different levels of heterogeneity during the training process of FL, especially when training with non-IID data. However, FedLAW al.) 2023a) does not consider the discrepancy among different layers, assigning single shrinking factor y all layers of the model, which could impact the global model performance negatively, as an inappropriate value of y can affect the generalization. Therefore, this paper aims to address the challenge of learning an appropriate "y without relying on a proxy dataset, while also adaptively determining the corresponding y value for each layer.
4 METHOD
In this section, we propose Federated Learning with Adaptive Layer-wise Weight Shrinking (FedLWS), which dynamically computes the layer-wise shrinking factor for each layer of the global model. Unlike prior approaches, FedLWS does not rely on a proxy dataset, thus addressing the privacy and practical deployment issues present in traditional Federated Learning (FL) systems. Additionally, our method seamlessly integrates into existing FL methods to enhance performance.
4.1 FEDERATED LEARNING WITH ADAPTIVE WEIGHT SHRINKING
We first analyze the balance between regularization and optimization. From Equation 3 we can see that the shrinking factor 7' governs the trade-off between regularization (the model tends to keep its current weights) and optimization (the model tends to update its weights according to the gradients). smaller shrinking factor 7² imposes stronger regularization, while a larger 7 favors optimization by allowing larger updates. The challenge lies in finding an appropriate 7 that dynamically adapts to the training process. An idcal 7 should be able to maintain a balance between the regularization term introduce intuitive hypothesis: the balance
and the optimization term. To address this problem, between the regularization term and the optimization term should be related to the variance of the gradients of local client models. Specifically, if local client models agree less with each other in terms of the gradient updates (i.c., higher variance of their gradients), the global model should be more conservative in terms of updating its weights, meaning that stronger regularization is needed (i.e., smaller 7).. Conversely, if the gradient variance is smaller, one should update the global model's weights more aggressively according to the gradients (less regularization and higher 7t). Specifically, denoting the gradient variance of the local models as T", we formulate it as follows:
K
K
g
g ginaan
1
Σ
(4)
K
K
k=
4

---

Published as a conference paper at ICLR 2025
where gi is the local gradient of the k-th client. Therefore, given the discussion above, assume the gradient variance of the local models is proportional to the balance between the regularization and optimization term, expressed as:
T'or', pd (1 75) >- s.t. 0< 7d <1, 7 By
(5)
To empirically substantiate this hypothesis, we conducted the experiments employing the test dataset to optimize 7², which is impractical in real-world cases but for demonstrating purpose here. After the optimal value of xt is obtained, we compute the right-hand-side term of Eq and denote its value as r". For the left-hand-side term of Equation 5 T', we compute its value with Equation 4 To test the correlation between rf and 7' with different magnitudes of data heterogeneity, Dirichlet sampling Dira. to simulate client heterogeneity. A smaller value of a indicates higher degree of non-lll) data distribution. The experimental results are illustrated in Figure [(a] can be observed that as the degree of data heterogeneity diminishes, both r and 7 exhibit a corresponding reduction. This suggests that, the case of IID data, y tends to facilitate larger updates. More exploration and theory analysis can be found in Section [..3 Appendix B. and D We also compute Pearson correlation coefficient (Cohen et al., 2009) between r and т, which 0.860. Together with the analysis of Figure ((a it validates our hypothesis.
adopt
We further introduce the scaling term /3 as a hyperparameter:
Sr (1 - 7*) [[w]]
s.t. 0 < <1
(6)
y |
Combining Equation with Equation shrinking factor 7 as follows:
4 we can readily deduce the computational expression for the
Br m B1
w/ll
During the model aggregation process, both wi and 1988 w r++ w/J are known. Therefore, after calculating 7¹ via Equation 4 the value of shrinking factor 7 can be directly obtained through Equation 7 Equation 7 enables a comprehensive consideration of both the global gradient and model parameters, ensuring balanced interplay between the optimization and regularization terms during the FI. training process and eliminating the need for optimization using additional proxy datasets.
FS is conducted after model aggregation, making it easily be combined with other Federated Learning methods. For example, when combined with FedAvg, the first four steps of our method Step 4 in Section Following this, our method computes the shrinking factors and applies layer-wise weight shrinking to the aggregated model. In communication round t, k= A w represent the aggregated model before weight shrinking and denote the aggregated model after weight shrinking w
align with the Step
we use w r+1 = >
4.2 LAYER-WEL WAYER-WISE EXTENSION OF FEDLWS
Now, we investigate the effect of layer-wise factors on the performance of the global model. Previous studies have demonstrated divergence across various layers of deep neural networks (Ma et al. 2022
Rehman et al. 2023 Lee et al. 2023). Therefore, applying a single shrinking factor, 7 to the entire model may not fully capture the benefits of global weight reduction. We validate this through experiments, as shown in Figure ((b) the model we used is a 5-layer CNN, the layer-wise y refers to employing distinct shrinking factors for each layer of the model, with values uniformly decreasing from to 0.96. Figure ((b) illustrates that considering 7 in layer-wise manner, rather than using single y for the entire model, can further enhance the model's generalization performance.
Considering that each layer in deep neural networks may vary differently, it might be beneficial to design a respective regularization strength for each layer of the global model. Therefore, we aim to calculate the shrinking factor for each layer, using of to represent the shrinking factor of the 1-th layer of the global model in communication round t. Benefiting from the Equation 4 and 7 we can achieve the goal by calculating of follows:
5

---

Published as a conference paper at ICLR 2025
Pearson: 6.865
2.0
70
86
80
y=1
y=0.99
16
y 36
y=0. 97
y=0.96
Commuriication Rounds
(a) Gradient variance r ratio r under different degrees of data heterogeneity.
(b) Accuracy curves with different fixed ry.
Figure I: Empirical observations on CIFAR- 10 CNN as backbone; see more results in Appendix B.2 In (a) a is the degree of data heterogeneity, with smaller a indicating more heterogeneous data, illustrates the correlation between r and ratio. (b) indicates that layer-wise 7 can enhance the global model's 's performance.
Laseowle
Weight
Madd
Appregaise
wi @
8 w
wi. @
Training
The overview of Fe. @ server broadcasts the parameters of the global model to each client; 2 clients perform local training on their dataset; @ clients upload the local models the server; 4 the server conduct model aggregation to generate new global model w 11+ & the server calculates the layer shrinking factors 7: @ the serve conduct layer-wise weight shrinking and obtain the final global model w
Figure 2:
2/
IBu
gmeant
(8)
Brt ln
w f
K
where 8/, and w represent the gradients and parameters of the 1-th layer of the global model. Through Equation g the server obtains the layer-wise shrinking factors [71 .... and then conducts layer-wise weight shrinking to generate the new global model w as follows:
71 (w((W) .... 7 (wt1) (tt)) 92
(w
where the "o" denotes the connection between different layers within the model, w gl is the 1-th
layer's parameter of 1+1 Having introduced our adaptive way setting 7f. we give an overview of our method, FedLWS, in Figure 2 The pseudo-code of FedLW'S is shown in Appendix C.I Algorithm where we highlight the additional steps required by our method compared to FedAvg.
4.3 DISCUSSIONS
Privacy. Our proposed FedLWSS is more privacy-preserving and adaptable in practice compared to
some of the previous works (Luo et al. 2021 Lin et al.) 2020 Li et al.) 2023a). FedLWS directly calculates the value of the shrinking factors 7 without the need for proxy dataset fine-tuning step. The calculation of the shrinking factor y is based on the parameters and the gradient updates of the global model, which can be casily obtained by the server during the training process, without the need to transmit any additional information. Hence, our proposed FedLWS avoids data leakage and more applicable real-world scenarios.

---

Published as a conference paper at ICLR 2025
Table 1: Top-1 test accuracy (%) on four datasets with three different degrees of heterogeneity.
Datuset
FashimMNNSST
CIFAR-10 10
CIFAR- 100
Tiny ImageNet
Average
Heterogencity
a=100
e=0.5
o=0.1
c=100
o=0.5
c=00.
c=e100
co:0.5
a=0.1
a=100
o.00.5
a=0,1
FedLAW (Li al. 20236
89.78
89.25
87.52
81.30
75.27
64.76
41.05
37.56
34.59
37.20
33.49
29.13
58.41
FodAvg McMahan ol. 2017
90.44
90.04
58.62
76.01
74.47
61.04
41.46
37.21
36.71
36.31
20.40
+LWS (Oers)
90.99
90.33
88.99
76,85
75.63
64.08
42.42
41.03
37.70
37.16
35,12
31.34
59.30
FedDisco et 2025
89.90
88.
74.72
62.86
46
35.
79
29. 20.39
58.
+LWS (Oers)
90.95
90.43
89.39
76.34
75.41
66.91
42.68
41.68
36.93
37.30
34.92
31.53
59,54
L-DAWA (Rchman
2025
89:97
86.01
79.05
75.57
75.61
62.87
42.38
39.81
36.31
33.55
31.43
30.02
49.30
+LWS
91.05
86.63
86.11
76.21
76.77
65.13
42.97
40.40
37.04
36.83
33.77
32.34
51.19
FoilProx eL. 7020h
91.24
901.6
88, 78
73.96
73:27
60.62
18.15
39.35
34.60
35.03
34.33
29.37
57.45
+LWS (Ours)
91.35
91.24
89.25
74.55
38.64
35.37
35.29
34,98
30.68
58.18
FedDyn Acar al. 2021
90.29
88.57
87.84
77.92
74,61
56.37
11.01
41.80
36.92
34.32
32.80
27.32
57.73
+L.WS
90.69
77.33
46.14
33.20
27.74
Modularity. Our proposed FedLWWS can be a plug-and-play module in many existing FL methods to further improve their performance. It has a broad range of applications. For the FL methods adjusting the client-side model (Li et al.) 2020b] Acar al. 202T; Li et al., 2021). FedLWS operates on the server-side, therefore it can be easily integrated with these client-side adjustment methods. Moreover, for the methods adjusting the server-side model, our FedLWS is conducted after the model aggregation, and most previous works (McMahan et al.) 2017: Ye et al. 2023 Hsu et al., 2019) commonly normalize the aggregation weights (the sum of weight is 1), which is orthogonal to our method. Hence, our FedLWS can also be incorporated with them by conducting weight shrinking after model aggregation. While both ours and FedLAW et 2023a) consider the shrinking factor, FedLAW lacks flexibility in combination with other methods, as it learns not only the shrinking factor but also the aggregation weight, which is non-orthogonal with the previous work on model aggregation adjustment. Moreover, it requires a proxy dataset, necessity not present in ours. More discussions can be found in Appendix A
5 EXPERIMENTS
5.1
EXPERIMENT SETUP
Dataset and Baselines. In this paper, we consider four image classification datasets: CIFAR-10 10 Krizhevsky et al. 2009), CIFAR- 100 (Krizhevsky et al. 2009), FashionMNIST (Xiao et al. 2017
and Inagene (Chrabaszez et al. 2017); and text classification datasets: AG News (Zhang et al.
2013), Sogou News (Zhang et al. 2015). ) and Amazon Review (Ben-David et al.] (006)) We compare our method with six representative baselines. Among these, 1) FedAvg (McMahan et al., 2017) is the standard algorithm of Federated I Learning; 2) FedProx (Li et al., 2020b) and FedDyn (Acar al. (02I) focus on the adjustment of the local model; 3) Feddisco (Ye et al. 2023) and L-DAWA (Rehman et al. 2023) focus on the adjustment of the aggregation scheme. Our approach can be easily integrated with the methods mentioned above. Additionally, we also demonstrate the performance of FedLAW (Li et al. 2023a). However, since it leverages additional data for fine-tuning that other methods do not, we present it only for reference. In the experiments, our method FedLWS is the layer-wise approach described in Section 4.2 unless indicated specifically [
Table 2: Results on text classification datasets.
Table 3: Average aggregation execution time (Sec) across different model structures.
Method
With LWS?
AGG News
Sugan
Amazon Revlew
0.1.
Festure Shif
Method
CNN
ResNet20
WRN56_4
DemeNet1221
RodAvg
74.86
72.32
70.37
90.56
67.66
91.53
92.76
$5.15
$R.62
FedAug
0,019
0.10
0.18
0.561
1.359
FedProx
75.24
65.07
77.18
74.50
9017
50.00
92.24
93.10
88.24
58.75
FedLAW
FedLWS (Ours)
4,830
0.035
7.11
0.12
9.80
0.21
0.832
20.08
27.25
1.756
Federated Simulation. To emulate the FL scenario, we randomly partition the training dataset into K groups and assign group k to client k. Namely, each client has its local training dataset. We reserve the testing set on the server-side for evaluating the performance of global model. In practical FL scenarios, the clients often exhibit heterogeneity, leading to non-lID characteristics among their data. In this paper, employ Dirichlet sampling Dir, synthesize client heterogeneity, it is widely used in FL. literature (Wang et al. 2020a Vurochkin et al. 2019 Ye et al.) 2023). The smaller the value of a, the greater the IID. We apply the same data synthesis approach to all methods for fair comparison. More implementation details can be found in Appendix
'The source code is available at httpp:///githhbb.cm////gggggonop /
7

---

Published as a conference paper at ICLR 2025
5.2 PERFORMANCE EVALUATION
Performance and Modularity, In this section, we consider six representative methods and report the test accuracy on all datasets before and after applying our FedLWSS under different heterogeneity settings in Table ! It can be observed that the application of FedLWS leads to increased accuracies for all methods across various datasets and heterogeneity settings, highlighting the effectiveness of our proposed method. This is particularly inspiring because FedLWS necessitates modification to the original federated training process. The accuracy enhancements can easily attained by simply proscescsssing the aggregated global model. As shown in Table I our method performs better on more complex or heterogeneous data. This can be explained as follows: FedAvg already achieves strong results on relatively simple tasks datasets with III) distributions (Li et al.) [020b] Ye et al. 2023). In such scenarios, the differences between client models are relatively small, and consequently, the yy computed using Equation 7 tends to be closer to 1. This means that FedLWSSs behavior aligns more closely with the baseline under these conditions. However, on datasets with greater complexity or heterogeneity, client model differences become more pronounced. FedLWS effectively addresses these differences through layer-wist wise weight shrinking, resulting in improved global model performance. Furthermore, can be observed that FedLAW performs better on the CIFAR-10 dataset and, in some scenarios, even surpasses our method. One possible explanation is the proxy dataset used in our experiments, which contains 200 samples, consistent with the original settings in (Li et al.) 2023a). For CIFAR-10, this provides 20 samples per class, providing sufficient information to optimize the shrinking factor effectively. In contrast, for datasets like CIFAR-100 and TinylmageNet, the proxy dataset contains only 2 and I sample per class, respectively. This limited representation makes it difficult for FedLAW to train an optimal shrinking factor, which may explain the variations in its performance across different scenarios. To verify that our method can also be applied to text modality, we conducted experiments on NLP datasets under different heterogeneity settings. Table 2 shows that FedLWS still consistently improves the baselines on text modality.
Computation Efficiency. In Table 3 we show the aggregation execution time of FedAvg, FedLWS, and the closely related work FedLAW (Li al., 2023a) across various model architectures. For FedLAW AW, the proxy dataset contains 200 samples and the server epoch is 100 (consistent with (Li ct al. 2023ah). It can be observed that, in comparison to FedLAW, our FedLWS requires significantly less execution time, as FedLAW necessitates an additional proxy dataset for fine-tuning, which we do not require. Although larger models may slightly increase the computational load, our method remains markedly more efficient than optimization-baser alternatives, further demonstrating the efficiency of our approach.
Table 4: Comparison of layer wise and model- Table 5: The performance of compared methods with wise weight shrinking effects.
different model architectures.
Dataset
CIFAR-10
CIFAR-100
Model
ResNer20
WRN56_4
DemsNNt121
VII
Dataset
FedL.A
75.72
$0.46
51,20
Heterogeneity
6=100
o=0.1
6=100
g=0.1
FodAvg
75.07
78.97
86. 14
51.31
CIFAR-10
+LWS
76.17
81.29
86.63
54.14
HedAvg
76.01
61.04
41.46
36.71
Fedi -S(MMdee-wise))
76.17
63.21
41.77
37.65
FedLAW
36.53
36.60
55.36
22.03
FedAvg
36.71
39.71
56.59
25.60
CIBAR-100
Fedl aseea
76.85
64.08
42.42
37.70
+LWS
38,01
42.37
57.41
26.05
1.00
0.96
Fe W5/dc
PeddWwg
0.8
0.94
0.15
Dersit
0.92
0.00
Range layer-vise
Mazn of lsyar-wise
0.8
of the layer
0.05
of tayer
080
Morisk-wiise
0.00
2
0.08
0.06
0.04
0.02
0.00
0.02
0.04
0.06
0.08
Communication Rounds
Parameter weight
(a) The value of model-wise and layer-wise shrinking factors during the training process (o=0.1).
(b) The histogram of final models' parameters.
Figure 3: The comparison between layer-wise wise FedLWS and model-wise FedLWSS (ResNet20).
8

---

Published as a conference paper at ICLR 2025
Acouracy improvement with LWS(()
Accursty impoveeeettt LMS(54)
Accuracy improvementtom LWS(S9)
ond
Fallymoss
Harranse LORDONA
Lood cauch
Purtopaion Tudo
(a) Client Number.
(b) Local Epoch.
(c) Partial Participation.
Figure 4: Accuracy improvement after applying LWS under different client numbers, local epochs,
and partial participation ratios.
5.3 ABLATION STUDIES
Effects of layer-wise shrinking factor. In this section, we conduct experiments to analyze the effectiveness of the layer-wise shrinking factor. As a point of comparison, we implement the model- wise FedLWS by computing a model-wise shrinking factor y for the entire model, which can be viewed as degraded version of layer-wise FedLWSS. The model-wise shrinking factor is calculated according to Equation y based on the entire model's 's parameters and gradient. Table 4 taking the FedAvg the baseline, we compare the performance of layer-wise and model-wise weight shrinking across different datasets and heterogeneity degrees. Table 4 demonstrates that layer-wise weight shrinking achieves higher model accuracy compared to model-wisg weight shrinking, indicating that the layer-wiss approach allows for more precise adjustments the model aggregation process, thereby yielding improved results. In addition, even introducing the degraded model-wiss shrinking factor into FedAvg can still achieve performance improvement, suggesting the effect of shrinking factor on the global model's generalization ability.
In Figure ([a] we compare the values of model-wise 7 and layer-wise yy during the training process in Resnet20, where we show the mean of layer-wise 7 of all layers, 7 at the first layer and ? the last layer due to the limited space. The variation of layer-wise for each layer is deferred to Figure P in Appendix. It can be observed that throughout the training process, the value of model-wise y and the mean of layer-wise "y are very close. Besides, the variance of the layer-wise Y gradually decreases with the training process. That to say, at the initial iterations, there is a significant difference in layer-wise y among different layers of the global model. As the training progresses, the gap between layers continues to diminish. This indicates that the difference between layer-wise wise y and model-wise y primarily occurs at the initial stages of training. Layer- at the initial training period. It can be inferred that layer-wise y can enhance better utilization of the global weight shrinking phenomenon and improve model generalization by assigning optimal y values for each layer of the global model. We show the histogram of the final models' parameters in Figure ((b] It can be seen that, in comparison FedAvg and model-wise FedLWS, layer-wise FedLWS makes more model parameters close to zero, which is similar to weight decay. This may explain why FedLWS enhances the generalization ability of the global model. More results about shrinking factor can be found in Appendix B
better adjust the global model
Effects of model architectures. In Table we evaluate our proposed FedLWS across various model architectures, including ResNet (He et al. 2016). Wide-ResNet (WRN) Zagoruyko & Komodakis 2016), DenseNet (Huang al. (017) and Vision Transformer (ViT) (Dosovitskiy 2020)). The results demonstrate the effectiveness of FedLWS across different model architectures, indicating its robust performance even as the network depth or width increases.
Effects of client number, local epoch, and partial participation. In this section, we tune three crucial parameters in FI.: the number of clients E (10, 30, 50), the number of local epoch E e (1, 10), and partial participation ratio R E 10.1, 0.3, 0.5]. We show the accuracy improvement brought by FedLWSS in Figure (a) 4(b) and [c)) respectively. The experiments consistently reveal that our proposed method consistently brings performance improvement across different FL settings.
Relation with Weight Decay. While our method differs from weight decay, it is still meaningful to compare it with adaptive weight decay methods. Therefore, we applied two adaptive weight decay methods, AWD (Ghiasi et al., (023) and AdaDecay (Nakamura & Hong 2019). to Federated Learning model aggregation process, and compared them with our approach, the result shown in

---

Published as a conference paper at ICLR 2025
Table 6: Top-1 I test accuracy (%) on CIFAR-10 and CIFAR-100 with different degrees of heterogeneity. Compared with weight decay methods.
Dataset
CIFAR-10
CIFAR-100
Average
Helergeneity
ID(r=100)
NID(o=1)
NHD(os0))
IiX==100)
NID(o=
NID
FedAug
76.01
75.18
61.04
41.46
41.62
36.71
55.34
Deaaeeay (Nakomura Hong 2019
76.21 (10.20)
(10.15)
61.34 (10.30)
42.15 (10.69)
(t0.1)
36.98
55.62 (10.28)
Ghiasi
76.30 (10.29)
75.64
(10.11)
42.8 ((135)
(1..6()
37.17 (10.46)
(10.4))
(10.54)
75.88 (10.70)
(*3.4)
42.42
(11.3(
(10.99)
56.64 (1136)
Difforences
Differences between local gradient
11.
11-
1.75
65
Cilent ID
8
125
Accufal
4.75
Fedfeg
50
Fedhhrrtatt
425
FedeegiuMS
FedDlsco-UWS
30
101337
10
10"
10:2
10a
Client ID
Client
Figure 5: Divergence local data and gradients
Figure 6: Accuracy with different |
Table It demonstrates that method significantly outperforms the adaptive weight decay methods in the Federated Learning scenario. More discussions can be found in Appendix A.2
Relationship between local dataset and local gradient. Our method the shrinking factor. To analyze the reasons for its effectiveness, we conducted additional experiments for further observation. Figure 3 we illustrate the divergence between different local datasets and between the correspondinglocal gradients (i.e., local model's 's parameter changes after local training). As shown, the distances local gradients closely resemble those of the local data distributions. This demonstrates that the local gradient can effectively capture relevant information about the local data. In other words, our calculation of the local gradient variance, т, also serves as measure the degree of data heterogeneity. For more heterogeneous data, stronger regularization is applied, making the training process more stable. More implementation details can be found in Appendix B.3
local gradients calculate
Hyperparameter. In this section, we explore the impact of varying the hyperparameter 3 to highlight the flexibility and adaptability of hyperparameter tuning in our proposed FedLWS. In Figure 6 we demonstrate the impact of different / values accuracy when FedLW'S is combined with various methods. As /3 approaches O, the ? value calculated using Equation 7 converges to 1, resulting in the model's performance degrading to that of the baseline. Conversely, when is too large, the calculated y becomes excessively small, which can cause model instability or even failure. Based on our experiments, we recommend safe range for & between 0.001 and 0.1.
CONCLUSION
In this paper, through empirical explorations, we show that layer-wise weight shrinking can further improve the generalization of the global model. Subsequently, we investigate the factors influencing the shrinking factor during the training process of Federated Learning. Based on these observations, we present FedLWS, method that directly computes layer-wis, shrinking factors without requiring any additional data. FedLWWS can be seamlessly integrated with the existing Federated Learning methods to further enhance the performance of the global model. Experiments demonstrate that FedLWS steadily enhances the state-of-the-an FL approaches under various settings. A potential limitation of our method that it is only applicable to scenarios where the client model architectures identical. The heterogeneity in client model architecture is common issue in FL. and this limitation is prevalent in many FL methods. We plan to address the application in scenarios with heterogeneous client models in future work.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We truly thank the reviewers for their great effort our submission. Changlong Shi, Jinmeng Li, Dandan Guo and Yi Chang are supported by the National Natural Science Foundation of China (No. U2341229, No. 62306125) and the National Key R&D Program of China under Grant (No. 2023YFF0905400).
REFERENCES
Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama. Federated learning based on dynamic regularization. In International Conference on Learning Representations, 2021. URL URL
hww.:/ee//e....... forum?
#d=B7v4QMR629w
David Alvarez-Melis and Nicolo Fusi. Geometric dataset distances via optimal transport. Advances in Neural Information Processing Systems, 33:21428-21439, 2020.
Shai Ben-David, John Blitzer, Koby Crammer, and Fernando Pereira. Analysis of representations for domain adaptation. Advances in neural information processing systems, 19, 2006.
Hong-You You Chen and Wei-Lun Chao. Fedbe: Making bayesian model ensemble applicable to
federated learning. In International Conference on Learning Reprexentations, 2021. URL https
//openreview net hww.://www.eeeettneeeeeeeeeeeaee pgt
Patryk Chrabaszez, Ilya Loshchilov, and Frank Hutter. A downsampled variant of imagenet as an alternative to the cifar datasets. arXiv preprint arXiv:1707.088: 08819, 2017.
Israel Cohen, Yiteng Huang. Jingdong Chen, Jacob Benesty, Jacob Benesty, Jingdong Chen, Yiteng Huang, and Israel Cohen. Pearson correlation coefficient. Noise reduction in speech processing, pp. 1-4, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations, 2020.
Jintong Gao, He Zhao, Zhuo Li, and Dan dan Guo. Enhancing minority classes by mixing: An adaptative optimal transport approach for long-tailed classification. Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL
www. forum?
1d Cd-M7FQpIdoOX
Liang Gao, Huazhu Fu, Li, Yingwen Chen, Ming Xu, and Cheng-Zhong Xu. Fedde: Federated learning with non-iid data via local drift decoupling and correction. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10112-10121, 2022.
Amin Ghiasi, Ali Shafahi, and Reza Ardekani. Improving robustness with adaptive weight decay. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non- non-identical data distribution for federated visual classification. arXiv preprint arXiv:1909.06330 2019.
Gao Huang, Zhuang Liu, Laurens van der Maaten, and Kilian Q. Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition (CVPR), July 2017.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Ad- vances and open problems in federated learning. Foundations and Trends@ in Machine Learning. 14(1-2):1-210, 2021.
11

---

Published as a conference paper at ICLR 2025
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.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Sunwoo Lee, Tuo Zhang, and A Salman Avestimehr. Layer-wise adaptive model aggregation for scalable federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 8491-8499, 2023.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018.
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, 2021.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE signal processing magazine, 37(3):50-66) 60, 2020a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429 450, 2020b.
Zexi Li, Tao Lin, Xinyi Shang, and Chao Wu. Revisiting weighted aggregation in federated learning with neural networks. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023a.
Zexi Li, Xinyi Shang, Rui He, Tao Lin, and Chao Wu. No fear of classifier biases: Neural collapse inspired federated learning with synthetic and fixed classifier. In Proceedings of the IEEF/CVF International Conference Computer Vision (ICCV), pp. 5319-5329, October 2023b.
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-2366 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:5972-5984, 2021.
Xiaosong Ma, Jie Zhang. Song Guo, and Wenchao Xu. Layer-wised model aggregation for person- alized federated learning. Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 10092-10101, 2022.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efiini
learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pp. 1282. PMLR, 2017.
Kensuke Nakamura and Byung-War Woo Hong. Adaptive weight decay for deep neural networks. JEEE Access, 7:118857-11888 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-peeformance deep learning library. Advances in neural information processing systems, 32, 2019.
Pian Qi, Diletta Chiaro, Antonella Guzzo, Michele lanni, Giancarlo Fortino, and Francesco Piccialli. Model aggregation techniques in federated learning: A comprehensive survey. Future Generation Computer Systems, 2023.
Yasar Abbas Ur Rehman, Yan Gao, Pedro Porto Buarque de Gusmao, Mina Alibeigi, Jiajun Shen, and Nicholas D. Lane. L-dawa: Layer-wise divergence aware weight aggregation in federated sulf-supe visual representation learning. In Proceedings of the IEEE/CVF International Conference Computer Vision (ICCV), pp. 16464-16473, October 2023.
12

---

Published as a conference paper at ICLR 2025
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Fed- erated learning with matched averaging. In International Conference on Learning Representations, 2020a.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33:7611-7622 2020b.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017.
Hangting Ye, Wei Fan, Xiaozhuang Song, Shun Zheng, He Zhao, Dan dan Guo, and Yi Chang. Ptarl: Prototype-based tabular representation learning via space calibration. In The Tweifth International Conference Learning Representations, 2024.
Rui Ye, Mingkai Xu, Jianyu Wang, Chenxin Xu, Siheng Chen, and Yanfeng Wang. Feddisco: Federated learning with discrepancy-aware collaboration. In Proceedings of the 40th International Conference on Machine Learning. ICML'2 23. JoLR.org. 2023.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In International conference on machine learning, pp. 7252-7261. PMLR, 2019.
Sergey Zagoruyko and Nikos Komodakis. Wide Residual Networks. In British Machine Vision Conference 2016, York, France, January 2016. British Machine Vision Association. doi: 10.5244/ C.30.87. URL https::////.. //es:////ppppc hal
Chen Zhang, Yu Xie, Hang Bai, Bin Yu, Weihong Li, and Yuan Gao. A survey on federated learning. Knowledge-Boud Bused Systems, 216: 06775, 2021.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. Advances in neural information processing systems, 28, 2015.
Ye Zhang and Byron Wallace. A sensitivity analysis of (and practitioners' guide to) convolutional neural networks for sentence classification. arXiv preprint 1rXiv:1510.03820, 2015.
13

---

Published as a conference paper at ICLR 2025
A MORE DISCUSSION
A.1 MODULARITY.
Our proposed FedLWS can be easily incorporated with many existing FL methods to further improve their performance. For the FL methods adjusting the client-side model (Li et al. 2020b] Acarr et al. 2021 Li al. (02I), which corresponds to the line 6 in Algorithm|| [ since
proposed FedLWS operates on the server-side with lines 11 and 12 in Algorithm II it can be easily integrated with these methods. Moreover, for the methods that conduct the server-side model adjustment, most previous works (McMahan et al.) 2017 Ye et al.] [023] Hsu et al. 2019) commonly normalize the aggregation weights (the sum of weight is 1), which is orthogonal our method. Hence, our FedLWS can also be incorporated with them by conducting weight shrinking after model aggregation, further enhancing the model's generalization performance. For instance. regardless of whether the server utilizes FedAvg (McMahan et al., (017) or FedDisco (Ye et al. (023) for model aggregation (i.e., achieving the global model W "/ like line 10 in Algorithm we can further strengthen the aggregation process by implementing the lines 11 and 12 in Algorthhm
A.2 RELATION WITH WEIGHT DECAY.
In T AWD and AdaDecay were not originally designed for the federated learning (FL) context. In our work, we adapted these methods to make them applicable to FL settings for a fair comparison with FedLWS. However, these methods do not account for the differences between local models, which are critical in heterogeneous FL scenarios. In contrast, FedLWS explicitly addresses these differences by considering the variance of client update gradients, combined with its layer-wise adaptability, leading to better performance in FL environments. Although the layer-wise weight shrinking (LWS) effect is analogous to weight decay, these two methods are distinet. 1) LWS employs a distinctive sparse regularization frequency, modifying model weights only in each round, resulting in stronger regularization. In LWS. - y is near 0.05, which is significantly larger than the value of the weight decay (typically 10-9), 2) LWS shrinks the global model rather than decaying the model by subtracting a decay term. 3) In this paper, LWS is conducted on the server side to adjust the aggregated model. At this stage, the model is not trainable. We utilize parameter variations as to compute the shrinking factor. In contrast, weight decay is applied to the local model training process on the client side. Consequently, the two methods are not confficted in FL..
A.3 RELATION WITH FEDLAW.
Learning aggregation weights of local models is an effective solution to improve Federated Learning, Recently, FedLAW (Li et al., 2023a) has gradually attracted the attention of the researchers, which identifies the global weight shrinking phenomenon and then learns the optimal shrinking factor y and the aggregation weights À on the server. Despite the effectiveness of FedLAW, it needs to optimize the y and À at the server with additional proxy dataset, which is assumed to have the same distribution as the test dataset. We have shown in Table 7 that only change the class distribution of the proxy datasct can reduce the performance of FedLAW lot, where proxy dataset and test dataset are still from the same dataset but with different class distributions. Considering data privacy is significant concern in Federated Learning, obtaining the proxy dataset with distribution identical to the test dataset in practice is challenging, limiting its application in real-world. The original FedLaw paper also conducted similar experiments on CIFAR-I0 with 1=100 and 0. but compared to theirs, our experiments are more comprehensive, covering both CIFAR-10 10 and CIFAR- 100 datasets with a values of 100 and 1. the latter indicating a higher degree of heterogeneity. Consequently, when comparing our results in Table Z with those in Table 6 of the FedLAW paper, only the scenario with with
100 is directly comparable. As observed, when a 100, the performance drop in the original paper is 2.26% (from 79.40 to 77.14), while in our corresponding scenario, the drop is 4.61% % (from 81.30 to 76.69). The performance after applying long-tailed proxy data is similar across both studies; however, the results in our experiment outperform those in the original paper when using balanced proxy data. This could be attributed to the higher quality of our randomly sampled balanced proxy data. This further highlights that the effectiveness of the FedLAW method is highly dependent proxy data quality. Besides, FedLAW learns the shrinking factor y and the aggregation weights at the server jointly, making it difficult to combine with other aggregation methods in FL. Last but
14

---

Published as a conference paper at ICLR 2025
not least, FedLAW ignores the variations across different layers of model for model aggregation. problems is quite
Therefore, designing an effective and flexible method to solve necessary for model aggregation in FL.
To this end, we propose novel model aggregation strategy, Federated Learning with Layer-wise Weight Shrinking (FedLWSS). It is non-trivial since deduce the expression of the shrinking factor and can calculate it directly through the easily available gradient and parameters the global model, which avoids demanding proxy dataset and optimization. Considering the distinctions among layers within deep models, we further design the shrinking factor in a layer-wise manner, which is feasible and effective due the deduced expression. Besides, ours can be easily integrated with most of related model aggregation methods for decoupling shrinking factor and aggregation weights. Therefore, ours is not just applying FedL each layer of the global model. Even we only consider single shrinking factor y for all layers, ours is not equal to FedLAW. We have provided our degraded model-wise version (shared shrinking factor across all layers) in Table [8] evident that model-wise version still performs betters than FedLAW, proving the difference between FedLAW and ours. Furthermore, since we do not need to optimize the shrinking factor on the server, ours has less computational cost than FedLAW; As shown in Table of the main text, the computational time required by FedLAW is nearly 60 times that our method. These observations and results show that our proposed method can without demanding proxy dataset and additional computational cost.
effective plug-and-play module in many existing FL methods,
To summarize, FedLAW is the state-of-the-an method in the line of the optimal shrinking factor for model aggregation in FL. We propose a new method in this line that improves over FedLAW via nontrivial efforts, which we believe has significant contributions.
Table 7: Performance in scenarios where proxy data are long-tailed and test data are balanced. The numbers in parentheses indicate the decrease compared to using proxy data with the same distribution as the test data.
Dataset
Cifar 0.1)
Cifarl0(o 100)
Cifar 0.1)
Cifar100(a 100)
Avg
FedLAW (LT) FedAAv+LWWS (Ours)
36.91()7885)
76.69))
27.57(17.02)
36.58((4.47)
49.44(15.9 55.26
64.08
76.85
37.70
42.42
Table 8: Evaluation of model-wise methods (a 0.1).
Dataset
Fmnist
Cifar100
Cifar10
TinyImageNet
Avg
FedLAW (Li et al. 2023a
87.52
34.59
64.76
29.13
54.00
Ours (Model-wise)
87.95
36.41
65.31
30.39
55.02
Ours (Layer-wise)
89.39
36.93
66.91
31.53
56.19
B EXTRA EXPERIMENTAL RESULTS
B. EXPERIMENTS ON PRETRAINED MODEL.
In this section, conducted experiments using a pretrained ResNet20 on various datasets and degrees of data heterogeneity. The experimen- tal results as shown in Table 9 It can be seen that when using a pretrained model for initial- ization, FedLWS can still further enhance the performance of the global model. This demon- strates the effectiveness of our method. When the global model initialized using a pretrained model, the shrinking factor of FedLWS is close to 1 (around 0.99) during the training process. This be attributed to the superior performance of the pretrained model, which requires only minimal adjustments during the training process. As result, there are smaller gradient updates and less necessity for strong regularization. Therefore, when the global model is initialized using
Table 9: The performance of FedLWS using pretrained ResNet20 model for initialization.
Dataset
CIFAR-10
CIFAR-
Heterogeneity
0.1
0.1
FedAvg
86.74
91.30
63.98
66.77
+LWS (Ours)
87.02
91.47
64.14
67.06
15

---

Published as a conference paper at ICLR 2025
1w2
4.0 302
162
24
Pearson:: 0.839
2,00
3.2
20
10
H
1.50
1.25
1.5
0.1 0.2
0.3
0.4
055
D4
0.7
0.8
0.
3
5
10
300
0.1
0.2
0.3
044
05
06
0.7
0.8
0.9
10 100
(a) ResNet20, CIFAR-10.
ResNet20, CIFAR-100.
163
1=2
1x2
Paarson:: 0.807
4.5
52
16
5.0
3.0
4.8
1.5
4.6
0.0
100
0.1
0.4
055
100
(c) ResNet18, Tiny-ImageNer
(d) VIT, CIFAR1O.
Figure 7: Gradient variance and ratio under different degrees of data heterogeneity.
pretrained model, our method exhibits limited influence on improving the global model. This finding is consistent with the insights presented in our paper.
B.2 RELATION BETWEEN THE GRADIENT VARIANCE AND THE BALANCE RATIO
To further validate our hypothesis, in addition to the experiments shown in Figure I(a) we conducted additional experiments on different datasets and model architectures. The results, as shown in Figure 7 indicate positive correlation between gradient variance and the balance ratio across different model architectures and datasets, further supporting our hypothesis. can observed that both and exhibit corresponding reduction the degree of data heterogeneity diminishes. It is also consistent with our view that gradient variance is relative with the balance ration r between regularization term and optimization term. Therefore, it inspires to assume the relationship between the unknown balance ration with the easily available gradient variance T.
Moreover, there is sudden decline of between 0.6 and 0.7 in Figure I(a] which was not observed in other datasets or model architectures. This suggests that the behavior observed in Figure I is likely specific to the characteristics the dataset or the model architecture used in that particular experiment. Additionally, considering that deep learning models can exhibit variability due to stochastic elements (e.g., initialization, sampling, or optimization). this decline might also be attributed random fluctuation specific to this experiment.
B.3 EXPERIMENT DETAIL.
In the experiment of Figure 5 we set up 12 clients, and the dataset we used was CIFAR-10. To observe the data differences more intuitively, we introduced extreme scenario of data heterogeneity: The local data of clients to contained only the first 5 classes of the CIFAR- 10, clients had only the left classes, and clients 12 had local datasets that included all classes. The size
16

---

Published as a conference paper at ICLR 2025
Table 10: Performance comparison when only adjust learning weight decay.
Dataset
Cifar10(a 0.1)
Cifar10(a 100)
Cifar100(a 0.1)
Cifar100(a 100)
Avg
WrdLWS/onll-ree
61.54
74.92
34.39
38.89
52.44
W/dnlW//o
62.28
74.72
33.74
39,21
52.49
FedLWS
64.08
76.85
37.70
42.42
55.26
of the local dataset for each client was the same. We calculated the distances between the clients' local dataset via optimal transport, a methodology that has garnered widespread application within
the realm of machine learning (Alvarez-Melis & Fusi 2020 Gao et al., 2023 Ye et al. 2024), and
the results are displayed in the left of Figure 5J Then, we conducted Federated Learning process, during which each client obtained its own local gradient after local training. We then compared the cosine distance between the local gradients, with the results displayed on the right of Figure 5 can be that the relationships between the client vectors closely resemble those of the local data distributions. For example, client I's client vector exhibits minimal differences with clients 2-4 due to their similar local data distributions. However, the differences between client and clients 5-8 are much larger because of their highly divergent data distributions: client I's local data contains only the first 5 classes, while clients 5-8 have only the last 5 classes. The differences between client 1 and clients 9-12 are smaller than those with clients 5-8, as clients 9-12 include data from all classes, making their distribution relatively closer client 1. This demonstrates that the local gradient can effectively capture relevant information about the local data.
B.4 ONLY ADJUST A SINGLE TERM.
As can be seen from the Equation 3 y influences both the optimization term and the regularization term. Our method calculates an appropriate value y balance the regularization and optimization terms during the training process. In other words, our method simultancously considers both aspects to determine the optimal value of y. investigate the impact of our method on a single item, we conducted experiments as shown in Table TO comparing two different variants: FedLW WS/nnlyyopppp and S/onyS/on,y.reg. FedLWS/only-opp indicates adjusting only the optimization term during training, and TedLWS/only-rep indicates adjusting only the regularization term. The results indicate that adjusting only one of the terms does not perform well. This proves that our method effectively improves the model's 's performance by balancing the regularization and optimization terms during the training process.
B.5 MORE EXPERIMENT OF LAYER-WISE WEIGHT SHRINKING
B.5.1 MODEL PARAMETERS
In Figure S we present the histogram of the final models' parameters when using CNN as the model architecture. Similar phenomena to those observed Figure of the main text be observed. Layer-wise FedLWWS drives a larger number of model parameters towards zero, thereby achieving an effect akin to weight decay and enhancing the model's generalization ability.
B.5.2 SHRINKING FACTOR FOR EACH LAYER.
In Figure 9 we illustrate the layer- shrinking factors calculated by our method for each layer during the training process across various model architectures. It can be observed that during the initial stages of training, there is a significant difference among the layers of the model, with the classifier exhibiting more pronounced distinction compared to other layers. Previous works (Luo) ct al.) 2021 Li et al. 2023b)) have demonstrated the specificity of the classifier Federated Learning. Our observations in Figure 9 further validate this point. As the training progresses, the differences in the shrinking factors between layers gradually diminish. This indicates that our method primarily adjusts the aggregation process of the model during the early stages of training. As the training converges, the differences between clients gradually diminish, resulting in the calculated shrinking factors approaching 1.
17

---

Published as a conference paper at ICLR 2025
0.35
Fe FMd
FedAvg
0.25
Density
0.20
0.15
0.10
0.00
-0.04
-0.03
-0.02
-0.01
0.00
0.01
0.02
0.03
0.04
Parameter weight
Figure 8: The histogram of final models' parameters. (CNN)
B.5.3 COMPARE WITH FIXED LAYER-WISS WISE SHRINKING FACTOR.
In Table I1 taking the FedAvg and FedDisco the baselines, compare the performance of our FedLWS and the method with fixed layer-wise 7 (the method in Figure ([b) under different datasets and degrees of heterogeneity. Fixed layer-wise sets fixed shrinking factor for each layer, and its value remains constant throughout the training process. Given that the CNN model comprises five layers, set the 7 values between each consecutive layer to differ by 0.01. Therefore, in fixed layer-wise the of each layer is uniformly decreasing from to 0.96. From Table T1] it can be observed that the performance of fixed layer-wise y shows significant improvement in some scenarios, even surpassing FedLWS, indicating that considering weight shrinking in a layer-wise manner can indeed enhance the model's performance. However, its performance significant deterioration is observed in certain situations (CIFAR-100 100, NIID(o=0.1) Therefore, compared to fixed layer- wise y, our approach is more robust. By dynamically adjusting the layer-wise shrinking factors in different scenarios, our method achieves better overall performance.
stable, and
Table 11:
Compare the performance of FedLWS and fixed layer-wise shrinking factor in both IID(o=100) and NIID(a=0.1) settings CIFAR-10 and CIFAR-100.
Dataset
CIFAR-10
CIFAR-100 100
Heterogeneity
ITD(a=100)
NID((=00.1)
IT((==100)
NIID(==0.1)
Mean Improvement
Model
CNN
CNN
CNN
CNN
FedAvg
70.19
59.66
32.25
30.18
Fidevg+Fixed layer-wise
73.86(1 3.67)
61.47(1 1.81)
36.47(1 4.22)
26.14() 4.04)
1.415
FodAvg+FodLLWS
71.79(11.60)
62.73(13.07)
33.24(10.99)
32.51(12.33)
1.998
FedDisco
70.66
61.78
32.61
30.28
FedDisco+Fixed layer-wise
73.83(1 3.17)
62.37(t 0.59)
36.48(1 3.87)
25. 25.73(L 4.55)
0.770
FedDisco+FedLWS
71.36(10.70)
64.79(13.01)
32.97(10.61)
32.46(12.18)
1.625
B.6 DIFFERENT LAYER TYPE.
In this section, we investigate the differences in shrinking factors calculated using our method across different types of layers. To ensure a fair and accurate comparison, employed the MI.P,
18

---

Published as a conference paper at ICLR 2025
ResNet20, and ViT architectures for federated learning the CIFAR-10 dataset. In the Table below, we presented the average shrinking factors for each layer type, where 5 is the mean of layer-wise shrinking factors, e.g., MLP has 3 layers, corresponding 3 layer-wise shrinking factors " (m taa)) ViT(Att.) represents the attention layers in VIT, and ViT(MLP)) represents the MLP
y in VIT.
3
Table 12: Average shrinking factor for different types of layers.
MLP
ResNet20 (Conv.)
ViT (Att.)
ViT (MLP)
7 0.980 + 0.011 0.920 + 0.026 0.995 + 0.003 0.950 + 0.037
can be observed that the y values obtained for different types of layers vary significantly. This also demonstrates that our method can calculate the corresponding shrinking factors for different layer types. Notably, the shrinking factors for (iT(A layers are closer to and exhibit smaller differences across layers (low variance), indicating that weaker regularization is required. This can be attributed to the extensive parameter size of these layers, which minimizes the impact of gradient changes. Furthermore. in both MLP and ViT(MLP), the shrinking factor of the last layer is smaller than that of the other layers, e.g., the shrinking factors for the three MLP layers are 0.988, 0.984, and 0.968 respectively. This trend can be attributed the fact that the gradient changes in the last layer of MLP are greater than those in the preceding layers, akin to the phenomenon of gradient vanishing. Consequently, the last layer requires stronger regularization (smaller shrinking factor). The experiments indicate that our method calculates the corresponding shrinking factor for different layer types, allowing for a more refined adjustment of the model aggregation process.
B.7 TRAINING PROCESS
To evaluate the convergence of the proposed method, we examined the variation in test accuracy over rounds across multiple datasets and experimental configurations. Figure 10 illustrates the accuracy curves under various settings, including different datasets, heterogeneity degree a, model architectures, client numbers, selection ratios, and local epochs E. As shown in Figure [0] our method has similar convergence speed to FedAvg, the accuracy consistently increases with the number of rounds across all datasets, eventually reaching stable plateau. This trend demonstrates the robustness and convergence of the proposed method under these configurations. In most cases, the accuracy exhibits rapid improvement during the initial stages of training. followed by a gradual stabilization as the model approaches convergence. Furthermore, the results reveal several insights into the impact of different configurations on convergence speed and final performance:
Impact of a: Larger values of result in smoother optimization processes, whereas smaller values may lead to more oscillations in accuracy during training. However, the method ultimately converges in both cases.
Selection Ratios: A smaller selection ratio results in slower improvement and more oscillations in training accuracy. Nonetheless, the overall trend remains stable, indicating the method's adaptability to varying levels of client participation.
Local Training Epochs: Increasing the number of local training epochs significantly accelerates global convergence, highlighting the importance of local updates in enhancing global optimization efficiency.
These observations collectively demonstrate the convergence properties of the proposed method under diverse experimental settings. The consistent upward trend in training accuracy and eventual stabilization across all scenarios confirm the effectiveness and robustness of the method in federated learning.
C IMPLEMENTATION DETAILS
In this section. provide details of the pseudo-code, experiments details, environment, datasets, model architectures, and hyperparameters the experiment. We include the source code for implementing FedLWW in the Supplementary Material.
19

---

15: end for Published as a conference paper at ICLR 2025
Algorithm Fe:LWS: Federated Learning with Adaptive Layer-wise Weight Shrinking
Input: Communication round T, local epoch E, local datasets [D1, Dx J. initial global model Wy w gt o .... w o "l. where wa) is the L-th layer of wg. combine with FL method M; g2
2: Output: Final global model w
3: for = to 7 do
4: Server sends global model w/ to each client;
# Clients execute:
5: for each client [K] do
6: If M = FedAvg then: w + ClientUpdate DR. E);
7: If M FedProx then: w ClientUpaatewith(Mm Dk. E)# Use the Loss of FedProx
8: Send w server;
9:
end for
# Server executes:
10:
Server aggregates the received model to generate the global model:
11:
WH1 Σ Aaw' A
Σim D,
D.
# If M adjusts aggregation weights, A is
calculated using M.
Additional Step of our FedLWS
12:
Computes the value of gig whi w
gi and o/
|
gmeanill
||;
Computes the value of of 3r 3r fl s + II
I
Server conducts layer-wise weight shrinking obtain the updated global model (
w++1 ( (W[+)) (2 92
C.1 PSEUDO-CODE
To facilitate a clearer understanding our methodology, we present the pseudo-code of FedLWS in Algorithm thereby providing more intuitive representation of our method.
C.2 ENVIRONMENT
We conduct experiments under Python 3.7.16 and Pytorch 1.13.1 (Paszke et al. 2019).
C.3 DATASETS
In the experiment, we utilized four image classification datasets: CIFAR-I0 (Krizhevsky et al.) 2009), CIFAR-101 (Krizhevsky et al. 2009), FashionMNIST (Xiao et al. 2017), and Tmny-ImageNer [Chrabaszez et al. 2017). which have been widely employed in prior Federated Learning methods (Li al. 2023a Ye 2023 Luo et al. 2021). All these datasets readily available for download online. To generate a non-lll) data partition among clients, we employed Dirichlet distribution sampling Dira in the training set of each dataset, the smaller the value of a, the greater the IID. In our implementation, apart from clients having different class distributions, clients also have different dataset sizes, which believe reflects more realistic partition in practical scenarios. We set a =0.1, 0.5, and 100, respectively. When a is set to 100, we consider the data to be distributed in an IID manner. The data distribution across categories and clients is illustrated in Figure [T] Due to the large number of categories, we did not display the data distribution of CIFAR-100 and Imag-em Their distributions are similar to the other two datasets.
C.4 HYPERPARAMETERS
If not mentioned otherwise, The number of clients, participation ratio, and local epoch set to 20, 1, and I, respectively. We set / = 0.1 for CNN models and 3 = 0.01 for ResNet models. We set the initial learning rates 0.08 and set decaying LRR scheduler in all experiments; that is, in each round,
20

---

Published as a conference paper at ICLR 2025
the local learning rate is (the learning rate of the last round). We adopt local weight decay in all experiments. We set the weight decay factor as 5e-4. We use SGD optimizer as the clients' local optimizer and set momentum as 0.9.
C.5 MODELS
For each dataset, all methods are evaluated with the same model architectures for fair comparison. In Table annd We use simple CNN for FashionMNIST, ResNet20 (He et al. 2016 for CIFAR-10 and CIFAR-100, ResNet18 for Tmny-ImageNet, and TexICNN (Zhang & Wallace, 2015) for AG News.
SimpleCNN. The SimpleCNN is a convolution neural network model with ReLU activations. In this paper CNN consists of convolutional layers followed by fully connected layers. The first convolutional layer is of size (3, 32, 3) followed by a max pooling layer of size (2, 2). The second and third convolutional layers are sizes (32, 64, 3) and (64, 64, 3), respectively. The last two connected layers are of sizes (644494, 64) and (64, num classes), respectively.
ResNet, WRN, DenseNet and VIT. We followed the model architectures used in (Li et al. 2023a Dosovitskiy al. 2020 Li et al. 2018). The numbers of the model names mean the number of layers of the models. Naturally, the larger number indicates a deeper network. For the Wee5-4es (WRN56 _) in Table 5 "4" refers to four times many filters per layer.
D THEORY ANALYSIS
In the following analysis, we omit to assist clarity. In Federated Learning, the global model is trained by aggregating client updates. The generalization gap is defined as:
g Be-D(((w;))]
Σ B D. [L(w;:)]]]
K
where:
(lowD[((w:ed) is the true risk (expected loss over the global distribution).
D [L(w;; x) is the empirical risk (average loss over participating clients). Lemma D.1. Assume that the loss function L(w) is Lipschitz-smoott and differentiable and its higher-order terms in the Taylor expansion are negligible. Then, the variance of client losses is proportional to the variance of client gradients:
[[[[[ww; Dz) - ((w:D)) x B[[ggg Ameen l/ [2].
Proof. We first approximate the client loss L(Wa; DR) using first-order Taylor expansion around the global model w:
L( D2) ~ L(w; D) + Vw (w ww),
where wy is the locally updated model on client kt and w - Wk x gk. We have:
(L(wa: De) L(w; D))² (VwL(w:D) (
Let us denote dy gh Imean and then we have:
(L(w: Da) - L(w; 2)w; (W:
Taking the expectation over all clients, we derive:
E[[((w; Da) - L D)) x Ellgg - 9mean 2/
21

---

Published as a conference paper at ICLR 2025
Theorem D.2. The following upper bound of the generalization gap exists:
g.C.
V K
(10)
Proof. Using Bernstein's 's inequality, the probability of G is bounded as:
K
Pr 2 () 2exp
2( Mc/3)
where o² - + Σx= of is the variance of client losses and M is the maximum deviation of individual client losses.
K
The expected generalization gap is bounded by integrating over e:
E / Pr([(((w)- RR(w) 2 E) de,
which can be further written
E v 20 K
Given Lemma D.I and r k Σ 11gx Imean the variance of client losses be approximated by the variance of client gradients, i.e., o² = Therefore, have:
2T
B S
K
The above variance-based generalization bound provides theoretical foundation for managing client heterogeneity FL. It suggests that when T is large, meaning that the client gradient variance is high, which also indicates that there is heterogeneity between the clients. In this case, the generalization gap increased as well, According to Eq. (7) in our main paper, higher r leads to lower y, which increases the strength of the regularization (Eq. (3) in our main paper). To summarize, our adaptive approach dynamically adjusts the strength of the regularization according to the current heterogeneity between the clients during training, which reduces the generalization gap in FL models.
Discussion:
22

---

Published as a conference paper at ICLR 2025
1.50
1.00
0.94
0.34
0.52
0.92
0.00
of byyr
offlayer
yof layer
y layer 36
byyer
7
ayer
Vot layer
v layer
oflaye
volf layer
Vet layer 18
0.56
kyw
layer
y layer
layer
oflayer
ayer
lnyer
Communiration Rounds
(a) ResNer20.
1.20
0.50
0.50
0.97
0.56
layer
of layer
of leyer
layer
of layer 5
10
20
40
50
Communication Rounds
(b) CNN.
1.000
0.356
0.936
0.900
0.994
0. 992
0.900
layer
of layer
layer
10
of layer
of layer
of layer 7 layer
10
20
Communication
(c) VIT.
Figure 9: The value each layer's shrinking factor for different model architectures during the training process, with the dataset being CIFAR-10.
23

---

Published as a conference paper at ICLR 2025
Aerurcy
Asquaacy
118
113
130
CDr30 Ansee22
Areist
Frrist.a00
Tirymagenn:
Accurscy
Accuracy
Fading
Fediog
Rounds
Rounds
Rounds
Rounds
Rounds
200. NaNNe13
AECUrac
RudMS
Fed
Facia5 FosMe)
Pedd
Fodiig
Fodiig
Rounds
Sounds
Rounds
Bounds
Bounds
Clar3o, R.m C+L Natomb.5
O'erdd Num-20 Ratioe:
Cltddd OHN Namma3, End Ratpo3
Charl8 CVN nuevoD, E-1, Jatoni
Clrdo. Mam-d0. Batlon3
Accursoy
SccuPaCy
FudWS
PedUV5 Fediig
Tscis5 TacAkg
Fedlit Fediig
TedUNS Toskng
Rounds
Rounds
Rounds
Rounds
Rounds
Figure 10: Training Accuracy Across Various Datasets and Experimental Configurations (client number, local epoch E, and select ratio).
24

---

Published as a conference paper at ICLR 2025
Clor
(a) FashionMNIST, = 0.1.
(b) FashionMNIST, 0.5.
(c) FashionMNIST, 100.
Disinbution Cleets
over
Dierbbuttn
3-
150
254
Cloc
(d) CIFAR-10, 0.1.
CIFAR-10, 0.5.
(f) CIFAR-10, 100.
Figure 11: Data distribution over categories and clients.
25


Reviews:
[{"review_id": "dkMRNrbDIv", "reviewer": "Reviewer_2U5t", "summary": "This work studies server-side aggregation algorithms for non-IID federated learning, proposing an adaptive layer-wise weight shrinking strategy. Compared to the existing weight shrinking method, the proposed approach does not require auxiliary datasets to adjust the weights and allows for different shrinkage across various layers of neural networks.", "questions": "1. The relationship between the ratio of updating terms in Equation (4) and the gradient variance in Equation (5) is not immediately clear. It would be helpful if the authors could provide an intuitive explanation or theoretical evidence that supports this connection.\n\n2. More guidance is needed on how to choose the hyperparameter $\\beta$ and its potential impact on performance.\n\n3. The current experiments only consider one text dataset; it would be beneficial to explore additional scenarios.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The layer-wise strategy is well-motivated, as non-IID data may have varying disparity effects on different layers/modules.\n\n2. The proposal can be integrated into many non-IID-resistant training paradigms, such as FedProx and FedDisco.", "weaknesses": "1. The proposal presents similarities to existing concepts in federated learning, particularly the idea of layer-wise weighted aggregation (e.g., [a], [b]) and the weight shrinking technique from Li et al. (2023a). To enhance the clarity and impact of the proposed method, it would be beneficial for the authors to explicitly outline the novel aspects of their approach in comparison to these prior works.\n\n[a] Lee, Sunwoo, Tuo Zhang, and A. Salman Avestimehr. \"Layer-wise adaptive model aggregation for scalable federated learning.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. No. 7. 2023. \n\n[b] Ma, Xiaosong, et al. \"Layer-wise model aggregation for personalized federated learning.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.\n\nA comparison highlighting how the proposed method differs from or improves upon the layer-wise aggregation techniques discussed in [a] and [b] would be valuable. What unique contributions does this method bring to this area?\n\n2. While the proposal shows strong empirical performance, it would be more persuasive with theoretical justification. Specifically, the authors might provide a theoretical analysis of the advantages (e.g., faster convergence rate) of layer-wise shrinking compared to the original weight shrinking approach from Li et al. (2023a).", "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": "", "extra_scores": {}}, {"review_id": "0tZrqD8MTC", "reviewer": "Reviewer_wdHM", "summary": "This paper presents FedLWS (Federated Learning with Adaptive Layer-wise Weight Shrinking), a novel model aggregation strategy designed to improve the generalization of global models in Federated Learning (FL). Recognizing that a shrinking factor (sum of aggregation weights) less than 1 can enhance model generalization, FedLWS adaptively adjusts the shrinking factor at each layer based on gradient dynamics without requiring a proxy dataset, thereby reducing privacy risks. The approach calculates layer-specific shrinking factors during training, leveraging the unique characteristics of each layer, and is compatible with existing FL methods. Extensive experiments demonstrate FedLWS’s consistent performance improvements over state-of-the-art methods across diverse scenarios. A noted limitation is its current applicability only to scenarios where client models share the same architecture. Future work will address this limitation by exploring heterogeneous model scenarios.", "questions": "- How does the proposed layer-wise dynamic adjustment method demonstrate superiority over traditional dynamic adjustment techniques in the context of weight aggregation?\n- Can the authors provide a detailed background on prior experimental findings to help readers understand the motivation behind the approach?\n- In the experimental section, why does the proposed method perform better, and can this be explained rather than just described?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper addresses privacy concerns raised in prior research regarding the use of proxy datasets, and adopts an intuitive assumption in federated learning to dynamically adjust the shrinking factor (γ) layer-by-layer based on gradient variance. This approach aims to enhance model generalization while preserving privacy. The paper is well-structured, with extensive experimental support provided in both the main text and appendices. The key strengths of this approach are as follows:\n\n- FedLWS avoids the need for proxy datasets, reducing potential privacy risks and making it more suitable for sensitive applications.\n- The layer-wise adjustment of the shrinking factor (γ) allows FedLWS to adapt to specific layer characteristics, resulting in better model generalization across diverse and heterogeneous datasets.\n- As a “plug-and-play” approach, FedLWS can be easily integrated with existing federated learning methods, improving their performance without requiring substantial modifications to the training process.", "weaknesses": "1.  The introduction does a good job of citing relevant works to contextualize the study. However, it might be beneficial to discuss more recent research on adaptive weight shrinkage and layer-wise adaptation in FL, if available. This could provide a clearer picture of how the proposed FedLWS method aligns with and advances the current research landscape in FL.\n2.  This paper introduces a thoughtful approach by linking the shrinking factor to gradient variance, which allows dynamic adjustment of regularization. This concept appears promising and well-motivated, especially as it addresses practical deployment issues by eliminating the need for a proxy dataset. However, the authors could further clarify how this method stands out from other adaptive approaches. Are there any existing methods that also use gradient variance for similar adjustments, and if so, how does FedLWS provide additional benefits?\n3. The paper presents an intuitive hypothesis on the relationship between regularization, optimization, and gradient variance. Additional theoretical support or references could help strengthen this hypothesis. While the authors propose a function to capture this relationship, further clarification on the rationale behind this formulation would be valuable. Could a theoretical proof or justification for this relationship be provided?\n4. FedLWS employs a dynamically adjusted shrinking factor γ to balance regularization and optimization, but it is unclear whether this adjustment strategy could negatively impact convergence in federated learning. Frequent adjustments to the shrinking factor during multiple communication rounds may lead to instability or slower convergence. The authors could analyze the convergence behavior of FedLWS over multiple communication rounds, especially under varying data heterogeneity or client participation rates. Providing a convergence analysis or additional experimental results would help ensure that the method maintains stable convergence properties in federated learning.\n5.  The modularity of FedLWS as a “plug-and-play” solution is a valuable contribution. However, the paper could benefit from a more detailed example of how FedLWS integrates with popular FL methods like FedAvg or FedProx. For instance, describing specific integration steps for one or two methods would help readers understand the practical applicability of FedLWS.\n6. FedLWS employs a strategy of assigning different shrinking factors to different layers, based on the assumption that each layer requires a distinct level of regularization. However, in federated learning, different types of layers (e.g., convolutional layers and fully connected layers) may exhibit significantly different behaviors, and a simple layer-wise assignment may not be suitable for all types of layers. Could further experiments be conducted to verify the effectiveness of layer-wise shrinking across various types of layers? Additionally, it may be worth considering adjustments to the shrinking factor allocation based on the function or structural characteristics of each layer (e.g., differences between shallow and deep layers). This could potentially help FedLWS more effectively optimize each layer of the model.\n7. The results in Tables 1 and 2 effectively demonstrate that FedLWS consistently improves accuracy across datasets and heterogeneity levels, especially for more challenging tasks. However, a more detailed explanation of why FedLWS performs differently across various datasets would be helpful. For instance, why does FedLWS perform better on datasets with higher complexity or heterogeneity? This is particularly noteworthy given that prior work had the advantage of an additional proxy dataset as prior knowledge.\n8. In Figure 1(a), we observe a sudden drop when α  reaches 0.7, followed by a return to stability. Has the author considered the reasons behind this occurrence? In Figure 1(b), even a very small decrease (e.g., 0.01) has a significant impact on the results. Why is this the case?  The use of a 5-layer CNN model with uniformly decreasing values from 1 to 0.96 is an interesting experimental design choice. However, the authors could clarify the rationale for this specific range and pattern. Was the decrease from 1 to 0.96 based on empirical observation, theoretical insights, or a predetermined hypothesis? Including a brief justification for this choice could add depth to the experimental setup.\n9. The calculation of layer-wise shrinking factors appears to significantly increase computational overhead, especially in deep or large models such as ViT or BERT. Could the authors provide an analysis of this aspect?\n10.  Although FedLWS shows performance advantages in most experiments, it may underperform compared to FedLAW in certain cases. The authors attribute this entirely to FedLAW’s use of a proxy dataset, but this explanation seems insufficient, as FedLAW consistently uses a proxy dataset across all experimental settings. The authors should analyze under which experimental conditions FedLAW outperforms FedLWS and examine whether specific variables might explain this phenomenon.\n11.  Although the paper includes multiple image and text datasets, it does not explain whether these datasets offer generalizability. In particular, AG News, as the sole text dataset, is insufficient to fully validate FedLWS’s performance on other types of text data or time-series data.\n12.  Privacy is an essential consideration in FL, and the authors have addressed this well by removing the reliance on a proxy dataset. However, additional details on how FedLWS manages potential privacy risks, especially during gradient and parameter calculations, would be useful. Additionally, it would be valuable to discuss if FedLWS has any compatibility with existing privacy-preserving techniques like differential privacy.\n13.  The comparison between FedLWS and adaptive weight decay methods (AWD and AdaDecay) is insightful, showing FedLWS’s advantage in FL. Adding a brief explanation of why FedLWS outperforms weight decay in the FL context (e.g., due to its layer-wise adaptability) would make this comparison more compelling. The authors do not explain this reason in Appendix A.2.", "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": "", "extra_scores": {}}, {"review_id": "NdWKZwV457", "reviewer": "Reviewer_Ab6D", "summary": "This work proposes a simple and practical model aggregation method to enhance the performance of various federated learning algorithms.", "questions": "1. The proposed weighting method used to determine the model aggregation accords with intuition but lacks theoretical proof. Can a theoretical proof be given to ensure that the weighting coefficient determined from the gradient variance perspective is optimal?\n\n2. In Table 1, the performance of FedProx is generally worse than that of FedAvg. What is the reason for this phenomenon? The experiment should ensure that FedProx uses appropriate hyperparameters and that the comparison between the proposed LWA-enhanced methods and other baselines is fair, that is, the hyperparameters that make baselines' performance optimal should be used.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The proposed scheme is a direct intuitive scheme, which is easy to realize in practice. From the experimental results, the proposed algorithm can indeed play a role in improving performance.", "weaknesses": "Whether the proposed scheme is theoretically optimal is worth considering.", "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": "", "extra_scores": {}}, {"review_id": "ARmEmuEddF", "reviewer": "Reviewer_motB", "summary": "This paper introduces a method, FedLWS, to overcome the limitations of previous research named FedLAW, which leverages a Global Weight Shrinking (GWS) effect using a learnable shrinking factor based on proxy data. The authors identify two key issues with FedLAW: privacy concerns and considering the divergence across different layers of a deep neural network. To address these issues, they propose an alternative approach for calculating the shrinking factor using the variance of local model gradients. This method is based on the empirically observed relationship between the ratio of the regularization term to the optimization term and local gradient variance. Furthermore, this calculation enables the authors to propose a layer-wise shrinking factor for each layer, considering inter-layer differences for improved model generalization. Extensive experiments, including ablation studies, are conducted across various scenarios to verify the effectiveness of the proposed method. Overall, the authors present an effective alternative to address the limitations of previous research, and the experimental results appear to support its effectiveness. However, a more detailed analysis of the experimental results is necessary.", "questions": "1. Is it correct to use the equal sign between the left and right sides of equation (4)? Would it be more appropriate to use a proportionality sign instead?\n\n2. What does the sudden decline in the value in the range of 0.6 - 0.7 in Fig 1(a) signify?\n\n3. The calculation of the layer-wise shrinking factor seems to increase the computational load, yet Table 3 reports only a 0.02-second increase. Is there a technique applied to compensate for the increase in computational load?\n\n4. In Figure 3(a), why does the value of layer-wise gamma increase monotonically from early to later layers?\n\n5. In section 5.3, the results in Table 4 and the narrative are inconsistent. While the performance has improved model-wise compared to FedAvg, it is stated otherwise. The author insists(line 448): on the CIFAR-10 dataset with NIID (α= 0.1), the performance of model-wise weight shrinking is not as effective as FedAvg, whereas the layer-wise weight shrinking effect outperforms FedAvg.\n\n6. In Table 5, is the result for FedLAW with WRN56_4, indicated as 18.44, correct? If so, why is the accuracy for this setting significantly lower than that of the others?\n\n7. In Fig 6, it seems to be no effect of beta. What is the reason for selecting beta as a hyperparameter?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. By proposing an alternative method for calculating the weight shrinking factor, the authors improve performance over the previous algorithm, reduce computation time on the server side, and eliminate privacy risks associated with the use of proxy data.\n\n2. The authors analyze the characteristics of weight updates resulting from the application of the weight shrinking factor in the model agregation process of FL, validating these characteristics through equations and experiments.\n\n3. The effectiveness of the proposed method is validated through extensive experiments conducted in various environments, as well as ablation studies under diverse conditions.", "weaknesses": "1. The validation of the proposed method’s effectiveness across various experimental conditions is a strong point; however, the results are presented as simple listings or repetitions of earlier sections without in-depth analysis, making it difficult to gain insights into how the proposed method operates from a representation perspective. For example, in lines 388–389, the authors describe the results in Table 1 by comparing them with the baseline method, FedLAW. The authors note that in some settings, the method underperforms compared to FedLAW, attributing this to the latter’s use of proxy data. However, as FedLAW utilizes proxy data in all cases, this explanation does not fully account for the variations in performance improvement across different cases.\n\n2. There are several typos across multiple lines, and sentences with similar meanings are repeated, leading to a lack of coherence in the overall narrative.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)