FedLASE: Performance-Balanced System-Heterogeneous FL via Layer-Adaptive Submodel Extraction

Response to [Q1]: Clarification on why clients with fewer resources achieve better performance in Table 1.

We thank the reviewer for pointing out this interesting observation. We provide the following possible explanations based on our analysis and supporting experiments:

• Training dominance of resource-limited clients in client-number imbalanced scenarios.

In Table 1, the scenario reflects a setting where resource-rich clients are few, while resource-limited clients are the majority. Due to their small number, resource-rich clients have fewer training samples and lower participation probability (under uniform client sampling). As a result, the overall training process is dominated by resource-limited clients. This causes their average performance to appear higher, as shown in Table 1.

• Performance in balanced resource distribution.

To further support our hypothesis, we report results (shown in the table below) under the balanced resource setting in numbers {1,1/4,1/16,1/64}_{25,25,25,25}, following the same presentation style as Table 1, with all other experimental configurations remaining unchanged. In this case, resource-rich clients contribute more due to more frequent participation and larger models. Consequently, their local and global test accuracy is typically higher than that of lower-resource clients. This supports our understanding that the observed phenomenon in Table 1 arises from client number imbalance.

• Some remarks.

It is worth noting that this does not necessarily mean that the model with resource level 1 achieves the highest performance in the balanced resource setting in numbers. As presented by the related study about pruning, moderately pruned models usually can outperform unpruned ones. In our results, clients with resource level 1/4 often perform better than that of others. This observation suggests that appropriate sparsity may lead to better generalization, which is also one of the reasons why dropout has been widely adopted in practice.

• Not likely caused by overfitting.

We also considered whether the observed pattern could be caused by overfitting of large models. However, both local and global test accuracy curves across different resource levels show a consistent upward trend and plateauing behavior, rather than the degradation typically associated with overfitting. This suggests that overfitting is unlikely to be the main cause. Due to rebuttal space constraints, we omit the local and global test accuracy plots.

We hope this explanation clarifies the underlying reasons behind the performance trend observed in Table 1.

Response to [Q2]: Practical considerations on GPU systems.

We appreciate the reviewer’s insightful question. Compared with unstructured submodel extraction methods such as FIARSE, our method theoretically offers computational advantages. For structured submodel extraction methods like FedRolex and HeteroFL, under the same number of training rounds, they may achieve better memory usage and computation speed in practice on GPU systems due to better alignment with parallel hardware architectures. However, it is well-known that structured submodel extraction methods often suffer from lower model accuracy compared to unstructured approaches. Therefore, in real-world applications, the choice between structured and unstructured submodel extraction methods should be made based on the specific trade-offs between computational efficiency and model performance. In addition, studying unstructured pruning methods helps us gain a deeper understanding of the network training process, even if such methods may not offer advantages in terms of hardware execution.

Response to [Q3]: The significance of the new contribution.

We thank the reviewer for their comments. Below, we clarify our key contributions and explain why our method addresses challenges not considered in [21].

1. More realistic scenario considered.

Our work targets a more practical federated learning scenario where resource-constrained clients are the majority, while resource-rich clients are few. In contrast, [21] focuses on balanced resource levels across clients, which is less reflective of real-world situations. This new setup introduces important challenges for collaborative training.

2. Limitations of [21] in practical settings.

The FIARSE method in [21] selects parameters based on global importance ranking and allows each client to choose its own important elements. However, in our realistic setting:

• Submodel structure is easily broken: Resource-limited clients can only select a small number of parameters, and global ranking often ignores structural consistency across layers. This leads to fragmented submodels and degraded local training.

• Client collaboration becomes difficult: Resource-rich clients have limited training data and receive models with inconsistent structures. As a result, it becomes hard for them to benefit from collaboration with low-resource clients. This worsens client drift and slows down convergence.

3. Simple yet effective solution.

To address the above issues, we propose a layer-wise submodel extraction strategy. Instead of selecting parameters purely by importance, we control the selection within each layer to maintain model structure, which makes it (i) help preserve model architecture consistency, (ii) reduce divergence among clients, (iii) allow resource-rich clients to better learn from others, (iv) improve overall training performance in unbalanced scenarios.

4. Key insight and contribution.

Our main insight is that structural inconsistency in submodels is a major cause of training divergence in realistic federated learning. A purely greedy, importance-aware approach, like in [21], can be harmful when client resources are highly unbalanced. By considering both parameter importance and structural consistency, our method achieves better collaboration among clients. This is a necessary design consideration for future FL methods aimed at realistic, heterogeneous deployments.

Response to [Q1]: Clarification and quantitative evidence for "inter-layer discrepancies".

(1) Empirical observation of inter-layer discrepancies

FIARSE measures parameter importance based solely on parameter magnitude. To analyze layer-wise differences, we computed the maximum and mean of parameter magnitudes within each layer (reported in the first two rows of the table below, excluding the first and last layers of ResNet-18, normalization layers, and biases). As shown, the maximum parameter magnitudes differ significantly across layers, and several layers exhibit mean magnitudes that diverge substantially from others. Additionally, the number of parameters varies widely across layers. Ranking all parameters globally, as FIARSE does, overlooks both (i) inter-layer differences in parameter distributions and (ii) disparities in parameter counts. This can potentially remove critical information from important layers or disrupt the overall structural balance of the network.

(2) Training dynamics: layer-wise extraction and parameter updates

We further compared FIARSE and our proposed FedLASE in terms of layer extraction ratio dynamics and cumulative parameter update counts during training (reported in the last two rows of the table), and observed the following:

• FIARSE maintains largely stable extraction ratios across rounds, which likely leads to repeated updates of the same subset of parameters while many others remain insufficiently trained.

• Our method FedLASE, by incorporating layer structural information into submodel extraction, adaptively adjusts extraction ratios. The higher variance in extraction ratios across rounds allows more parameters to participate in training, resulting in a more balanced update distribution. This is further confirmed by our cumulative parameter update count analysis (visual comparisons omitted here due to rebuttal limitations).

(3) Performance validation

The superior performance of FedLASE over FIARSE further demonstrates the importance of accounting for inter-layer discrepancies. By dynamically balancing parameter training across layers, FedLASE achieves improved convergence and accuracy.

The experiments in the first two rows of this table are conducted under the system {1}{100}. mean_j|\theta^t{l_i,j} | and max_j|\theta^t_{l_i,j}| with t=599 represent the maximum and mean of parameter magnitudes within $l_i$-th layer in round t, respectively. 'ratio_std' denotes the standard deviation of submodel extraction ratios over the last 600 training rounds.

Notes: Unless otherwise specified, all experiments are conducted for CIFAR-100 under ResNet-18 using the following default settings: heterogeneous system is {1, 1/4,1/16,1/64}_{10,20,30,40}, learning rate is $0.01$, momentum is $0.8$, the training round is $T=1000$, and the data distribution is Dir(0.3). Other settings remain the same as described in the main paper.

Response to [Q2]: Design for the layer importance score.

• The use of mean: How to effectively aggregate parameter importance into a meaningful layer-wise importance score is a critical design choice. A few straightforward strategies include computing the mean (FedLASE), maximum (FedLASE w max), or assigning equal extraction ratios (FedLASE w eq) to all layers. Through extensive experiments, we found that using the average importance for the parameters within that layer provides a more stable and representative measure of that layer’s contribution. Some representative experimental results are shown in the following table.

• The use of logarithmic transformation: Once the preliminary layer importance scores are computed, a new challenge may be encountered: if a layer has a substantially larger number of parameters and correspondingly higher layer importance, it may dominate the submodel extraction process. As a result, layers with relatively lower total importance may end up with very few preserved parameters, potentially leading to structural imbalance and damage to the overall network architecture. To mitigate this issue, we apply a logarithmic transformation to smooth the importance scores across layers. This helps reduce the variance and prevents drastic parameter imbalance. The ablation results presented in the following table further demonstrate the effectiveness of this smoothing strategy and highlight the critical role of properly normalized layer importance scores.

• Other functions: Although linear normalization might seem simpler, it often lacks the desired smoothing effect, and its hyperparameter tuning can be non-trivial. Alternatively, learning a non-linear mapping function from data is a possible solution, but raises further challenges in terms of privacy leakage, dataset selection, and computational cost. Based on these observations, we adopt a log-based heuristic function, which is both simple and effective, to define layer importance. Designing more principled and adaptive layer scoring functions using richer information remains an important direction for future work.

Response to [Q3]: Additional experiments on NLP tasks.

Thank you for your insightful comment regarding the generalizability of our proposed method. We fully agree that evaluating on more diverse tasks and model architectures is crucial for demonstrating broader applicability. To address this concern, we extend our experiments to the AGNews dataset, a widely used large-scale text classification benchmark in the NLP domain. We adopt a pre-trained RoBERTa-based transformer model as the backbone. The training is conducted with $2$ local epochs, $300$ global rounds, a learning rate of $0.0001$, using the AdamW optimizer with momentum parameters of $(0.9, 0.95)$, and a batch size of $20$. We compare FedLASE with FIARSE, which is the most competitive baseline. As shown in the following table, FedLASE consistently outperforms FIARSE in this NLP setting as well, achieving both higher global accuracy and better client-level fairness. These results further validate the generalizability and effectiveness of FedLASE beyond small-scale image classification and convolutional architectures.

Dear Reviewer iCC5:

We hope this message finds you well. We would like to express our sincere gratitude for your insightful and constructive comments on our manuscript. As we approach the conclusion of the author-reviewer discussion period, we kindly request that you review our response and let us know if it adequately addresses your concerns. If you require any further information or clarification, please feel free to contact us. We would be glad to provide any additional details you may need.

Thank you once again for your time and valuable feedback.

Best regards, Authors

Response to [Q1] and [W1]: The significance of the new contribution.

We thank the reviewer for their comments. Below, we clarify our key contributions and explain why our method addresses challenges not considered in [21].

1. More realistic scenario considered.

Our work targets a more practical federated learning scenario where resource-constrained clients are the majority, while resource-rich clients are few. In contrast, [21] focuses on balanced resource levels across clients, which is less reflective of real-world situations. This new setup introduces important challenges for collaborative training.

2. Limitations of [21] in practical settings.

The FIARSE method in [21] selects parameters based on global importance ranking and allows each client to choose its own important elements. However, in our realistic setting:

• Submodel structure is easily broken: Resource-limited clients can only select a small number of parameters, and global ranking often ignores structural consistency across layers. This leads to fragmented submodels and degraded local training.

• Client collaboration becomes difficult: Resource-rich clients have limited training data and receive models with inconsistent structures. As a result, it becomes hard for them to benefit from collaboration with low-resource clients. This worsens client drift and slows down convergence.

3. Simple yet effective solution.

To address the above issues, we propose a layer-wise submodel extraction strategy. Instead of selecting parameters purely by importance, we control the selection within each layer to maintain model structure, which makes it (i) help preserve model architecture consistency, (ii) reduce divergence among clients, (iii) allow resource-rich clients to better learn from others, (iv) improve overall training performance in unbalanced scenarios.

4. Key insight and contribution.

Our main insight is that structural inconsistency in submodels is a major cause of training divergence in realistic federated learning. A purely greedy, importance-aware approach, like in [21], can be harmful when client resources are highly unbalanced. By considering both parameter importance and structural consistency, our method achieves better collaboration among clients. This is a necessary design consideration for future FL methods aimed at realistic, heterogeneous deployments.

Response to [Q2] and [W3]: Comparison with [22], [23], [13], [24].

We appreciate the reviewer’s comment regarding comparisons with related works addressing system heterogeneity. We provide detailed clarifications for each of these methods below:

• Depth [13]: Depth extracts submodels purely based on network depth. However, in realistic cross-device scenarios with a large proportion of resource-constrained clients (e.g., minimum resource level 1/64), depth-based truncation fails to generate valid submodels under such tight constraints. For instance, truncating ResNet-18 by depth alone cannot produce functional submodels small enough to fit these devices while retaining meaningful structure. In contrast, our method is more flexible and adapts submodel extraction to diverse resource profiles, making it better suited for real-world settings dominated by resource-limited clients.

• FjORD [24]: FIARSE has already demonstrated superior performance over FjORD [24], as reported in [21]. Building on FIARSE, our method further achieves significant improvements across diverse heterogeneous scenarios, demonstrating clear advancements beyond [21]. Additionally, FjORD did not release official code, and its implementation details are insufficiently specified, making faithful reproduction challenging and potentially unreliable. Therefore, we focus on FIARSE—an importance-aware submodel extraction baseline with publicly available code—as the most relevant and fair baseline for empirical comparison.

• FedLMT [23]: FedLMT uses low-rank decomposition, but its applicable resource range is constrained by architectural limits. For example, with ResNet-18, the lowest feasible rank ratio is $\alpha \geq 0.015625$. As noted in [23], for a resource level of 1/4, FedLMT requires $\alpha = 0.03$; consequently, for 1/16 resources, $\alpha$ must be less than 0.0075, which contradicts the condition $\alpha \geq 0.015625$. This limitation prevents FedLMT from accommodating extreme yet realistic heterogeneous scenarios, such as when devices differ widely in computational capacity (e.g., 1 vs. 1/64). Moreover, enforcing a rank ratio at the limit severely constrains model capacity, undermining its training effectiveness. Due to these constraints, FedLMT is unsuitable for the highly resource-imbalanced scenarios targeted in our work, and direct comparison would be neither practical nor meaningful.

• FedDSE [22]: Since the code for this work was not publicly available, we did not include it in our initial comparisons. Recently, we obtained the code from the authors and conducted experiments under the data heterogeneity setting L=4 as considered in their paper. The results shown below illustrate that our method significantly outperforms FedDSE across various system heterogeneity scenarios, demonstrating the superiority of our layer-adaptive submodel extraction approach. Moreover, our method can also be combined with FedDSE to design more effective layer pruning ratios for FedDSE.

• Justification of our experimental focus on FIARSE[21]: Given these considerations, we focus on FIARSE, which (i) is reproducible, (ii) is recognized as a strong and recent baseline for importance-aware submodel extraction, and (iii) directly aligns with our method’s design goals. Importantly, our method consistently outperforms FIARSE across various heterogeneous settings, demonstrating substantial empirical gains. This provides strong evidence of our contribution even without direct comparisons to methods that are either infeasible or not reproducible in our targeted scenarios.

We will explicitly add these clarifications in the revised manuscript and, where feasible, provide supplementary experiments or discussion to better position our method relative to these works.

Response to [Q3] and [W2]: Additional experiments on NLP tasks.

Thank you for your insightful comment regarding the generalizability of our proposed method. We fully agree that evaluating on more diverse tasks and model architectures is crucial for demonstrating broader applicability. To address this concern, we extend our experiments to the AGNews dataset, a widely used large-scale text classification benchmark in the NLP domain. We adopt a pre-trained RoBERTa-based transformer model as the backbone. The training is conducted with $2$ local epochs, $300$ global rounds, a learning rate of $0.0001$, using the AdamW optimizer with momentum parameters of $(0.9, 0.95)$, and a batch size of $20$. We compare FedLASE with FIARSE, which is the most competitive baseline. As shown in the following table, FedLASE consistently outperforms FIARSE in this NLP setting as well, achieving both higher global accuracy and better client-level fairness. These results further validate the generalizability and effectiveness of FedLASE beyond small-scale image classification and convolutional architectures.

Notes: All experiments are conducted for CIFAR-100 under ResNet-18 using the following default settings: learning rate is $0.01$, momentum is $0.8$, the training round is $T=500$. Other settings remain the same as described in the main paper.

Response to [Limitations]

Thank you for the valuable suggestion. We agree that discussing the limitations in the main paper will provide clearer insights and benefit readers’ understanding. We will revise the manuscript to move the discussion of limitations from the appendix to the main text in the next version.

Thanks the authors for the response and the additional experiments results. The efforts are highly appreciated.

However, my following concerns remain:

1. Significance of the contribution compared to [21]. While the authors provide a thorough comparison against [21], the fact remains that the main proposal is an extension of [21]. Essentially, this paper focuses on a narrower niche of system heterogeneity—specifically, more “skewed” resources, with a larger proportion of constrained, small devices. In practical deployments, it is not entirely clear how common such a significantly “skewed” environment is.

In addition, prior work has studied FL under skewed system heterogeneity [1]. Why not simply using these techniques?

[1] Zhang, Jiayun, et al. "How few davids improve one goliath: Federated learning in resource-skewed edge computing environments." Proceedings of the ACM Web Conference 2024. 2024.

2. Lack of comparison. While the authors provide more discussions and results, the paper would be stronger and more convincing with a comprehensive comparison.

Response to [Q1] and [W2]: Proxy for parameter importance.

We sincerely thank the reviewer for raising all these insightful and interesting questions. Using the parameter magnitude as a proxy for parameter importance has been widely adopted in the area of pruning for both centralized and federated learning, due to its simplicity and practical effectiveness.

• Theoretically, this choice can be motivated through a second-order Taylor expansion of the global loss function in federated learning:

$F(\theta + \Delta\theta) \approx F(\theta) + \nabla F(\theta)^T \Delta\theta + \frac{1}{2} \Delta\theta^T H \Delta\theta$

with $H$ being the corresponding Hessian matrix. In typical federated optimization settings, local models are assumed to converge reasonably well before aggregation, implying $\nabla F(\theta) \approx 0$. Then, the loss difference is dominated by $\Delta\theta^T H \Delta\theta$. If we further assume $H \approx I$, then the impact of zeroing the $i$-th parameter $\theta_i$ can be approximated as $\theta_i^2$, suggesting that parameters with smaller magnitude usually have less contribution to the loss and can be safely pruned.

• Empirically, we adopt the same strategy as FIARSE [21], which also uses post-aggregation parameter magnitude to guide submodel construction under heterogeneity. Their results, along with our experiments, demonstrate that this simple metric is robust across many settings, even though it may not be exact.

Response to [Q2] and [W1]: Reason for selecting logarithmic layer normalization.

• Reason: After the predefined layer importance scores are computed, a new challenge may be encountered: if a layer has a substantially larger number of parameters and correspondingly higher layer importance, it may dominate the submodel extraction process. As a result, layers with relatively lower total importance may end up with very few preserved parameters, potentially leading to structural imbalance and damage to the overall network architecture. To mitigate this issue, we apply a logarithmic transformation to smooth the importance scores across layers. This helps reduce the variance and prevents drastic parameter imbalance.

• Ablation study: The ablation results shown in the following response further demonstrate the effectiveness of this smoothing strategy and highlight the critical role of properly normalized layer importance scores.

• Other selection: Although linear normalization might seem simpler, it often lacks the desired smoothing effect, and its hyperparameter tuning can be non-trivial. Alternatively, learning a non-linear mapping function from data is a possible solution, but raises further challenges in terms of privacy leakage, dataset selection, and computational cost. Based on these observations, we adopt a log-based function, which is both simple and effective. Designing more principled and adaptive layer scoring functions using richer information remains an important direction for future work.

Response to [Q3] and [W3]: Ablation study.

To better understand the contributions of different components in FedLASE, we conducted additional experiments that isolate the effects of: (i) STE-based gradient approximation, (ii) personalized overlapping aggregation, (iii) full retention of specific layers, and (iv) logarithmic layer normalization. Specifically, we evaluate the following variants:

• FedLASE w/o STE: removes the STE trick used during submodel training.
• FedLASE w/o overlap: replaces the personalized aggregation (which uses overlapping weighted averaging of masked parameters) with a FedAvg-style update: $\theta = \frac{1}{|\mathcal{A}|} \sum_{n \in \mathcal{A}} M^n \odot \theta^n$.
• FedLASE w all layers: extracts submodel extraction across all layers (excluding batch normalization and bias terms), instead of selectively preserving key layers.
• FedLASE w/o log: disables the logarithmic layer normalization when computing layer importance.

Notably, for fair comparison with FIARSE [21], we retain STE and overlapping aggregation in our default setup. This allows us to isolate the impact of our proposed layer-aware extraction scheme. We can see from the table below that:

• Removing logarithmic normalization will bring a large performance drop, yet 'FedLASE w/o log' still significantly outperforms FIARSE, demonstrating the importance of processing layer-level statistics beyond raw parameter magnitudes.
• 'FedLASE w all' shows slightly lower performance than our default version, suggesting that retaining the first and last layers intact indeed helps preserve critical information necessary for stable training and final prediction. Of course, we can also fix the extraction ratios of the first and last layers. However, to avoid introducing additional hyperparameters, we retain all parameters in these layers.
• The other two variants (w/o STE, w/o overlap) exhibit similar trends to those for FIARSE, reaffirming that our improvements stem from our core contribution, i.e., the sample and effective layer-adaptive submodel extraction framework for heterogeneous FL.

Notes: Unless otherwise specified, all experiments are conducted under the following default settings: heterogeneous system is {1, 1/4,1/16,1/64}_{10,20,30,40}, learning rate is $0.01$, momentum is $0.8$, the training round is $T=1000$, and the data distribution is Dir(0.3). Other settings remain the same as described in the main paper.

Response to [Q4] and [W4]: Impact of random.

Thank you for pointing this out. All results in Tables 1 and 2 are averaged over three different random seeds. We agree that it is important to understand the robustness of the proposed method under different sources of randomness, such as client sampling, data partitioning, and initialization. Due to space constraints in the rebuttal, we evaluate FedLASE and the baselines under more independent random seeds to more clearly demonstrate the impact of randomness. For each method, we report the global test accuracy under each seed, along with the corresponding mean and standard deviation. The results show that while there is some variance due to randomness, our method consistently outperforms FIARSE and other baselines across all seeds. This confirms the stability and robustness of FedLASE under realistic federated conditions.

Response to [Q5] and [W5]: Scenarios for online client dropouts.

• Results in Tables 1 and 2: We thank the reviewer for raising the important issue of online client availability, which is indeed a key challenge in real-world federated learning systems. In our original submission, we have already included experiments simulating typical client dropout scenarios: specifically, we considered 100 clients, with only 10% randomly selected in each round to simulate partial availability. These results are shown in Tables 1 and 2 and demonstrate the effectiveness of FedLASE. To further explore more realistic and complex system dynamics, we conduct two additional experiments to reflect heterogeneous and time-varying participation:
• Static scenario: We increase the total number of clients to 200 and randomly select 10% (i.e., 20 clients) to participate in each round.
• Dynamic scenario: We again use 200 clients, but the client participation ratio is dynamically sampled from {10%, 15%, 20%} in each round, simulating more complex real-world deployment scenarios with fluctuating client participation. The results in the following table illustrate that our method still consistently outperforms all baselines across both static and dynamic participation settings, highlighting its robustness to client dropout and its adaptability to more practical FL environments.

Response to [Limitations]

Although our proposed strategy is simple and effective, designing better alternatives to the logarithmic layer normalization function remains a challenging and open problem. Moreover, incorporating system heterogeneity information to design more fine-grained and adaptive metrics for parameter and layer importance is a promising direction for future research. We will include these limitations in the revised manuscript.

Dear Reviewer 1hNE:

We hope this message finds you well. We would like to express our sincere gratitude for your insightful and constructive comments on our manuscript. As we approach the conclusion of the author-reviewer discussion period, we kindly request that you review our response and let us know if it adequately addresses your concerns. If you require any further information or clarification, please feel free to contact us. We would be glad to provide any additional details you may need.

Thank you once again for your time and valuable feedback.

Best regards,

Authors