Why Go Full? Elevating Federated Learning Through Partial Network Updates

Authors: haolin  wang
xuefeng  liu
jianwei  niu
wenkai  guo
shaojie  tang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 09:27:54

Abstract: Federated learning is a distributed machine learning paradigm designed to protect user data privacy, which has been successfully implemented across various scenarios. In traditional federated learning, the entire parameter set of local models is updated and averaged in each training round. Although this full network update method maximizes knowledge acquisition and sharing for each model layer, it prevents the layers of the global model from cooperating effectively to complete the tasks of each client, a challenge we refer to as layer mismatch. This mismatch problem recurs after every parameter averaging, consequently slowing down model convergence and degrading overall performance. To address the layer mismatch issue, we introduce the FedPart method, which restricts model updates to either a single layer or a few layers during each communication round. Furthermore, to maintain the efficiency of knowledge acquisition and sharing, we develop several strategies to select trainable layers in each round, including sequential updating and multi-round cycle training. Through both theoretical analysis and experiments, our findings demonstrate that the FedPart method significantly surpasses conventional full network update strategies in terms of convergence speed and accuracy, while also reducing communication and computational overheads.

Main paper: 
 introduction: 
 Federated learning is a machine learning framework that protects data privacy, which has attracted widespread attention from researchers in recent years [McMahan et al., 2017 , Kairouz et al., 2019 , Li et al., 2019] . In traditional federated learning, after receiving the global model sent by the server, each client uses their local data to update the entire model parameters set for several iterations; then, the server averages the updated models to obtain a new global model and broadcasts it to all clients, starting the next training round.While this approach has proven effective in many applications [Hard et al., 2018 , Rieke et al., 2020] , its convergence speed and ultimate performance are often lower than those of centralized schemes [McMahan et al., 2017 , Zou et al., 2023] , even when data across clients are independently and identically distributed (i.i.d.) . This suggests that while full network updates and sharing enrich each model layer with more knowledge, they may also introduce factors that negatively impact final performance. To further investigate the underlying reason, we conduct an experiment to visualize the update step sizes during each iteration. Typically, in centralized learning, the update step sizes of the model show a downward trend, indicating that the model is gradually converging. However, in federated learning, as shown in Fig. 1a , the update step sizes significantly increase after each parameter averaging. This suggests that after averaging, the gradients calculated by subsequent layers become notably large, indicating inadequate cooperation among layers within the global model, a phenomenon we term as layer mismatch. The cause of this issue is illustrated in Figure 2a . The middle section of the figure depicts the local models of each client, which have undergone sufficient local training. Within these local models, the layers cooperate effectively, demonstrating match. However, upon aggregating the parameters of each layer, the averaged layers may struggle to maintain this cooperation, resulting in mismatch. This layer mismatch can result in two key issues: first, it may prevent the global model in federated learning from converging to the optimal point of the global loss function, thereby compromising performance. Second, the persistent mismatch disrupts the federated learning process at the server level, significantly reducing training efficiency. To address the aforementioned problems, we propose FedPart, which employs partial network updates.Our main motivation is illustrated in Fig. 2b . In this toy example, we assume that only the i-th layer of the network is trainable in a given round. During local training on each client, this trainable layer can naturally align with the fixed parameters of other layers, which serve as anchors that constrain its update direction. This makes the averaged layers better align with other layers. To validate this approach, we conduct experiments and visualize the results in Figure 1b . The curves demonstrate that partial network updates significantly reduce the increase in update step sizes after averaging, confirming their role in alleviating layer mismatch.However, as a trade-off, training and transmitting only a portion of the parameters at a time might limit the efficiency of knowledge learning and sharing. Through thorough analysis, we identify a solution to this challenge, which involves carefully selecting trainable parameters. This solution is based on two key principles. The first principle is sequential updating. We train the network layers sequentially, from shallow to deep, one layer at a time. This design is based on the observation that the shallower layers of a neural network typically converge to their final parameters faster than the deeper ones [Raghu et al., 2017] . To align with this natural order, we adopt a similar sequential strategy for layer selection. The second principle is the multi-round cycle training strategy. Our method emphasizes the importance of repeating the process of training from shallow layers to deep layers multiple times. During the original full network updates, shallow layers often learn low-level features, while deep layers learn high-level semantic features [Zeiler and Fergus, 2014, Erhan et al., 2009] . To preserve this property, inspired by the idea of Block Coordinate Descent (BCD) [Poczos and Tibshirani] , we propose the multi-round cycling training method to retain this characteristic to the greatest extent.In addition, FedPart offers improved computation and communication efficiency, making it highly suitable for edge computing scenarios [Wang et al., 2019a ,b, Abreha et al., 2022] . This is because FedPart only needs to train a part of the neural network at each training round, thereby significantly reducing the computational overhead of each client in each iteration. At the same time, since clients only need to upload and download the parts of the model that need updating, the amount of parameters to be transmitted is also greatly reduced.To validate the effectiveness of FedPart, we explore its performance from both theoretical and experimental perspectives. Theoretically, we demonstrate that FedPart has a superior convergence rate under non-convex settings compared to FedAvg. Experimentally, we perform extensive evaluations on various datasets and model architectures. The results indicate that the FedPart method significantly improves convergence speed and final performance (e.g., an improvement of 24.8% on Tiny-ImageNet with ResNet-18), while also reducing both communication overhead (by 85%) and computational overhead (by 27%) simultaneously. Furthermore, our ablation experiments demonstrate how each of the proposed strategies contributes to enhancing the overall performance of FedPart. We also conduct comprehensive visualization experiments to illustrate the underlying rationale of FedPart. In summary, the contributions of this paper are as follows:• We identify the issue of layer mismatch in federated learning, which arises from updating and aggregating all parameters in each training round. This phenomenon can potentially impact the model's convergence speed and overall performance.• To mitigate the effects of layer mismatch, we introduce FedPart, which implements partial network updates. Additionally, we develop corresponding strategies for selecting trainable parameters.• We theoretically analyze the convergence rate of FedPart in a non-convex setting, demonstrating its advantages over full network updates.• We perform extensive experiments, showing that FedPart achieves significant improvements across multiple evaluation metrics compared to the full network updates scheme. Additionally, ablation and visualization experiments enhance our understanding of the rationale behind FedPart. 
related work: 
 Current research on partial parameter training or aggregation in federated learning has led to various applications, broadly categorized into three types:Train all parameters, aggregate partial parameters. Also known as personalized federated learning, this approach involves each client training all parameters but only aggregating some of them [Tan et al., 2022] . For example, FedPer [Arivazhagan et al., 2019] and FedBN [Li et al., 2021b] personalizes classification and batch-normalization layers respectively, FedRoD [Chen and Chao, 2021] applies both global and local classifier heads. Other works may only upload a low-rank space of parameter matrices [Wang et al., 2023 , Wu et al., 2024a] . Although these methods achieve impressive results in data-heterogeneous scenarios, they usually exhibit performance degradation when datasets across clients are distributed in an i.i.d. manner. Moreover, they do not effectively reduce computational overhead, and any reduction in communication overhead is minimal, as the personalized components are typically small.Train partial parameters, aggregate all parameters. This category refers to each client training a different part of the model, with a global update to the entire model during aggregation. For example, PVT [Yang et al., 2022] and FedPT [Sidahmed et al., 2021] strategically assigns specific model layers to each client, Federated Dropout [Caldas et al., 2018] and FedPMT [Wu et al., 2023] randomly assign a neurons in different layers to clients, HeteroFL [Diao et al., 2020] and FjORD [Horvath et al., 2021] deterministically decide a trainable subnetwork based on client computational power, FedRolex [Alam et al., 2022] further introduces a sliding window method, and CoCoFL [Pfeiffer et al., 2022] introduces a quantization technique for overhead reduction. The primary objective of these methods is to reduce client overhead by dynamically leveraging varying computational capacities among clients. However, compared to full network updates, these methods often result in performance degradation and slower convergence speed.Progressive Training. This approach starts with a small model and gradually increases its size until the entire network is trained [Rusu et al., 2016] . This training paradigm has gain attention in the field of federated learning as its efficiency in reducing resource consumption (e.g., ProgFed [Wang et al., 2022] and ProFL [Wu et al., 2024b] ). However, because these methods eventually train a full model, they are not able to solve the layer mismatch problem. Moreover, while these methods aim to reduce resource consumption, they often lead to performance losses compared to full network training. To the best of our knowledge, our FedPart is the first approach to simultaneously enhance both convergence accuracy and efficiency.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We have discussed our limitations in the conclusion section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
method: 
 Generally speaking, FedPart is based on partial network updates, which trains and aggregates only a few layers of the global network model for each training round. At the beginning of each training round that requires partial network update, the server first determines which layers need to be trained and sends this information to all clients. Subsequently, each client trains the corresponding layers, transmitting them to the server for aggregation, and the server broadcasts the averaged results to each client for next training round. We elaborate on two key components of FedPart in the following subsections: partial network updates and the strategic selection of trainable layers.The partial network updates involve training and aggregating only a few layers of the global network model in each later communication rounds. Specifically, we partition the layers of global model into trainable ones and frozen ones. For each training iteration t and client i, the optimization objective is:arg min w t i E x∼D i [L(x| ŵt i , wt i )], where ŵt i and wt i respectively denotes parameters of trainable and non-trainable layers, w t i ≜ { ŵt i , wt i } represents the total parameter set, D i represents the local data distribution of client i and L(•) refers to the loss function. To optimize this objective, we adopt the following gradient descent formula:EQUATIONHere, γ is the learning rate, S t i is a binary mask that selectively enables updates only for trainable parameters and ⊙ denotes element-wise product. After performing several local training iterations, the parameters of these selected layers are sent to the server and globally averaged at iteration t = T : wT = 1 N N i=1 w T i , where N represents the number of clients. For the sake of simplicity in formulation, the above equation aggregates and calculates the gradient for all parameters. However, in practical implementation, we only update and transmit the trainable components, significantly reducing both computational and communication costs.Although training only a subset of parameters can largely mitigate the layer-mismatch issue, it may limit the efficiency of knowledge learning and sharing. Therefore, we propose to carefully select trainable layers to address this limitation. As illustrated in Fig. 3 , following the initial full network updates, we train parameters layer by layer from the shallowest to the deepest. Subsequently, we cycle back to the shallowest layer and periodically repeat this process. Generally, this strategy is driven by two key principles:Sequential updating. This principle refers to training model layers in sequence, from shallow to deep layers one at a time. Our motivation is based on the fact that the convergence of neural networks follows a natural intrinsic order, with shallower layers typically converging earlier than deeper ones [Raghu et al., 2017] . By partially updating network in accordance with this inherent training order, we can largely replicate the convergence process of full network updates while preserving training efficiency simultaneously.Multi-round cycle training. This principle refers to repeating the process of updating the neural network layers from shallow to deep multiple times. To illustrate our motivation, consider a fully trained neural network: shallow layers primarily focus on low-level semantic features (such as the edges in images), while deeper layers focus on higher-level semantic features (such as the main objects in images). However, during partial network updates, because the deeper layers are initially non-trainable, shallow layers are forced to learn complex high-level semantic features, which disrupts the original information hierarchy in the neural network. Through multi-round cycle training, we return to the shallow layers after training the deep layers. This strategy can reduce the burden on the shallow layers and helps approximate the final results of full network updates. To analyze the convergence of FedPart, we first introduce some definitions and notations. Let each client use a uniform loss function L(x|w) with parameters w, for the data x to calculate the loss function value, ) ] which is the expected loss function of client i. In this setting, the overall optimization goal of federated learning can be written as the sum of the expected loss functions of each client, that is:f i (w) = E x∼D i [L(x|wf (w) = 1 N N i=1 f i (w).Additionally, for notational simplicity, we denote the parameters of the i-th client at time t as w t i , the computed stochastic gradient vector as G t i , and the average of all client models at time t as wt . To represent partial network updates, we add a binary matrix S t i as a mask for each update process of each client. To keep consistency with the methods section, we assume that for each mask S t i , 1 M of the elements are set to 1, and the rest are set to 0. Before proving the convergence of our FedPart under non-convex conditions, we first introduce three assumptions:Assumption 1: The expected loss function of any client is L-smooth, namely:||∇f i (w) -∇f i (u)|| ≤ L||w -u||, ∀i, w, u.(2)Assumption 2: The variance and second-order moments of the gradients are bounded, that is:E x∼D i [||∇L(x|w) -∇f i (w)|| 2 ] ≤ σ 2 , ∀i, w, x ∈ D i , E x∼D i [||∇L(x|w)|| 2 ] ≤ G 2 , ∀i, w, x ∈ D i .Assumption 3: The variance of the gradients is approximately equal under all permissible masks:E x∼D i [||S 1 ⊙ (∇L(x|w) -∇f i (w))||] E x∼D i [||S 2 ⊙ (∇L(x|w) -∇f i (w))||] ≤ k, ∀i, w, x ∈ D i , S 1 , S 2 .(3)The first two assumptions are common in the literature, ensuring certain necessary characteristics of the loss function. The third assumption ensures that the specific choice of the mask matrix does not have too much impact on the final update variance, provided that the mask meets the required criteria.A further discussion about the third assumption can be found in Appendix G.Based on these assumptions, we can analyze the convergence of FedPart. In terms of the approximate convergence rate, we align with related literature [Alistarh et al., 2017 , Lian et al., 2017 , Ghadimi and Lan, 2013] , using the average magnitude of the expected gradient over iterations, and finally obtain the following theorem:Theorem 1: Under assumptions 1-3, with the total number of clients as N , and all parameters divided into M groups of trainable parameters, the convergence rate of FedPart satisfies:EQUATIONwhere ⊙ denotes element-wise product. A detailed proof of this theorem is provided in Appendix B.The results show that compared to the convergence rate of full network updates O( 1 √ N T ) [Yu et al., 2019] , FedPart demonstrates significantly better convergence performance. This advantage becomes more pronounced as the choice amount of partial parameters per instance is reduced, which aligns with our original intention to reduce layer mismatch. However, it should be noted that convergence analysis can only indicate the difficulty of converging to a stationary point, and cannot measure the model's performance after convergence. Therefore, it is not advisable to arbitrarily reduce the number of parameters trained in each iteration.Communication Cost. Suppose FedPart divides all layers into M groups, and only one group is trained during each partial training session, with one communication round for each group. It is easy to verify that the averaged communication costs of FedPart during the partial network update phase is reduced to 1 M of the original costs. Computational Cost. Assuming that computational costs are uniform across all layers, our method can reduce the overall computational expense during the partial network update phase by 1 3 . The primary reason for this reduction is that in FedPart, there is no need to compute gradients for the layers preceding the trainable parameters. To analyze this quantitatively, suppose the total overhead for both forward and backward propagation in a complete model satisfies D tot = D f or + D bak . The ratio of computational costs over a single training cycle can then be written asComp P N U Comp F N U = M * D f or + M i=1 i * D bak M M * D f or + M * D bak = M * D f or + M i=1 (M +1) * D bak 2M M * D f or + M * D bak .Moreover, it is widely accepted in the literature that the computational cost of backward propagation is approximately twice that of forward propagation [Rasley et al., 2020, Hobbhahn and Sevilla, 2021] . Therefore, the above equation can be rewritten as:Comp P N U Comp F N U ≈ M * D f or + (M + 1) * D f or M * D f or + 2M * D f or ≈ 2 3 (5)4 ExperimentsIn the experimental setup, we primarily choose 40 clients, with local epochs to be 8. We test the global model on a balanced set. Unless specifically stated otherwise, the training datasets across all clients are independently and identically distributed (i.i.d.). We utilize the Adam optimizer [Kingma and Ba, 2014] with a learning rate of 0.001, which is determined to be the optimal learning rate. In line with prior references [Li et al., 2021b , Chen et al., 2022] , we refrain from uploading local statistical information during model aggregation. Each experiment is conducted three times with different random seeds to ensure robustness. The experimental results for additional scenarios, including learning-rate tuning and client sampling, are presented in Appendix F.When choosing experimental metrics, we employ three distinct measures to capture various aspects of the benefits. These metrics include: Best Acc., which represents the ultimate accuracy achieved in classification tasks; Comm., indicating the total upstream transmission volume required by each client for a given training round (in GB); and Comp., which illustrates the total floating-point computation required by each client (in TFLOPs). All experiments are conducted on a server equipped with 8×A100 GPUs, and we provide the complete source code in our supplementary material.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: The full set of assumptions is provided in Section 3.3, and a complete proof is available in Appendix B.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
main properties: 
 Comparison with full network updates. We apply the FedPart method to three classic federated learning algorithms: FedAvg [McMahan et al., 2017] , FedProx [Sahu et al., 2018] , and FedMOON [Li et al., 2021a] , and compare the results with their full network updates (FNU) counterparts. We utilize ResNet-8 [He et al., 2016] (detailed in Appendix A) and update only one layer in each two consecutive training rounds (denoted as 2 R/L). Additionally, we insert five rounds of full network training between each cycle in our FedPart. We conduct experiments on the CIFAR-10 [ Krizhevsky et al., 2010] , CIFAR-100 [Krizhevsky et al., 2009] , and TinyImageNet [Le and Yang, 2015] datasets. The results in Table 1 show that our FedPart method demonstrates rapid convergence and consistently outperforms traditional FNU methods across all training cycles C, ultimately achieving significantly higher accuracy (e.g., improving FedAvg on Tiny-ImageNet by 21%). At the same time, its communication and computational costs are only 28% and 73% of those required by FNU. Furthermore, we observe that in some scenarios, the performance improvements of other federated learning algorithms even surpass those observed with FedAvg. This highlights that the layer mismatch problem identified in this paper is novel and cannot be addressed by any existing methods. However, our results on CIFAR-10 are less impressive. This suggests that in simpler datasets, the primary issue might be the client drift problem explored in previous studies, whereas the layer mismatch problem becomes more prominent in complex datasets.FedPart with deeper models. To evaluate the effectiveness of FedPart with deeper networks, we conduct experiments on ResNet-18 (detailed in Appendix A). This presents a more challenging scenario, as the proportion of trainable parameters significantly decreases in each round. Our experimental setup also follows the 2 R/L pattern, with five additional full network updates inserted between cycles. The results, displayed in Table 2 , show that in deeper networks, FedPart not only maintains its advantages in convergence speed and accuracy but also provides even greater reductions in communication and computational costs (by 85% and 27% compared to full network updates). FedPart for language modality. We also extend the FedPart method to the field of natural language processing and evaluate it on AGnews and So-gouNews [Zhang et al., 2015] datasets.We choose the transformer architecture [Vaswani et al., 2017] for experiments. As shown in Table 3 , the results indicate that FedPart performs well on language tasks, not only maintaining comparable performance as FNU, but also reducing communication and computational overhead by 66% and 25%, respectively. This demonstrates the method's scalability.FedPart under data heterogeneity. We also evaluate the performance of FedPart under scenarios involving data heterogeneity. The results in Table 4 show that our FedPart consistently improve final performance (e.g., an improvement of 3.4% on Tiny-ImageNet) in the presence of data heterogeneity. However, the extent of performance improvement is relatively smaller. This suggests that client drift [Karimireddy et al., 2020] may have a more pronounced negative impact on our method. We conduct experiments with extreme data heterogeneity (α = 0.1) in Appendix F.3. Training rounds per layer. In our FedPart, the training rounds per layer (denoted as R/L) is an important hyperparameter. A larger R/L value means more thorough training in each cycle, but it also results in a decrease in the number of cycles within the same number of training rounds. We explore the performance of FedPart under different R/L. From the results in Table 5 , when R/L=1, the outcome shows limited final performance due to insufficient training for each layer. However, further increasing the R/L value not only fails to improve the final performance but also reduces the convergence speed. In extreme cases, when R/L=10, only one cycle is conducted overall, which significantly affects both convergence speed and final accuracy. This indicates that generally, increasing the number of cycles is more effective than extending their duration. This aligns with the motivation behind our proposal of multi-cycling training. Rounds of initial warm-up updates. To explore the impact of the duration of the initial full network updates phase (i.e. warm-up stage), we conduct experiments with this stage set to lengths of 0, 5, and 60. In Table 6 , the term state refers to the period before or after partial network updates, which follow the warm-up phase. The experimental results clearly show that initial full network updates is crucial to the final model's accuracy. Notably, extending the full network update phase yields diminishing returns. However, even when the model is trained with FNU until no further accuracy improvement is observed (60 init.), utilizing FedPart still enhances the model's accuracy. This confirms FedPart's capability to improve the convergence of the final global model and reduce layer mismatch.Different orders for selecting trainable layers. We experiment with three different orders for selecting trainable parameters: sequential, reverse, and random. Sequential is the default configuration of FedPart, selecting layers from shallow to deep. In contrast, the reverse sequence selects layers from deep to shallow, while the random sequence selects layers randomly in each round. The results of the experiments are depicted in Table 7 , demonstrating that the effectiveness of the three methods ranks as follows: sequential > reverse > random. This aligns with the intrinsic convergence order of neural networks and meets our experimental expectations.In this section, we conduct experiments to demonstrate why our proposed parameter selection strategy can enhance final performance, and what the impact of layer-wise information exchange has on privacy leakage. Our experiments are based on ResNet-8 and the CIFAR-100 dataset. We analyze the models obtained from four different methods: 1) FedAvg-100, which represents training with full network for 100 rounds; 2) FedPart(No Init. 1C), which represents using FedPart for one cycle without initial full network updates; 3) FedPart(1C), which involves initial full network updates followed by one cycle of FedPart training; 4) FedPart(5C), which involves initial full network updates followed by five cycles of FedPart training. The visualization results are as follows.Activation maximization visualization. Activation maximization [Erhan et al., 2009] involves finding an input that maximizes a specific activation value within a neural network, reflecting the feature patterns the neuron focuses on. We use this method to explore the visual patterns captured by different models and measure their similarity using SSIM (Structural Similarity Index Measure) [Hore and Ziou, 2010] . The results in Table 8 show that, without initial full network updates and multiple cycles, the features captured by the FedPart model significantly differ from those of the FedAvg model. However, this discrepancy decreases after applying our layer-selection strategy, suggesting that the model better recognizes the hierarchical nature of different semantic information, thus enhancing its performance. Additional visual results are provided in Appendix C.We also analyze how different models extract semantic information by visualizing the convolutional kernels. We find that in the full network updates represented by FedAvg-100, the shallow convolutional kernels primarily function as edge/corner detectors. However, direct training of partial networks disrupts this property. Further, by initially employing full network updates and adding multiple training cycles, we gradually restore this characteristic. This effectively explains the impact of the parameter selection strategy on the final model formation. For specific visualization results of the convolutional kernels, please refer to Appendix D.We next demonstrate that FedPart offers enhanced privacy protection, as it transmits less information in each communication round. Formally, we can abstract the model training process (for both full and partial parameter training) as a mapping: (∆w 1 , ∆w 2 , ..., ∆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 attacker to leverage). Therefore, we believe that partial network training generally leaks less information.To verify this experimentally, we conduct several rounds of federated learning using both full network and partial network updates. We employ DLG (Deep Leakage from Gradients) [Zhu et al., 2019] to attempt the recovery of original images and use PSNR (Peak Signal-to-Noise Ratio) [Hore and Ziou, 2010] to measure the extent of privacy leakage. DLG is a classic privacy leakage scheme, which aims at finding an input that produces gradients most similar to the gradients calculated from a given sample. In this way, DLG can approximately recover the input sample. Let the original model input be x, then the specific formula for recovering the input x is as:min x ||∇ xL(x|w) -∇ x L(x|w)|| 2 (6)We use PSNR to measure the quality of the reconstructed image. Given that x denotes the original image and x denotes the reconstructed image, the PSNR is calculated as follows:EQUATIONwhere MSE(x, x) denotes the mean square error between m × n matrices x and x, given by:EQUATIONA larger PSNR value means a better quality of the reconstructed image, which further implies a higher risk of privacy leakage.The results in Table 9 show that, for different trainable layers, our method consistently exhibits better privacy protection in both average and worst-case scenarios compared to full network updates. Attacking examples are provided in Appendix E. Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We release our complete code in the supplementary materials.Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
conclusion and limitation: 
 We observe that the model averaged in federated learning is not directly applicable to the specific tasks of each client, a situation we refer to as layer mismatch. To address this issue, we propose the FedPart method, which introduces a strategy for selecting and training partial networks. We validate the effectiveness of FedPart both theoretically and experimentally. In future work, we plan to evaluate our method on a wider range of model architectures and apply it to larger-scale datasets to further investigate its effectiveness and scalability.#i represents the i-th layer of the model, with detailed partitioning method is presented in Appendix A.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: This full code is released in the supplementary materials, and the datasets are also public.Guidelines:• The answer NA means that paper does not include experiments requiring code. Justification: Descriptions about the resources required is stated in Section 4.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: The research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: This paper may have positive societal impacts due to its better ability for privacy protection, which is discussed in Section 4Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).11. Safeguards • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In Section 4, we primarily adopt ResNet and language transformer for experiments, whose architectures are illustrated in Fig. 4 and Fig. 5 , respectively.We also demonstrate the detailed partitioning method in our FedPart. Taking ResNet-8 (on the left in Fig. 4 ) as an example, we divide the trainable parameters of the model into 10 layers, corresponding to the numbers #1-#10. Among these, the trainable parameters of #1-#9 include not only the weights of the convolutional layers but also the weights and biases of the accompanying BN layers after the convolutional layers. The other models follow the same representation method of layer partitioning. During the sequential training phase of the FedPart method, we select one single layer to train in the order of their numbering #i. Before beginning the proof, we need to analyze the upper bound of the gradient variance after parameter selection. According to Assumption 3, we know that for any mask matrices S 1 , S 2 , it holds that:EQUATIONConstructively, we set a series of mask matrices S 1 , • • • , S M that have no overlapping '1' elements at the same positions and their sum exactly forms an all-one matrix. Clearly, each of these mask matrices meets our requirements. Therefore, we can derive:E x∼D i [||S j ⊙ (∇L(x|w) -∇f i (w))|| 2 ] ≥ 1 k 2 * E x∼D i [||S 1 ⊙ (∇L(x|w) -∇f i (w))|| 2 ] (10) Summing over j from 1 to M , the left-hand side of the inequality is exactly the variance of the gradient without any mask matrices. Therefore:M j=1 E x∼D i [||S j ⊙ (∇L(x|w) -∇f i (w))|| 2 ] = E x∼D i [||∇L(x|w) -∇f i (w)|| 2 ] ≥ M k 2 * E x∼D i [||S 1 ⊙ (∇L(x|w) -∇f i (w))|| 2 ]For following proof, we mainly refer to Yu et al. [2019] . According to Assumption 2, the upper limit of the left side of the inequality is σ 2 , so we finally obtain a general upper bound for the gradient with masks:EQUATIONWith the groundwork laid, we are now ready to begin the formal proof process. First, based on Assumption 1, as the loss function is L-smooth, we have:E[f ( wt )] ≤ E[f ( wt-1 ] + E[ S t i ⊙ ∇f ( wt-1 ), wt -wt-1 ] + L 2 E[|| wt -wt-1 || 2 ] (12)Next, we analyze the third term on the right-hand side of the inequality above to derive the following inequality:E[|| wt -wt-1 || 2 ] = γ 2 E[|| 1 N N i=1 S t i ⊙ G t i || 2 ] = γ 2 E[|| 1 N N i=1 S t i ⊙ (G t i -∇f i (w t-1 i ))|| 2 ] + γ 2 E[|| 1 N N i=1 S t i ⊙ ∇f i (w t-1 i )|| 2 ] = γ 2 N 2 N i=1 E[||S t i ⊙ (G t i -∇f i (w t-1 i ))|| 2 ] + γ 2 E[|| 1 N N i=1 S t i ⊙ ∇f i (w t-1 i )|| 2 ] ≤ γ 2 σ 2 k 2 M N + γ 2 E[|| 1 N N i=1 S t i ⊙ ∇f i (w t-1 i )|| 2 ]The last inequality comes from the derived Eq. 11. Next, we analyze the second term on the right-hand side of Eq. 12:E[ S t i ⊙ ∇f ( wt-1 ), wt -wt-1 ] = -γE[ S t i ⊙ ∇f ( wt-1 ), 1 N N i=1 S t i ⊙ G t i ] = -γE[ S t i ⊙ ∇f ( wt-1 ), 1 N N i=1 S t i ⊙ ∇f i (w t-1 i ) ] = - γ 2 E[||S t i ⊙ ∇f ( wt-1 )|| 2 + || 1 N N i=1 S t i ⊙ ∇f i (w t-1 i )|| 2 - ||S t i ⊙ ∇f ( wt-1 ) - 1 N N i=1 S t i ⊙ ∇f i (w t-1 i )|| 2 ] ≤ - γ 2 E[||S t i ⊙ ∇f ( wt-1 )|| 2 + || 1 N N i=1 S t i ⊙ ∇f i (w t-1 i )|| 2 - ||∇f ( wt-1 ) - 1 N N i=1 ∇f i (w t-1 i )|| 2 ]Further expanding the right-hand side of the above inequality, we obtain:E[||∇f ( wt-1 ) - 1 N N i=1 ∇f i (w t-1 i )|| 2 ] = E[|| 1 N N i=1 ∇f i ( wt-1 ) - 1 N N i=1 ∇f i (w t-1 i )|| 2 ] = 1 N 2 E[|| N i=1 (∇f i ( wt-1 ) -∇f i (w t-1 i ))|| 2 ] ≤ 1 N E[ N i=1 ||∇f i ( wt-1 ) -∇f i (w t-1 i )|| 2 ] ≤ L 2 N N i=1 E[|| wt-1 -w t-1 i || 2 ]In the above derivation, we have used the assumption of L-smoothness and Jensen's inequality. Next, we will continue to estimate the upper limit of this term. Assuming that the last parameter aggregation occurred at time t = t 0 , and the next aggregation will take place at t = t 0 + E, then:E[|| wt -w t i || 2 ] = E[||γ t τ =t0+1 1 N N i=1 S τ i ⊙ G τ i -γ t τ =t0+1 S τ i ⊙ G τ i || 2 ] = γ 2 E[|| t τ =t0+1 1 N N i=1 S τ i ⊙ G τ i - t τ =t0+1 S τ i ⊙ G τ i || 2 ] ≤ 2γ 2 E[|| t τ =t0+1 1 N N i=1 S τ i ⊙ G τ i || 2 + || t τ =t0+1 S τ i ⊙ G τ i || 2 ] ≤ 2(t -t 0 )γ 2 E[ t τ =t0+1 || 1 N N i=1 S τ i ⊙ G τ i || 2 + t τ =t0+1 ||S τ i ⊙ G τ i || 2 ] ≤ 2(t -t 0 )γ 2 E[ t τ =t0+1 N i=1 1 N ||S τ i ⊙ G τ i || 2 + t τ =t0+1 ||S τ i ⊙ G τ i || 2 ] ≤ 4(t -t 0 )γ 2 G 2 ≤ 4Eγ 2 G 2Substituting all the above inequalities into the right side of Eq. 12, we can finally obtain that when using a learning rate 0 ≤ γ ≤ 1 L , it satisfies:E[f ( wt )] ≤ E[f ( wt-1 )] - γ -γ 2 L 2 E[|| 1 N N i=1 S t i ⊙ ∇f i (w t-1 i )|| 2 - γ 2 E[||S t i ⊙ ∇f ( wt-1 )|| 2 + 2γ 3 E 2 G 2 L 2 + L 2N M γ 2 σ 2 k 2 ≤ E[f ( wt-1 )] - γ 2 E[||S t i ⊙ ∇f ( wt-1 )|| 2 + 2γ 3 E 2 G 2 L 2 + L 2N M γ 2 σ 2 k 2After rearranging the above inequalities, we obtainE[||S t i ⊙ ∇f ( wt-1 )|| 2 ≤ 2 γ (E[f ( wt-1 )] -E[f ( wt )]) + 4γ 2 E 2 G 2 L 2 + L N M γσ 2 k 2 (13)Finally, summing the inequalities from t = 1, • • • , T , and multiplying both sides by 1 T , we obtain:1 T T i=1 E[||S t i ⊙ ∇f ( wt-1 )|| 2 ≤ 2 γT (f ( w0 ) -f * ) + 4γ 2 E 2 G 2 L 2 + L N M γσ 2 k 2 (14)Selecting a learning rate γ =√ N M L √T , we obtain:1 T T i=1 E[||S t i ⊙ ∇f ( wt-1 )|| 2 ≤ 2L √ N M T (f ( w0 ) - f * ) + 4N M E 2 G 2 T + σ 2 k 2 √ N M T . Furthermore, by choosing E ≤ T 1/4(M N ) 3/4 , we can derive the following corollary:1 T T i=1 E[||S t i ⊙∇f ( wt-1 )|| 2 ≤ 2L √ N M T (f ( w0 )-f * )+ 4G 2 √ M N T + σ 2 k 2 √ N M T = O( 1 √ N M T). This proves the convergence rate of FedPart.For better visualizing the semantic information recognized by each layer in the different models, in Fig. 6 , we present representative results from the first and last layers of models under four scenarios: FedAvg-100, FedPart(No Init. 1C), FedPart(1C), and FedPart(5C).From the visualization results, it can be observed that FedAvg-100, due to being a full network update, captures low-level semantic features (such as clear boundaries) in shallow layers, while deeper layers capture complex semantic information. However, the results of FedPart(No Init. 1C) exhibit noticeable differences in color and structural features compared to the full network update. This confirms our belief that partial network updates are detrimental to establishing a hierarchical information extraction approach, resulting in the model converging to possible local minima. Additionally, we observe that by including the initial phase of full network updates and multiple rounds of sequential training, the similarity of semantic information obtained by the model gradually approaches that of FedAvg. Therefore, the results sufficiently demonstrate that although we only train one layer of the network each time, by employing an appropriate layer selection scheme, we ultimately achieve results comparable to those of full network updates.To visually depict the characteristics of the convolutional kernels in the first convolutional layer of different models, we conduct kernel visualization. The four models we select come from the following scenarios: FedAvg-100, FedPart(No Init. 1C), FedPart(1C), and FedPart(5C).In Fig. 7 , we present a comparison of results for planes in the first convolutional layer. It can be seen that the kernels in the first convolutional layer of the FedAvg-100 model are mostly edge and corner detectors. In contrast, the results of FedPart(No Init. 1C) and FedPart(1C) appear more random and irregular. However, after training to convergence, the results of FedPart(5C) are noticeably more similar to those of FedAvg-100, and start to exhibit characteristics of simple feature extractors. This indicates that through partial network updates, the layers of the model gradually coordinate with each other, yielding a cooperative effect. In this section, we conduct experiments with an α = 0.1 setting as data heterogeneity is more severe. The experimental results are shown in Table 12 .It can be seen that, in this extreme non-IID scenario (α = 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. In Assumption 3 in Section 3.3, we assume that for any mask matrices S 1 , S 2 , it holds that:EQUATIONRegarding the value of k on the right side of the equation above, recall Eq. 14 in Appendix B, the convergence rate of FedPart satisfies:1 T T i=1 E[||S t i ⊙ ∇f ( wt-1 )|| 2 ≤ 2 γT (f ( w0 ) -f * ) + 4γ 2 E 2 G 2 L 2 + L N M γσ 2 k 2Therefore, theoretically, the smaller the value of k, the smaller the value on the right side of this inequality, leading to improved convergence of Hence, it is important to carefully examine the value range of k in practice.We begin with analysing the lower bound of k. Since S 1 and S 2 are arbitrary, it is possible that S 1 = S 2 , indicating a lower bound of 1 for the value k. As for approximating the upper bound of the k, we conduct Monte Carlo simulations on real-world nueral networks.We test the k values in three neural networks at different training stages. For each neural network, we conduct Monte Carlo simulations to collect 10,000 samples to accurately approximate the value of k. The experimental results are shown in Table 13 . We can see that k is close to 1 under different settings, which proves that the effect of applying different masks to the variability of gradient is similar, thus strongly supporting Assumption 3. Justification: The following content of this paper is centered on the partial network updating method introduced in the abstract and introduction.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: This paper proposes a new approach for federated learning, where only parts of the networks are updated through communication and averaging between clients. The proposed method reduces the layer mismatch between locally learned models and the global model, which is a nice observation as pointed out by reviewers. Also, the experiment results are extensive  and strong, indicating that the proposed method is better than the full network FL baselines under multiple metrics.

Reviews: review: paper_summary: paper_summary: The paper proposes a new and novel method to partially train networks to achieve better training efficiency but also, in some cases, better performance.
strength_weakness: strength_weakness: As the authors claimed, while the experiments are sufficient - they are rather limited in terms of dataset sizes given the target applications. Further, synthetic data experiments are missing - I would have personally appreciated to see how FedPart performs in the setting of both IID and non-IID data.
questions: questions: Based on my comments above, I would like to ask the following questions,

 - How would the method perform on IID and non-IID data?
 - Would it be possible to add some experiments regarding the previously mentioned question?
limitations: limitations: The authors have sufficiently addressed limitations of this work.
paper_summary: paper_summary: This paper discovered the layer mismatch challenge in federated learning due to the full network update. To mitigate this challenge, the authors proposed the FedPart method. Specifically, the FedPart method would ask the clients do full network update in the beginning communication rounds, and then it would ask each client to do only train one layer in each round from the shallowest to the deepest. The experiment results indicate that the proposed method is better than the full network FL baselines.
strength_weakness: strength_weakness: 1. 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.

2. 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.

3. 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.

4. 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.

1. Zhang, Tuo et al. “TimelyFL: Heterogeneity-aware Asynchronous Federated Learning with Adaptive Partial Training.” 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW) (2023): 5064-5073.
2. Lee, Sunwoo et al. “Embracing Federated Learning: Enabling Weak Client Participation via Partial Model Training.” IEEE Transactions on Mobile Computing (2024): n. pag.
questions: questions: Please see the weakness above.
limitations: limitations: Please see the weakness above.
paper_summary: paper_summary: In this paper, the authors suggest a new approach for the network update step in federated learning. Considering that traditional federated averaging updates and aggregates all parameters at once, leading to a divergence between the global model and the local solution to a client’s specific task, they simply propose to only update parts of the network. They show that the layer mismatch between parameters learned locally and their corresponding global version is largely alleviated through this approach.

[Edit after rebuttal: raised score to accept]
strength_weakness: strength_weakness: There are two main weaknesses, (1) the suggested layer-wise update means that there is only one (linear) weight matrix that is changed between the global and local models through learning on the local data, which seemingly makes the models much *more* vulnerable to data reconstruction compared to an update of the full model. After all, we only need to reconstruct data from a change in a linear map, rather than a change of a deep, non-linear model. While the authors show that for some specific attacks that were tailored for deep model architectures the approach does fine, a simple attack based on the change in the single layer might be much more effective. At least a discussion on this matter would be important. (2) The motivation and initial results seem convincing, yet 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. This seems strange, the motivated layer mismatch should happen consistently in every round. Maybe this is a misunderstanding, a clarification is certainly needed.

There are several other issues raised in the Question below, however, I do believe the paper has value and I will (significantly) raise my score if the concerns are answered.
questions: questions: Major:

- In Figure 1, the update step size after aggregation seems extremely sharp. I was wondering if the momentum terms (or higher order statistics) of the optimizer were properly reset once the client received the global set of parameters? I briefly checked the submitted source code and it seems the reset of Adam or other optimizers is missing, maybe I overlooked it.
- In Section 3.3 could you please comment on the assumption and put into perspective how realistic they are? A brief discussion on the assumptions would be great – taking them as they are just because they are in the literature feels a bit rough.
-  Regarding the communication costs, how do clients know the mask $S^t$? It seems that this is missing from the analysis.
- In Table 6, the ablation on warm-up rounds, it seems that the larger (60 round) warm-up performs significantly better than what was reported before (in Table 2). How does table 2 look when run with longer warm-up? What is the reason not to do that?
- In the section on the warm-up ablation you mention that FedPart still enhances the model’s accuracy even after so many warm-up rounds (second to last sentence). Where can I see this in the presented numbers?
- According to the paper checklist you provide repetitions to assess experimental statistical significance. This is not given for Table 6-9, which is in my opinion fine for these ablations, but the checklist argument should hence be No. You can add in the description that you assessed the significance for the main experiments with an indication which one these are.
      
Minor:
      
- In Figure 2, it is a bit unclear to me why the layers have these different shapes, could you please clarify?
- The activation maximization approach is extremely old and yields arguably useless results. There is more than a decade of advancements in this field, with MACO [1], yielding great visual representations that would be much more useful in accessing the meaning.

[1] https://arxiv.org/abs/2306.06805
limitations: limitations: There is almost no discussion of current limitations of the method at hand. The manuscript would benefit from a slightly more critical discussion in the future.
paper_summary: paper_summary: The authors observe 'layer mismatch' phenomenon in federated learning, and propose FedPart that uses partial network updates to address the former issue. They show that FedPart outperforms previous methods and also reduced communication and computation overhead.
strength_weakness: strength_weakness: 1. The 'layer mismatch' term does not convince me. The authors try to use Figure 1(a) to show that the 'update stepsize' (to be clarified later) increases after each averaging and thus it indicates layer mismatch, I don't fully understand the logic here, in particular why does this draw the conclusion that it is caused by 'inadequate cooperation among layers'? I do see there is a comparison in Fig1(b), but on one hand it does not differ that much (except for the 50 to 60 iterations region), on the other hand the logic that this is due to 'layer mismatch' is still not perfectly sound even if Fig1(b) shows greater difference. Are there other evidences or maybe related observations in previous literature to support this claim?

2. The proposed strategy to select trainable layers does not provide much insight. Before getting to that section, I was expecting some unique or more informative way to choose what layer to update. It seems like the proposed strategy is the first thing anyone would try, thus I would like to see some argument/support for the effectiveness of choosing this strategy.

3. The experiment setup, if I understand correctly, involves 5 rounds of full network training between each cycle, which I assume is a lot given that partial update is about 8 or 18 rounds per cycle (for Resnet-8 and Resnet-18 respectively)? This is where I concern about the impact of these full network trainings. 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. 

4. Regarding communication cost, I might argue that most communication cost are overhead in each communication, which I believe sharing 1/8 of the number of parameters BUT requiring more communication rounds is not necessarily saving communication cost. And my intuition is that only updating one layer at a time increases the total number of communication rounds?
questions: questions: 1. What exactly is the update stepsize in Fig1? Is it the learning rate? If not, what is it? I might have missed the exact definition of it so I would appreciate if the authors can clarify it here.

2. The comparison in Fig1b really does not look like a big difference to me, for some iterations the partial update actually has a slightly bigger swings? And also I don't understand why Fig1a does not match Fig1b in terms of the blue curve, shouldn't they both be full network update and should be the same?
limitations: limitations: N/A

Comments: title: Read your rebuttal
Response rebuttal
comment: 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.
Dear Reviewer,

We greatly appreciate your positive and insightful feedback. As the first to introduce and explore the concept of "layer mismatch," we recognize that there is still room for improvement, especially in developing a deeper understanding of the underlying reasons behind this phenomenon. We are committed to further enhancing our work by strengthening the foundation of our findings with both theoretical and empirical evidence. Thank you once again for your feedback.
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 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.

| Experiment Configurations                    | CIFAR-10 | CIFAR-100 | Tiny-ImageNet |
| -------------------------------------------- | -------- | --------- | ------------- |
| I.I.D                                        | +1.9%    | +1.9%     | **+3.7%**     |
| Non-I.I.D. ($\alpha=1$)                      | +1.0%    | +2.2%     | **+3.4%**     |
| I.I.D. (client\_num=150, sampling\_rate=0.2) | +2.1%    | +1.6%     | **+3.4%**     |

**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 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 am not convinced by the author's response on W2.

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. 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 a implication that the proposed FedPart only works well on the easy dataset with iid setup. 

As a result, I would maintain my score.
Thank you for your thoughtful review and for taking the time to consider our clarifications. We greatly appreciate your positive feedback and helpful comments.
Thank you for the clarifications provided. Taking in account the rest of the reviews/responses thus far, I will be keeping my score as is. I wish the authors the best of luck with the final decision process.
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 your response.

I appreciate all answers, they clarify many details that strengthen the paper, and I would like to see these additional discussions (esp. on the assumptions and privacy) covered in the final version of the manuscript.

I raise my score to accept and wish the authors best of luck for the remaining review process.