Why Go Full? Elevating Federated Learning Through Partial Network Updates

Thank you for the reviewers' comments. Next, we will address each of the issues raised one by one.

W1 & Q1: What exactly is the update stepsize in Fig1? ... The authors try to use Figure 1(a) to show that the 'update stepsize' increases after each averaging and thus it indicates layer mismatch, ... why does this draw the conclusion that it is caused by 'inadequate cooperation among layers'?

We address W1 and Q1 together as they are closely related. In this paper, the step size refers to the magnitude of model parameter changes between iterations. Since the learning rate is fixed in this experiment, this value directly reflects the norm of the gradient. Therefore, during a normal neural network training process, the step size gradually decreases towards zero, leading the model to converge to a fixed point. However, we found that in federated learning, the step size significantly increases after each parameter averaging. This increase in step size suggests that while the global model parameters might be generally robust and contain more knowledge, they cannot cooperate well to complete local tasks immediately (i.e., there is a parameter mismatch). Additionally, because the gradient calculation follows the backpropagation algorithm, the gradient of a parameter in one layer is only directly related to the parameters in the subsequent layers, with no direct correlation among parameters within the same layer. Therefore, we believe this mismatch mainly exists between layers and we refer to this phenomenon as layer mismatch.

W2: The proposed strategy to select trainable layers does not provide much insight ...

We appreciate the reviewer's comments. In fact, although our method seems simple and intuitive, it is grounded in significant reasoning. The Sequential Training principle addresses the observation that neural networks converge shallow layers before deeper ones, a trend that is even more pronounced in federated learning [1]. From this perspective, in current federated learning, clients need to train deep networks without the shallow layers being fully converged, resulting in wasted computational power. Our sequential training strategy can effectively address this issue, but can lead to significant degradation in neural network performance. Therefore, we proposed another principle, Multi-round Cycle Training, which is inspired by the classic optimization method BCD (Block Coordinate Descent). By repeating cycles of training multiple times, it effectively alleviates the performance degradation caused by layer-wise training.

In addition, before arriving at this final solution, we tried many different approaches, such as various parameter selection strategies (both randomly, reversely) and optimizer designs, some of which are presented in the ablation studies. Furthermore, while the final method's design is relatively simple, it still involves many details, such as determining the timing for full parameter updates and balancing the length of each cycle with the total number of cycles.

W4: Regarding communication cost, ... (does) only updating one layer at a time increases the total number of communication rounds?

We will first address the reviewer's fourth weakness, as we believe the third point will be better addressed after clarifying this one. The reviewer misunderstood our experimental results and our method does NOT increase the required number of communication rounds. In all experimental tables, we ensure the number of communication rounds is identical when comparing full parameter and partial parameter training. For example, in Table 1, the row with C=1 shows the performance of each method at the round when partial network training completes the first cycle, ensuring a fair comparison across all algorithms.

W3: The experiment setup ... involves 5 rounds of full network training between each cycle, which I assume is a lot ... Since the amount of full network trainings is significant in FedPart, one cannot confidently tell if it is really the partial network trainings that are truly beneficial.

The reviewer's understanding of the experimental setup is generally correct, but we do not agree with the skepticism. Based on the response to weakness 4 above, our method does not increase the communication rounds. Therefore, from an equivalent perspective, in the process where all rounds involve full parameter training, we modified some rounds to partial parameter training and observed significant performance improvements. Logically, this clearly indicates that partial network updates contributed to the improvement.

Q2: The comparison in Fig1b really does not look like a big difference to me ... And also I don't understand why Fig1a does not match Fig1b in terms of the blue curve

Regarding the experimental results in Fig. 1b, the reviewer felt the effect was not significant. We respect the different opinions raised by the reviewer. However, we believe that, at least it is fair to say that FedPart can significantly reduce the degree of layer mismatch in extreme cases and generally maintain the original effect in typical cases, which still supports our claim. Concerning the reviewer's comment on the difference in the full parameter training curve, we believe there is some misunderstanding. Indeed, the blue curves in Fig 1a and Fig. 1b are NOT calculated in the same way. As the step size is the total magnitude of model updates, if not all dimensions are updated, this magnitude will naturally decrease, but this is an unfair comparison. To make the comparison as fair as possible, the full parameter training curve in Fig. 1b actually shows the "update magnitude of the same parameters as those updated in partial parameter training," ensuring a fair comparison. We apologize for the confusion and will clarify it in a later version.

[1] Unlocking the Potential of Federated Learning for Deeper Models.

Dear Reviewer 6JAi,

Could you please respond with how you think about the authors' response? Please at least indicate that you have read their responses.

Thank you, Area chair

I have read the authors' responses and discussions with other reviewers. I believe the authors have addressed some of my concerns (e.g. W3 W4) and questions, so I will update my score. However, I still have some concerns about the vaguely proposed term 'layer mismatch', and the authors' responses do not convince me. I asked if there are other evidences or related observations in previous literature to support the claim, but authors did not answer it. This skepticism remains also due to the lack of explanation about Fig 1(b) where I asked why the difference only seems big for the 50 to 60 iterations region but not else where. Considering that Fig 1(b) seems to be the only experimental support for the claim, not addressing questions regarding it makes the 'layer mismatch' claim unsure to me.

Thank you for the reviewers' comments. Here are our responses.

Q1: In Figure 1, the update step size ... seems extremely sharp ... if the momentum terms ... were properly reset?

We did reset the optimizer's state before the start of each local round. The relevant code is at line 131 in the file fling/component/client/base_client.py.

Q2: In Section 3.3 ... please comment on the assumption and put into perspective how realistic they are.

Thank you for the suggestion. Here are more detailed comments about the assumptions.

Assumptions 1 means that the gradient of the loss function does not change too abruptly as we move in the parameter space. This is reasonable for most models, where the loss changes gradually as the parameters are adjusted. Assumption 2 suggests that the variability in the gradient estimates is controlled. This is generally true in datasets without much extreme values, where individual data points do not drastically affect the overall gradient.

Assumption 3 means that the effect of applying different masks to the variability of gradient should be similar. As this assumption is first proposed in this paper, we next give some evidence. We conducted Monte Carlo simulations (10,000 samples) to approximate the value of $k$ in Assumption 3. As is shown the the table below, $k$ is close to 1 under different settings, supporting Assumption 3.

Q3: Regarding the communication costs, how do clients know the mask ?

Actually, the server does not need to transfer the mask $S^t$ to each client. Since the trainable parameters are determined at the layer level, the server only need to transfer the indexes of trainable layers, resulting in almost no communication overhead. Such a notation is just for easier formulation.

Q4: ... the larger (60 round) warm-up performs significantly better ... How does table 2 look when run with longer warm-up? What is the reason not to do that?

Extending the warm-up period leads to higher final model accuracy but also increases computational and communication overhead. For instance, to achieve extreme performance, one might need to train with FedAvg until full convergence before adding several rounds of partial network training. The reason why we did not use such a setting in Table 2 is that, we want to show a balanced improvement for both model accuracy and system overhead.

Q5: In the section on the warm-up ablation ... FedPart still enhances the model’s accuracy after so many warm-up rounds. Where can I see this in the presented numbers?

The relevant figures can be found in Table 6. In the table, bef. refers to the accuracy before performing partial network updates (at the end of the warm-up stage), and aft. refers to the accuracy after partial parameter training. For example, the accuracy after full parameter training at round 60 in CIFAR-10 is 58.92, while after one cycle of partial parameter training, this number become 66.18.

Q6: ...repetitions to assess experimental statistical significance is not given for Table 6-9 ...

We appreciate the reviewer’s suggestions. We will amend the checklist arguments accordingly.

Q7: In Figure 2, it is a bit unclear to me why the layers have these different shapes

Sorry for the confusion. In Figure 2, we used different shapes to represent different network parameters. Misalignment of shapes between layers indicates a mismatch, while aligned shapes indicate a match between layers.

Q8: The activation maximization approach is old ..., with MACO, yielding great visual representations that would be much more useful in accessing the meaning.

We appreciate the suggestion. Due to time constraints, we cannot provide an improved visualization immediately but will adopt modern methods in future versions.

We next address the concerns raised in the "Weakness" section.

W1: ... (the proposed method) seemingly makes the models much more vulnerable to data reconstruction compared to an update of the full model ... a discussion on this matter would be important.

Thanks for the comment. Compared to full-parameter training, we believe our proposed method will not increase the risk of privacy leakage. We can abstract the model training process (for both full and partial parameter training) as a mapping: $(\Delta w_1, \Delta w_2, ..., \Delta w_n) = f(x)$, where the left hand side denotes the updates to each model parameter, and $x$ is the training data. From a privacy attack perspective, the goal is to find the best $x$ such that the updates to $w$ are as close as possible to the actual updates in each dimension. This resembles solving a system of equations, where $x$ are the unknowns, and each dimension of $w$ update represents an equation. With partial network training, the unknowns $x$ remain unchanged compared with full parameter training, but the number of equations decreases (i.e., less information for the attack to follow). Therefore, we believe partial network training leakages less information in general.

W2: ... the experimental setup is counter-intuitive: you “insert five rounds of full network training between each cycle in FedPart”, which means you do more full training than partial updates ...

There may be a misunderstanding. A "round" refers to one communication between the server and clients, while a "cycle" refers to the entire process of training a partial network from shallow to deep layers, including many rounds. For instance, in ResNet-18 with $R/L=2$, one cycle includes 36 rounds, indicating far more partial-parameter training than full-parameter training.

Thank you for your positive feedback. We appreciate your helpful comments and we will make sure to include the additional discussions on assumptions and privacy in the final manuscript.

Thank you for the reviewers' comments. Below, we address the concerns raised in the "Weakness" section.

W1: I suggest the author provide more evidence and discussion of the proposed layer mismatch challenge in the paper. In the current version, there is no experiment to validate the proposed challenge and how this challenge happens during the training.

We conducted experiments to verify the issue of layer mismatch, as shown in Fig. 1. Here we detail the logic of these experiments. In Fig. 1a, we observed that the step size significantly increases after each model averaging, which is detrimental to the final model convergence. We believe this phenomenon is caused by layer mismatch, and thus we designed a method targeting reducing layer mismatch (i.e., layer-wise training). The results show that after using our method, the step size variation is significantly reduced (as shown in Fig. 1b), and the final performance is improved (as shown in Table 1). Therefore, we believe that layer mismatch does exist, affects model convergence, and consequently impacts federated learning performance, which is the core challenge addressed.

Reviewers might have questions about the relationship between step size and layer mismatch, so we provide a more detailed explanation here: Formally, step size resembles the norm of the gradient in each model update. In a normal neural network training process, the step size gradually decreases towards zero during training, leading the model to converge to a fixed point. However, in federated learning, the step size significantly increases after each parameter averaging (Fig. 1a), indicating that the global model parameters do not collaborate well to complete local tasks (i.e., there is a parameter mismatch). Additionally, because the gradient calculation follows the backpropagation algorithm, the gradient of a parameter in one layer is only directly related to the parameters in the subsequent layers, with no direct correlation among parameters within the same layer. Therefore, we believe this mismatch mainly exists between layers (i.e., layer mismatch).

W2: In the experiment part, the paper mainly focused on the iid setup. Even in the ablation study (line 273), the paper adopts alpha equals 1 as the non-iid setup. In most of the FL papers, the experiment would adopt an alpha equal to at least 0.1 to simulate the extreme non-iid cases. As a result, the experiment in the paper could not provide useful support to the authors' argument. Also, the client number is very limited (only 40 clients), and the experiment does not contain any client sampling in the discussion.

Thank you for the comments. We conducted additional experiments with 150 clients, randomly sampling 20% of the clients for training and aggregation in each communication round. The results are shown below, indicating that FedPart still performs better than FedAvg in this scenario.

We also added experiments with an alpha=0.1 setting as suggested by the reviewer. The results are as follows. It can be seen that partial network training still exhibits much faster convergence compared to full network training and achieves comparable accuracy to full parameter training. Here, it can be seen that in extreme data heterogeneity, we were unable to achieve accuracy improvements because the main issue in this setting is client drift rather than layer mismatch. Since our method is not specifically designed to solve data heterogeneity, in future work, we may introduce other solutions for data heterogeneity to improve performance further.

W3 & W4 & W5 & W6: Even though the paper proposed many related studies, the experiment part contains no other layer-wise FL training baselines. As a result, it is really hard to judge whether the proposed method could stand for the SOTA performance. I noticed that some recent papers also proposed similar layer-wise training methods [1,2]. I suggest the author to include them as baselines to compare with.

We appreciate the reviewers' comments. However, these works cannot be used as benchmarks for comparison, as they are designed for a significantly different scenario from ours. Their purpose in using layer-wise training methods is to enable effective training of networks with different architectures in each client. However, if their methods are applied in our scenario (i.e., the same model is used on all clients, and the computational capacity is consistent), these methods will degrade to full parameter training (i.e., FedAvg, which is compared in our paper). Therefore, we believe comparing with these methods is unnecessary.

Thank you for the reviewer's time and patience. However, we believe there are still some misunderstandings regarding this issue, so we would like to clarify further:

Q1: Specifically, in your newly-added results, the performance between FedAvg and FedPart is very similar in the challenging dataset (CIFAR-100 and Tiny-ImageNet), and the performance is under the iid setup (alpha = 1) based on my understanding.

Thank you for the reviewer’s feedback. In fact, the performance difference between FedAvg and FedPart on challenging datasets is not similar. As shown in the results, on CIFAR-10, CIFAR-100, and Tiny-ImageNet, the final accuracy achieved by our method is 67.08%, 42.12%, and 26.75%, respectively, while FedAvg achieves 65.00%, 40.55%, and 23.33%, respectively. This means our method outperforms FedAvg by 2.1%, 1.6%, and 3.4%, respectively. These show that the improvement on more challenging datasets is still substantial, even exceeding the gains on simpler datasets. Please note that in our first table provided during the rebuttal phase, "C" represents different training time points, and when comparing model performance, the final state (C=5) should be considered.

We summarize the performance improvement of our method over FedAvg in three different scenarios, which is shown in the table below. It can be observed that our method generally shows greater improvements on more challenging datasets, which demonstrates the advantages of our approach.

Q2: Also, the provided non-iid test indicates that the FedPart nearly does not provide a performance upgrade compared to the FedAvg, the performance is so close when C equals 2. In addition, the author does not provide the non-iid test on the CIFAR-100 and Tiny-ImageNet. Based on the newly-added results, it would provide an implication that the proposed FedPart only works well on the easy dataset with iid setup.

Thank you for the reviewer’s feedback. We have the following points to clarify:

• First, the reviewer mentioned that the proposed FedPart only works well on the easy dataset with iid setup. For our understanding, the reviewer may take $\alpha = 1$ as iid or nearly iid. However, we would like to clarify that $\alpha=1$ represents a scenario with relatively significant data heterogeneity, and it is widely used in the literature [1,2,3,4]. In contrast, $\alpha = 0.1$ is a extreme non-iid setting, as the reviewer noted in the comment. Under the setting of $\alpha=1$, as shown in Table 4 of the main text, our method achieves final performance improvements of +1.0%, +2.2%, and +3.4% on CIFAR-10, CIFAR-100, and Tiny-ImageNet, respectively. These performance gains are not negligible.
• We agree with the reviewer that, in this extreme non-IID scenario ($\alpha = 0.1$), the model accuracy of our method is roughly on par with that of the full parameter method. However, this does not imply that FedPart offers no performance advantages—the benefits primarily arise from reduced communication and computation costs. The results indicate that FedPart can achieve similar accuracy to FedAvg while significantly reducing communication and computation costs (these metrics are consistent with those observed in the IID scenario). As shown in Table 1, when training on Tiny-ImageNet, FedPart reduces communication overhead by 72% and computation overhead by 27%. Therefore, we believe that even in such an extreme scenario of data heterogeneity, our method still holds practical value.

The above is our further response to the reviewer's comments. Thank you very much for the reviewer's patience, and we hope that if there are any further questions, the reviewer can feel free to ask us directly.

[1] Xu, Jian, Xinyi Tong, and Shao-Lun Huang. "Personalized federated learning with feature alignment and classifier collaboration." International Conference on Learning Representations (ICLR). 2023.

[2] Oh, Jaehoon, Sangmook Kim, and Se-Young Yun. "Fedbabu: Towards enhanced representation for federated image classification." International Conference on Learning Representations (ICLR). 2022.

[3] Tan, Yue, et al. "Federated learning from pre-trained models: A contrastive learning approach." Advances in neural information processing systems 35 (2022): 19332-19344.

[4] Zhang, Jianqing, et al. "Fedala: Adaptive local aggregation for personalized federated learning." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. No. 9. 2023.

Dear Reviewer,

Thank you once more for the time and effort you invested in reviewing our paper, as well as for your insightful follow-up feedback regarding the non-iid setting. In response to your comments, we have made several clarifications, which we hope address at least some of the concerns you raised. We acknowledge that there is still room for improvement under the extreme non-iid setting (e.g., $\alpha=0.1$). Please let us know if there are any additional questions or if further clarification is needed.

Thank you for the reviewers' comments. Here are our responses.

Q1: How would the method perform on IID and non-IID data?

Thank you for the reviewer's comment. We have conducted additional experiments to enrich our analysis of non-IID data scenarios. We added experiments with an alpha=0.1 setting as data heterogeneity is more severe. The results are as follows. It can be seen that partial network training still exhibits much faster convergence compared to full network training while maintaining the final accuracy. This further shows the robustness of our proposed FedPart under data heterogeneity.

W1: ... limited in terms of dataset sizes ...

We acknowledge your concern regarding the dataset sizes used in our experiments. As we mentioned in the paper, we did not use particularly large datasets to test our method. However, the datasets we used are comparable in size to those used in related works in the field [1, 2], and we believe that these datasets are sufficient to demonstrate the efficacy of FedPart. We will aim to incorporate more extensive datasets in future iterations of our research.

W2: ... synthetic data experiments are missing ...

We appreciate your suggestion to include synthetic data experiments. This is indeed a worthwhile addition that can provide more insights into the performance of FedPart in various settings. However, due to time constraints, we cannot immediately incorporate synthetic data experiments. We will seriously consider adding synthetic data experiments in both IID and non-IID settings in future versions of our work.

[1] McMahan, Brendan, et al. "Communication-efficient learning of deep networks from decentralized data." Artificial intelligence and statistics. PMLR, 2017. [2] Li, Qinbin, Bingsheng He, and Dawn Song. "Model-contrastive federated learning." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.