Energy and Memory-Efficient Federated Learning with Ordered Layer Freezing and Tensor Operation Approximation

We would like to clarify some points regarding the assumption about Centered Kernel Alignment (CKA) and its applicability in the FedOLF approach.

First, we would like to clarify that FedOLF does not "assume" that low-level layers across local models exhibit higher degrees of CKA similarity. Specifically, the conclusion that low-level layers tend to have higher CKA similarity is based on the theoretical results presented in prior works, notably in [5-6, 9].

As stated in lines 279-282 of our paper, we rely on the findings of these works, which empirically and theoretically demonstrate that low-level layers across different models tend to exhibit higher CKA similarity than higher-level layers. Therefore, this is not an assumption, but rather a proven result derived from prior research.

For a deeper theoretical understanding of why CKA is an appropriate metric for comparing the similarity of layer outputs, we encourage you to refer to Section 3 in [9], which provides a detailed explanation. It highlights why CKA is a robust metric due to its ability to capture important properties like non-invertible linear transformations and isotropic scaling, which are crucial for comparing feature representations in neural networks, we’ve also cited this work in lines 282-283.

Regarding your concern about the performance of individual layers under diverse data distributions, we refer you to several works that provide empirical evidence of how low-level layers exhibit more consistent CKA similarity across various datasets and conditions. For example, Figs. 1-3 in [5], Figs 2-3 in [6], and Fig. 7 in [9]. These results show that lower layers universally obtain higher CKA similarities than upper layers across different datasets, non-iid data distributions, or model architectures. Therefore, we believe that the idea of low-level layer sharing will continuously work well even on highly diverse data, we have included this justification in lines 280-283.

We agree that both memory and computational constraints are crucial when considering the scalability of a federated learning approach. However, as discussed in [8], memory constraints are typically more critical than computational and communicational constraints in real-world federated learning scenarios. This is because memory limitations can directly prevent a device from participating in the federated learning process, whereas computational and communicational constraints mainly affect efficiency, slowing down the training process or increasing costs. In contrast, memory constraints are more restrictive, as a device unable to store sufficient model parameters cannot participate in training, leading to exclusion from the FL process altogether. This is especially problematic in non-IID environments, where excluding important clients can result in significant performance degradation, we introduce this point in lines 37-40.

In addition, even though we do not particularly target addressing computational and communications constraints, our method still has higher computation and communication efficiency than existing methods due to its superior performance. As shown in Fig.4, FedOLF naturally improves the computation and communication efficiency by obtaining higher accuracy with the same amount of overall (computation+communication) energy consumption. Based on these points, we believe that considering the memory footprint is a sufficient design for FedOLF.

Thank you for your insightful comment. We appreciate your suggestion to include a detailed layer-by-layer analysis of error propagation and to visually track error dynamics during training. Below, we address your concerns and provide additional clarification on the vanishing error assumption for frozen layers.

For in-depth analysis of vanishing representation layers, please refer to Appendices A and B, in formula, the norm of error $\sigma_{l}$ is upper bounded: $\|\sigma_{l}\| \leq \prod_{q=1}^{l-1}B_{q} \|\sigma_{1}\|$, where $\sigma_{1}$ is the representation error of the input layer and $B_{q}$ is the Lipschiz-continuity parameter (see Assumptions 1 and 1.1). We empirically show that the term $\prod_{q=1}^{l-1}B_{q}$ is decaying and therefore the error $\|\sigma_{l}\|$ is likely to be vanishing. For illustration, Figs 9-10 (Appendix A) present the visualization of layer-by-layer error propagation. In these figures, we exhibit the representation errors of each layer and find that $\sigma_{l}$, i.e. the representation error on a layer $l$, is likely to decay as $l$ increases.

However, if you feel that the empirical analysis is insufficient, we would like you to refer to Eq (16) and Appendix B.2 in paper [10] where an equivalent assumption is made. This paper first shows that the overall gradient error is bounded by a linear function of the last layer's gradient error, i.e. $\|\nabla f - \nabla f'\| \leq c_{1} \|\nabla f^{(N)} - \nabla f'^{(N)} \| + c_{2}$. ($c_{1}, c_{2}$ are constants and $N$ is the total layer number). Then the paper assumes $\|\nabla f'^{(N)} - \nabla f^{(N)} \|$ is bounded to therefore bound $\|\nabla f - \nabla f'\|$. This assumption is equivalent to the vanishing representation error justification in this paper. To prove this, first note $\|\nabla f^{(N)} - \nabla f'^{(N)}\| \leq L \|\sigma_{N-1}\| = L \prod_{q=1}^{N-1}B_{q}\|\sigma_{1}\|$ due to the $L$-smoothness of $f$. Then, if $\|\nabla f^{(N)} - \nabla f'^{(N)}\|$ is bounded, $\prod_{q=1}^{N-1}B_{q}$ must be shrinking, meaning that the representation error must be vanishing. Because by contradiction, if the errors are non-vanishing, $\prod_{q=1}^{N-1}B_{q}$ will explode as $N$ increases and $\|\nabla f^{(N)} - \nabla f'^{(N)}\|$ cannot be bounded.

Therefore, we believe that vanishing representation error is a reasonable assumption as an equivalent assumption has already been made by existing works (we mention this work in lines 782-783, Appendix A).

Thank you for your thoughtful questions.

For Q3, as discussed in our previous response to Weakness 1, uniform client clustering is a common and widely used approach in federated learning experiments.

For Q4, Regarding Section 3.5 and the convergence of FedOLF, we would like to clarify that FedOLF does not necessarily require any client to perform full-model training for convergence. This is supported by the findings (such as Figure 7) in the CKA paper [9]. It shows that for any two models, the CKA similarities between the very first few layers are usually high, even if one model is untrained and the other is fully trained. This means that, even if the first layer never gets trained, the output representation error will be quite limited compared with a trained model. Therefore, if none of the clients can do full-model training and the first few layers are forever frozen, FedOLF can still use these layers for forward propagation with limited training error. Moreover, the convergence analysis in Section 4 remains valid regardless of how many clients can do full training. Therefore, FedOLF can still converge effectively without requiring clients to perform full-model training. The trade-off is that if fewer clients have full-model training capability, the accuracy will be somewhat lower due to increased gradient errors $D$ from frozen layers (see Section 4). However, the system will still converge to a solution, albeit with lower accuracy than if more clients were able to perform full training.

For Q5, as stated in our response to Weakness 1, FedOLF is already evaluated in a scenario with a large proportion (80%) of resource-constrained clients and highly skewed class distributions (Dirichlet 0.1). For data-level assumptions, please refer to Assumption 3 (Eq. 9) which bounds the difference between local and global gradients. This is a common assumption that is pre-made by many existing works, such as [13, 14, 19, 20, 21].

As shown in Figs. 11 and 13 in Appendix C, FedOLF has equivalent, even less training cost (FLOPs and energy) compared with other methods (the cost of training time is proportional to FLOPs). Moreover, due to its superior performance, FedOLF is able to achieve higher accuracy with the same amount of energy consumption (Fig. 4), which significantly improves both computation and communication efficiency.

1. Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos Venieris, and Nicholas Lane. Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout. In NeurIPS, 2021.

2. Enmao Diao, Jie Ding, and Vahid Tarokh. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. In ICLR, 2021.

3. Minjae Kim, Sangyoon Yu, Suhyun Kim, and Soo-Mook Moon. DepthFL : Depthwise federated learning for heterogeneous clients. In ICLR, 2023.

4. Kilian Pfeiffer, Martin Rapp, Ramin Khalili, and Joerg Henkel. CocoFL: Communication- and computation-aware federated learning via partial NN freezing and quantization. Transactions on Machine Learning Research, 2023b.

5. 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. In NeurIPS, 2021

6. Weishan Zhang, Tao Zhou, Qinghua Lu, Yong Yuan, Amr Tolba, and Wael Said. Fedsl: A communication efficient federated learning with split layer aggregation. IEEE Internet of Things Journal, 2024.

7. Sebastian Caldas, Jakub Konecny, H. Brendan McMahan, and Ameet Talwalkar. Expanding the reach of federated learning by reducing client resource requirements. In NeurIPS Workshop on Federated Learning for Data Privacy and Confidentiality, 2018.

8. Kilian Pfeiffer, Ramin Khalili, and Joerg Henkel. Aggregating capacity in FL through successive layer training for computationally-constrained devices. In NeurIPS, 2023a.

9. Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In ICML, 2019.

10. Mirzasoleiman, Baharan, Jeff Bilmes, and Jure Leskovec. ”Coresets for data-efficient training of machine learning models.” In ICML, 2020.

11. Menachem Adelman, Kfir Levy, Ido Hakimi, and Mark Silberstein. Faster neural network training with approximate tensor operations. In NeurIPS, 2021.

12. Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. Fedproto: Federated prototype learning across heterogeneous clients. In AAAI, 2022.

13. 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, 2020.

14. Ang Li, Jingwei Sun, Pengcheng Li, Yu Pu, Hai Li, and Yiran Chen. Hermes: An efficient federated learning framework for heterogeneous mobile clients. In Proceedings of ACM MobiCom, 2021.

15. R. Tamirisa et al., FedSelect: Personalized Federated Learning with Customized Selection of Parameters for Fine-Tuning, CVPR, 2024

16. Ziru Niu, Hai Dong, A.K. Qin. FedSPU: Personalized Federated Learning for Resource-constrained Devices with Stochastic Parameter Update, arXiv:2403.11464

17. Kaiming He and Jian Sun. 2015. Convolutional neural networks at constrained time cost. In CVPR, 2015.

18. Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, and Soumith Chintala. 2020. PyTorch Distributed: Experiences on Accelerating Data Parallel Training. Proc. VLDB Endow. 13, 2020.

19. Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In ICML, 2020.

20. Zhida Jiang, Yang Xu, Hongli Xu, Zhiyuan Wang, Jianchun Liu, Qian Chen, and Chunming Qiao. Computation and communication efficient federated learning with adaptive model pruning. IEEE Transactions on Mobile Computing, 2023.

21. Canh T. Dinh, Nguyen H. Tran, and Tuan Dung Nguyen. Personalized federated learning with moreau envelopes. In NeurIPS, 2020.

22. Tan, Jiahao, et al. "pFedSim: Similarity-Aware Model Aggregation Towards Personalized Federated Learning." arXiv preprint arXiv:2305.15706.

Thank you for your valuable questions.

For Q1, as shown in Theorem 3 in the original TOA paper [11], the TOA error is proportional to 1/s − 1. Therefore, in terms of accuracy, a larger s always leads to better accuracy, and the optimal value of s is naturally one. However, in practice, the optimal value of s depends on how a client would like to balance accuracy between the actual memory footprint and communication cost, which is hardware-dependent.

For Q2, in practice, changing the value of s has very little impact on the computation overhead. Because TOA is only applied to frozen layers, and varying s only affects the workload of forward propagation. However, forward propagation accounts for the overall computation overhead much less than forward propagation [16-18]. Therefore, it is unnecessary to vary s dynamically based on clients’ different hardware. Besides, the universal hyperparameter setting and grid search are common and unquestionable settings in many FL papers, such as [7, 8, 12, 13, 14, 15]

Dear reviewer,

Thank you for your insightful feedback. Below, we address each point raised in your comment.

For weakness 1, as stated in the experiment setup (lines 400-408), the experiments are already conducted in a scenario where most (80%) are resource-constrained and unable to train an entire model. In practice, the definition of "resource-constrained" can be complicated and multi-aspect (e.g. computation, communication and memory). In the experiment simulation, the resource constraints can be reflected by the number of frozen layers for each client as discussed in [4], and we follow the same setting here as justified in lines 460-461.

For uniform clustering, this paper follows the same experiment setting of FjORD which uniformly divides clients into five clusters [1]. Besides, most existing works employ a uniform client clustering policy in the experiment settings. For example, HeteroFL [2] divides clients into five uniform clusters with sub-model ratios 0.0625, 0.125, 0.25, 0.5 and 1.0. DepthFL [3] has four uniform clusters with 1-4 residual blocks respectively. CoCoFL [4] has three uniform client clusters that train 1/3, 2/3, and all of the layers respectively. Therefore, uniform client clustering is a quite common and sufficient selection to simulate a real-world network with heterogeneous clients in FL (see our justification in lines 454-456). Furthermore, to the best of our knowledge, none of the existing works utilize a non-uniform clustering strategy. If you have specific references that demonstrate the significance of non-uniform clustering for federated learning, we would be eager to explore this further.

For skewness, the evaluation has already been conducted with a skewed class distribution among clients. As introduced in the experiment setup, this paper follows [5] and allocates data to clients unevenly. Specifically, the total number of samples per client, and the number of categories per client both follow a Dirichlet Distribution with parameter $\alpha$. The smaller $\alpha$ is, the more skewed class distributions we have. In this paper, we set $\alpha=0.1$, which is already considered an extreme case for class imbalance according to [5] (see lines 452-454). Given this setup, we believe that our evaluation already captures the challenges posed by class imbalance. If you could provide more insight into why you think further evaluation with more skewed distributions is necessary and what kind of skewness setting can represent typical real-world scenarios, we would be happy to consider expanding the experiments.

Dear reviewer,

Thanks for your response.

First, we would like to justify that the statement “This statement clearly states the limitations of the proposed technique” may not be appropriate. Accuracy decay along with the reduction of full-model clients is a common property of SOTA works and is not unique in our approach. This has been verified by extensive experiment results in the baselines' papers, such as: Fig.3 in Feddrop [7], Fig. 7 in FjORD [1], Table 5 in DepthFL [3], Table 1 in HeteroFL [2], Table 2 in CoCoFL [4] and Table 3 in SLT [8].

Second, to answer the questions “Also, can this technique be used when most of the devices are significantly memory constrained. If yes, what are the limits? Also, how does the proposed technique compare with the state-of-the-art techniques under such scenarios." and "The non-uniform clustering is just to identify the limitations of the proposed technique in the context where only a very limited number of devices (e.g., just one or none) have full model training capability":

we’ve reduced the number of full-model clients to one (they are replaced by clients who can only train 20% of the entire model) and evaluated FedOLF on CIFAR-100 / ResNet-20 and got 37.62% accuracy, which is only 0.23% lower than the uniform case (see Table 1), which shows the robustness of FedOLF to extreme non-uniform cases where most clients are resource-constrained. Due to a limited response time, we are unable to provide the accuracy of the baselines in the non-uniform case. However, we can see that FedOLF in non-uniform case still outperforms the baselines in uniform case (Table 1):

It has been well-studied that the baselines’ performance always decrease when it goes from uniform case to non-uniform case (see examples above). Therefore, the SOTA’ accuracy in non-uniform case will definitely be lower than the uniform case, and hence lower than FedOLF in non-uniform case. To our best knowledge, we haven’t seen any possible limits in the non-uniform scenario.

Lastly, the question “If state-of-the-art techniques outperform the proposed under such scenarios, what scenarios are ideal for the proposed technique?” has also been addressed as SOTA methods will not outperform FedOLF.

Regards

We respectfully address the points raised in your latest comments, ensuring clarity and rigor in our explanations.

1. Harsh Memory Constraints

Your statement that “the authors have not tested the concept under harsh memory constraints” appears to overlook our initial response, where we explicitly outlined an experimental setup involving 80% memory-constrained devices—a scenario widely recognized in the literature as a harsh memory-constrained setting [1-3]. Additionally, we conducted further experiments as you requested, reducing the number of full-model clients to just one, and our results remained strong. If this does not satisfy the standard for "harsh memory constraints," we kindly request you to clarify what would constitute an acceptable threshold, along with relevant citations.

2. Computation and Communication Constraints

Regarding the claim that "memory being the most important constraint" was supported by only a few references, we again encourage you to provide the additional references that demonstrate the equal importance of computational, memory, and communication constraints in FL). Although FedOLF focuses on memory constraints as they are the most pertinent factors to determine whether edge devices can do full-model training, the computation and communication constraints are NOT overlooked in this paper. First, OLF simultaneously mitigates computation overhead by minimizing the length of backward propagation path (Fig. 1). Second, we propose TOA to further reduce computation and communication overheads. Thirdly, your statement that we completely ignore the computation and communication constraints is contradictory to your initial review, wherein it is clearly written “Applying TOA to frozen layers preserves model accuracy while reducing the communication cost”.

3. Diverse Non-iid Data distributions

As stated in your initial review, the goal of testing FedOLF with diverse non-IID data distribution is to validate the claim that "low-level layers universally obtain higher CKA similarity." In our last response, we showed that this concern is misplaced, as the CKA similarity property has already been substantiated in multiple prior works ([5,6,9]). Therefore, CKA similarity is a verified finding, not an untested assumption. In this case, no extra tests are needed for any verification, and reiterating the requirement of diverse non-iid evaluation seems to be unreasonable.

Furthermore, all datasets we employed (EMNIST, CIFAR10, CIFAR100, and CINIC10) are NOT “additional”, as the CKA theory is proven to work well on these datasets [5,9,22].

4. Comparison to CoCoFL

You mentioned that layer freezing were proposed in CoCoFL ([4]). However, the main contribution of our paper lies in ordered layer freezing rather than layer freezing only. We have introduced many aspects wherein our method (ordered freezing) differs from CoCoFL, including:

• Backward gradient path minimization (Figure 1)
• Low-level layer sharing (Section 3.5)
• Limited training error (Section 3.6)
• Tensor Operation Approximation (TOA) (Section 3.3)

These innovations result in significant performance improvements, as demonstrated in our evaluations. Thus, FedOLF represents a distinct and superior approach, not merely an extension of CoCoFL. Besides, “layer freezing and quantization have been proposed in COCOFL” is a new concern that you didn’t raise in the original review, which should not be proposed in the discussion phase. Introducing new critiques during the discussion phase risks undermining the review process's integrity.

5. Handling Computational Constraints

Regarding devices with limited computation power, our method (and SOTA) employs a synchronous FL protocol—a standard practice in the field. Under this protocol, the server waits for all clients to complete training. We have demonstrated that our method does not increase overall training time compared to baselines (Figure 11, Appendix C). While latency is a limitation of synchronous FL, it is not unique to our method and should not be grounds for rejection.

6. Inconsistency Between Comments

By saying “all computational, memory, and communication constraints have to be considered for the system design. See [4]”, you admit that [4]’s setting (layer freezing, uniform clustering) presents an eligible system design in overall aspects. This is contradictory to your another statement “such constraints have been completely ignored by the authors”, as our work follows the setting of [4].

7. Review Integrity

Lastly, we sincerely would like you to explain why you didn’t propose these questions in your last response. As a matter of integrity, reviewers are expected to raise all questions once, allowing authors sufficient time to address them comprehensively.

Dear reviewer,

Thank you for reviewing the paper.

We appreciate the reviewer’s observation regarding the prevalence of layer freezing in machine learning. While it is true that layer freezing has been explored extensively in centralized machine learning, our contribution specifically addresses its application in the FL paradigm, which presents unique challenges and dynamics not encountered in centralized settings. To the best of our knowledge, our work is the first to propose Ordered Layer Freezing (OLF) in FL, emphasizing the systematic ordering of layer updates based on task-specific importance—a key differentiator from existing works.

We acknowledge that high-level layers often contribute more to task-specific learning than low-level layers, as noted. However, our approach goes beyond this established understanding by (1) providing a principled strategy for ordering updates in FL to optimize both resource efficiency and learning performance, and (2) demonstrating its benefits with rigorous quantitative analyses, including memory footprint reduction (Fig. 1), generalized feature sharing across low-level layers (Section 3.5), and detailed convergence behavior (Section 4). These analyses validate the effectiveness and novelty of OLF in the FL context, where communication constraints and heterogeneous client data necessitate innovative adaptation of layer freezing.

Regards

For Q1, as described in Appendix C.2, the communication energy consumption is measured using a plug-in power monitor. In each communication round, we use the monitor to measure the energy consumption of the computer to acquire this round’s communication energy consumption, and the total communication cost is the summation of all rounds’ communication energy cost.

For Q2, we thank the reviewer for raising this important point. We have updated our submission and attached the accuracy vs. round graphs to the appendix (Figs 14 and 15 on page 21). The graphs show that FedOLF (orange curve) has higher accuracy and better convergence in most cases.

For Q3, we attribute this result to the unique property of SLT. For ResNet, SLT progressively trains each block in a bottom-first order. When the model increases from ResNet20 to ResNet44, the number of layers in each block also increases. In this case, the accuracy of SLT significantly increases as it trains more layers in the same rounds. Even though the performance gap between SLT and FedOLF gets narrowed as the model grows, FedOLF still achieves higher accuracy and energy efficiency as shown in Table 1 and Figure 4.

Regards

Dear reviewer,

Thank you for your insightful feedback. Below, we address each point raised in your comment:

For weakness 1, we appreciate the reviewer’s observation regarding the prevalence of layer freezing in machine learning. While it is true that layer freezing has been explored extensively in centralized machine learning, our contribution specifically addresses its application in the FL paradigm, which presents unique challenges and dynamics not encountered in centralized settings. To the best of our knowledge, our work is the first to propose Ordered Layer Freezing (OLF) in FL, emphasizing the systematic ordering of layer updates based on task-specific importance—a key differentiator from existing works.

For weakness 2, we appreciate the reviewer’s suggestion regarding trading off time for training versus memory footprint through on-demand re-evaluation of results. To address this trade-off in the context of OLF, we analyze the interplay between communication efficiency, memory footprint, and training quality. As shown in Figure 7b, decreasing the TOA scaling factor $s$ significantly reduces the size of the transmitted frozen layers, thereby lowering memory and communication costs. However, as TOA becomes more aggressive (i.e., $s$ is smaller), the model’s performance deteriorates proportionally. This decline in performance demands more training iterations, ultimately leading to increased energy consumption, as demonstrated in Figure 4, where the energy cost per accuracy increment rises with aggressive TOA settings. Furthermore, Figure 4 highlights that FedOLF consistently achieves the highest communication and computation efficiency among all methods without requiring additional compromises in training quality. Given that our approach already addresses memory constraints effectively while maintaining high training performance, adopting on-demand re-evaluation would introduce unnecessary overhead and inefficiency.

For weakness 3, we appreciate the reviewer’s observation. To address this, we note that our analysis accommodates scenarios where $B$ is not universally smaller than one, such as in networks like ResNet where $B$ is not smaller than one (see Appendix A). In this scenario, even though $B$ is not universally smaller than one across all layers, the representation errors still exhibit a decaying trend toward high-level layers.

To further support this argument, we refer to Eq (16) and Appendix B.2 in paper [1] where an equivalent assumption is made. This paper first shows that the overall gradient error is bounded by a linear function of the last layer's gradient error, i.e. $\|\nabla f - \nabla f'\| \leq c_{1} \|\nabla f^{(N)} - \nabla f'^{(N)} \| + c_{2}$. ($c_{1}, c_{2}$ are constants and $N$ is the total layer number). Then the paper assumes $\|\nabla f'^{(N)} - \nabla f^{(N)} \|$ is bounded to therefore bound $\|\nabla f - \nabla f'\|$. This assumption is equivalent to the vanishing representation error justification in this paper. To prove this, first note $\|\nabla f^{(N)} - \nabla f'^{(N)}\| \leq L \|\sigma_{N-1}\| = L \prod_{q=1}^{N-1}B_{q}\|\sigma_{1}\|$ due to the $L$-smoothness of $f$. Then, if $\|\nabla f^{(N)} - \nabla f'^{(N)}\|$ is bounded, $\prod_{q=1}^{N-1}B_{q}$ must be shrinking, meaning that the representation error must be vanishing. This is because, by contradiction, if the errors are non-vanishing, $\prod_{q=1}^{N-1}B_{q}$ will explode as $N$ increases and $\|\nabla f^{(N)} - \nabla f'^{(N)}\|$ cannot be bounded.

1. Mirzasoleiman, Baharan, Jeff Bilmes, and Jure Leskovec. ”Coresets for data-efficient training of machine learning models.” In ICML, 2020.

For weakness 4, to clarify, the effectiveness of FedOLF does not depend on the quality of the initial model parameters, either theoretically or empirically. Firstly, our convergence analysis in Section 4 does not rely on any assumptions related to model initialization. In FedOLF, given any arbitrary original point, the global model will eventually converge to a ϵ-critical point with bounded error. Secondly, in the experiment, FedOLF is always evaluated with random model initialization and still produces high accuracy. Therefore, the original model’s quality does not have any effect on FedOLF either theoretically or empirically.

For w2, we believe that training time vs. memory tradeoff will not work well. First, we’ve shown that training lower layers isn’t quite memory-efficient (lines 55-86). For example, if you only train the first bottom layer while freezing the rest, the activation values in the frozen layers are needed to compute the gradient of the first layer, and still have to be stored in memory as discussed in [4]. Given that activation takes a great proportion of the overall memory usage, the memory usage of low-level layer training can be very close to full-model training. Therefore, for a device with strong computation resources but very limited memory space, training the first few layers will NOT take a long time for training, instead, it will raise the “out-of-memory” error and terminate immediately.

On the contrary, in FedOLF, only the top active layers need to store the activation as lower frozen layers do not involve any gradient computation (see Fig. 1). Therefore, compared with other layer freezing schemes, FedOLF always minimizes the total activation size stored in memory given the same amount of frozen layers.

Second, increasing the training time simultaneously increases the computation cost, which is not quite efficient. Also, as mentioned in Section 3.5, lower layers generalize well across different local datasets and can be shared between clients. Therefore, it is unnecessary for memory-constrained clients to train lower-layers, as they can directly use the layers that are pre-trained by other clients.

For w3, we appreciate your constructive feedback. We have updated our paper and removed the “B<1” claim. As a replacement, we re-write Assumption 1 as “each layer $W_{l}$ is $B_{l}$-lipschitz continuous”. That is, we change $B$ to $B_{l}$, meaning that the bound is no longer universal and depends on the layers. We find that even though $B_{l}$ can be large occasionally for some layer in some particular model structures (e.g. ResNet), the accumulative product $\prod\_{l} B\_{l}$ usually presents a overall decaying trend (This claim is supported by the results in Appendix A). This means that the overall training error, which is proportional to $\sum^{N}\_{l} \prod^{l}\_{q}B_{q}$ is always bounded, and our convergence analysis still holds (see Appendix B).

For w4, we have attached the graphs reflecting accuracy per iteration in Appendix D (Figs 14-15, page 21). In all scenarios, we use an random initial model with accuracy close to zero (i.e. the initial model is very far from the optimal model). The graphs show that FedOLF (orange curve) has higher accuracy and better convergence in most cases, regardless of the quality of the initial model.

[4] Kilian Pfeiffer, Ramin Khalili, and Joerg Henkel. Aggregating capacity in FL through successive layer training for computationally-constrained devices. In NeurIPS, 2023a.

Regards.

Dear reviewer,

Thanks for your response.

For w1, we apologize for the confusion. The unique challenge in FL is about how to preserve accuracy under non-iid local data and resource-constrained edge devices.

In centralized ML, all data are stored in a central server. The server usually has a strong capacity and therefore a higher training cost. In this setting, layer freezing is applied to reduce the server's training cost. Since the powerful server usually has full-model training capacity, it can conduct a few full-training steps first to access the heuristic of all parameters, like weight or gradient. Then, the server can selectively freeze the unimportant parameters in each layer to reduce the computation cost with minimal accuracy loss. In centralized ML, accuracy and efficiency can be easily achieved simultaneously.

In FL, however, devices with important data may not be able to train a full model owing to memory constraints, and therefore cannot participate in FL. When the global model cannot learn important knowledge from these devices, performance degradation will occur. To address this issue, we need to find a way to enable these important but memory-constrained devices to join FL.

We provide an insightful analysis to show that ordered layer freezing (OLF) is the most effective method (compared with SOTA) to address this challenge with higher accuracy and smaller memory requirements, without using any heuristic (full-model training).

For accuracy, prior works [1-3] show that high-level layers are more sensitive and customized to different local datasets. These layers need to be frequently trained to ensure that the global model learns sufficient knowledge from diverse local data. On the other hand, low-level layers are more generalized across different data distributions, i.e. they do not need to be trained by every client and can be shared among clients (we discussed this point in Section 3.5). Moreover, OLF largely preserves accuracy due to the vanishing representation error property (Section 3.6). Because of these advantages, FedOLF achieves higher accuracy than SOTA methods (see Table 1).

For memory, we’ve shown that ordered-layer freezing significantly reduces memory footprint by minimizing the backward propagation path’s length (Fig. 1). Compared with existing layer freezing methods CoCoFL and SLT, FedOLF has much lower memory usage (Figs. 5-6) and hence is more compatible with memory-constrained edge devices.

The expression “a systematic ordering of the layer updates based on the task” means that the hyperparameter $l_{k}$, i.e. the number of frozen layers , can be determined adaptively based on the concrete task. As shown in Section 3.4, in real scenarios, we can calculate the memory footprint based on the task’s specific content (e.g. what dataset / model / algorithm we use), and choose the proper $l_{k}$ to ensure that the memory usage does not exceed the limit of the physical device. Meanwhile, in the baselines, the hyperparameters are determined arbitrarily in the simulation, which is unformulated and difficult to use in real-world applications.

[1] Weishan Zhang, Tao Zhou, Qinghua Lu, Yong Yuan, Amr Tolba, and Wael Said. Fedsl: A communication efficient federated learning with split layer aggregation. IEEE Internet of Things Journal, 2024.

[2] 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. In NeurIPS, 2021.

[3] Mirzasoleiman, Baharan, Jeff Bilmes, and Jure Leskovec. ”Coresets for data-efficient training of machine learning models.” In ICML, 2020.

Dear reviewer XVEh,

It's been a couple of days since we replied to your latest response. As the discussion deadline is close, we would like you to provide a feedback as soon as possible.

Regards