FedEL: Federated Elastic Learning for Heterogeneous Devices

We appreciate the reviewer’s time and effort in providing feedback, though it appears that the comments may have been intended for a different submission. We welcome any relevant insights they can share regarding our paper.

As the Area Chair updated the comments after the rebuttal deadline, we response them in the official comment.

W1: New Coming Devices

FedEL supports new devices through two simple strategies:

A. Offline Profiling (Default): New devices first profile tensor computation time locally using the received model. This takes a few local epochs. After profiling, they join training using the selective strategy. This avoids disruption and ensures fairness.

B. Online Profiling (Optional): If immediate participation is preferred, the device begins training with the full model while collecting profiling data. After a few epochs, it switches to selective training. This introduces minor overhead but avoids delaying training rounds.

W2: Memory and Power Consumption

Figures 8 and 9 initially showed only mean memory and power use. To improve transparency, we now report mean and variance over five runs.

FedEL shows the lowest memory and energy use among all baselines, with low variance. This reflects its robustness under runtime fluctuations, achieved by selective training that reduces memory footprint and shortens training—saving energy overall.

W3: Auxiliary Heads and Overhead

The number of auxiliary heads in FedEL is limited due to our block partitioning strategy, which is governed by a per-round training time threshold. This threshold constrains how many blocks can be considered during training and thus limits the number of early exit heads.

Each auxiliary head is a lightweight fully connected layer that maps the flattened output of its associated block to the prediction space. Since feature lengths may vary between blocks, each head is customized to the output shape of its corresponding block.

Importantly, these auxiliary heads do not add training-time overhead. Our selective training process proceeds in two steps:

1. Sliding Window Selection: The window defines a candidate region that includes all tensors in the current block and its auxiliary head.

2. Tensor Optimization: The optimization procedure (Problem (1) in our paper) then selects the most valuable tensors from this region, ensuring the total backward time remains within the training threshold.

As a result, FedEL maintains per-round training time within the original limit, regardless of the number of auxiliary heads. These heads are sparsely inserted, lightweight, and only activated when their corresponding block is the last one in the sliding window.

W4: Transformer Compatibility

As detailed in Section B of the Supplementary Material, we implemented FedEL on ALBERT, a compact Transformer model, using the block decomposition strategy from [39]. Auxiliary heads are placed after each encoder–classifier stack to support variable-depth execution—this setup integrates seamlessly with our sliding window mechanism for dynamic tensor selection.

We handle Transformer-specific traits as follows:

1. Shared Embeddings: Treated as a single shared tensor. We trace gradients from output layers back to the embedding layer to profile its backward cost. The weight update is counted once per round, regardless of reuse across layers.

2. Multi-head Attention: Each projection (query, key, value, output) is treated as a separate tensor for profiling. If weights are shared across layers, we aggregate their gradient cost and update them only once per round.

FedEL is architecture-agnostic. It only requires: Per-tensor backward cost profiling, Tensor selection, and Fine-grained control to freeze or update tensors.

These apply to both CNNs and Transformers. We also evaluated FedEL on MobileBERT and ViT-Tiny, showing improved accuracy and reduced training time:

W5: Dynamic Participation

Our experiments already include partial client participation—a subset of clients is randomly sampled each round, following common FL practice (FedAvg, HeteroFL, FIARSE).

We acknowledge this was not stated clearly and will revise the Experimental Setup section.

Crucially, FedEL supports dynamic participation by design. Each client performs local profiling and selection independently, with no need for global sync. Even if a client drops out or rejoins, FedEL functions seamlessly. The reported results already reflect this, demonstrating FedEL’s robustness to client heterogeneity and availability.

W1: More Related Works

We thank the reviewer for highlighting this important point. We have expanded our comparison with the cited works and summarize the key distinctions in the table below:

[1] NeurIPS 2023 offers convergence guarantees but applies global pruning without local data or neuron importance awareness, limiting its adaptability at the client level.

[2] TNNLS 2022 supports distributed pruning but incurs significant communication overhead due to the need to transmit importance scores and pruning decisions each round.

[3] DAC 2021 (Helios) enables dynamic model adaptation but requires frequent communication and additional optimization overhead, especially for straggler mitigation.

[4] TMC 2024 does not perform partial training; instead, it trains the full model and uploads only a subset of parameters, which is orthogonal to our focus on efficient local training.

In contrast, FedEL performs lightweight, local tensor selection using an importance-based mechanism, incurs no extra communication, supports fine-grained dynamic adaptation, and avoids heavy optimization—offering a practical and scalable solution to device heterogeneity.

We will add these to the related works in the final paper.

W2: Model Partition

We thank the reviewer for the insightful question. Our approach incorporates two key modules: (1) sliding window and (2) important tensor selection, which together enable multi-granularity model partitioning.

1. Sliding Window (Layer-Level Slicing): This is a coarse-grained strategy that slices the model at specific layers. As detailed in Section 4.1, only semantically meaningful layers (e.g., block or stage boundaries) are eligible slice points—not all layers. This ensures the model remains structurally sound and stable during training. The main goals of this module are to (i) enable compatibility with early-exit classifiers and (ii) limit the search space for finer-grained selection.

2. Important Tensor Selection (Width-Level Slicing): Within each layer or block selected by the sliding window, we further apply fine-grained slicing by selecting only important tensors for training. This allows each client to update only a subset of parameters, making the training process more efficient and adaptive to device capabilities.

3. Transformer Models: As noted in Section B of the supplementary material, we adopt the block design from the Albert model [39], placing early-exit classifiers after each encoder–classifier block (see [39], Figure 1). The sliding window works at the block level, enabling depth-wise flexibility, while tensor selection operates within each block to control the width of updates.

W3: Offline Tensor Time Profiling

We thank the reviewer for the comment. As shown in the table below, our empirical results across diverse neural network architectures indicate that profiling performed over a single training epoch incurs less than 1% additional overhead, which aligns with observations in ElasticTrainer. This low cost enables us to periodically re-profile the system if needed—e.g., when data size or background workloads change—thus addressing concerns about time variance at the block level.

W4: Dynamic Participation and Handling Stragglers

We thank the reviewer for highlighting this important point.

First, our experiments use partial client participation, where a random subset of clients is selected in each round, based on availability—not device performance. This reflects standard FL practice and aligns with methods like DepthFL, HeteroFL, and FIARSE. FedEL runs on this selected subset each round. We will revise the main paper to make this clearer.

Second, we acknowledge that in cases of extreme heterogeneity, some straggler clients may contribute fewer updates, possibly under-training some tensors (as discussed in Section C of the supplementary material). To address this, one promising direction is client clustering—grouping clients by similar compute capacities. Each group can run FedEL independently, and asynchronous updates between groups can help balance training and reduce the impact of extreme stragglers.

W5: Implementation of Block-Based Training

We appreciate the reviewer’s question. To enable training only on selected blocks, we use a simple and efficient layer-level masking approach based on the model structure.

For each local training round, we create a binary mask over model layers: Layers in the selected blocks get a mask value of 1. All other layers get a mask value of 0.

In PyTorch, we apply this by setting each layer’s requires_grad attribute: If mask = 1: requires_grad = True, so the layer is trainable. If mask = 0: requires_grad = False, so the layer is frozen.

This approach avoids any changes to the model architecture or complex control logic. It is lightweight, flexible, and works well with our sliding window and tensor importance selection modules.

W6: How to Select Blocks?

We thank the reviewer for this insightful question. Currently, the block selection in our method is not formulated as an explicit optimization problem. Instead, we use a practical heuristic guided by two components: offline tensor time profiling and a sliding window mechanism. Specifically, the model is first partitioned into blocks based on its computational structure, as determined through offline profiling. During each training round, a sliding window is used to define the active set of blocks: A new block is added to the window, and previous blocks may be removed based on their historical contribution to training effectiveness. This defines a dynamic candidate region, within which important tensors are then selected for training using our fine-grained selection module.

While this heuristic works well in practice (as supported by our experimental results), we agree that formulating block selection as an optimization problem is a compelling direction for future research. One potential formulation could be derived from the term $O_1$ in our convergence analysis, enabling principled selection of blocks under resource or latency constraints.

W7: Overhead of Adjusting Tensor Importance

We appreciate the reviewer’s concern about potential overhead. As shown in the table below, the total time spent on adjusting tensor importance—including the sliding window update, importance scoring and adjustment, and tensor selection—is minimal:

This represents only 2.4% of the total round time, which we consider negligible in the overall context of federated training. The lightweight design of our mechanism ensures scalability and practicality without compromising system efficiency.

W1: Model and Datasets

We appreciate the reviewer’s comment regarding the scale and recency of the models and datasets used in our evaluation. Our choice of models (e.g., VGG16, ResNet50, Albert) and datasets (image classification, speech recognition, lightweight NLP) was guided by two primary considerations: (i) alignment with prior work in heterogeneous federated learning (e.g., DepthFL, HeteroFL, and FIARSE), and (ii) practical deployment on real edge devices (Jetson Orin, Xavier), where larger models such as full-scale Transformers or multimodal architectures are currently infeasible due to hardware constraints. This setup enables fair comparisons while maintaining realism in resource-constrained federated environments.

That said, we agree that demonstrating FedEL’s scalability to larger, more recent models is an important next step. We are currently extending it to support larger vision and language models, MobileBERT on the GLUE benchmark and ViT-tiny on the NICO++ dataset, running on more capable hardware platforms (Jetson AGX). Results below indicate that FedEL continues to provide consistent performance benefits over FedAvg, even as model complexity and task diversity increase. These findings validate the generalizability of our selective training framework beyond small-scale settings.

W2: Overhead of Profiling and Tensor Selection

We appreciate the reviewer’s attention to overhead analysis. As shown in the first table, profiling during a single training epoch adds less than 1% overhead across all tested models. This step is done offline, before FL training begins, and its results are reused, so it introduces no runtime overhead during training.

For the tensor selection step during each FL round, we measured a small runtime cost—only 2.4% of the total round time, as shown below. Communication also accounts for just 2.7%. These results confirm that FedEL has minimal overhead and is efficient enough for real-world use.

W3: Adapting to Device Dynamics

We thank the reviewer for the valuable feedback. When runtime dynamics occur, profiling can be updated online with negligible cost (less than 1% overhead per epoch). This allows adaptive updates without impacting training efficiency. Dynamic re-profiling is a potential direction for future work.

W4: Partial Client Participation

We thank the reviewer for pointing out this clarification need. Our experiments do adopt partial client participation, which reflects realistic FL settings. Specifically, in each communication round, a random subset of clients is selected to participate, consistent with common practice in prior works such as FedAvg, HeteroFL, and FIARSE.

We will revise the manuscript to explicitly clarify this in the experimental setup. Additionally, FedEL's design—especially its local tensor selection and sliding window mechanism—naturally supports partial participation, as it makes no assumptions about client availability or full participation. Our reported results already reflect performance under this dynamic client selection.

W5: Uniform Network Channel

We thank the reviewer for raising this important point. We isolate computational heterogeneity to evaluate our method fairly, following practices in HeteroFL and DepthFL. However, our approach naturally extends to communication heterogeneity.

During offline profiling, we estimate tensor sizes and include communication time in the selection objective:

$\max_{\boldsymbol{A}} \boldsymbol{A} \cdot \boldsymbol{I} ~~s.t. T_{fw} + T_{bw}(\boldsymbol{A}) + T_{tx}(\boldsymbol{A}) \le T_{th}$

Here, $T_{tx}(\boldsymbol{A})$ represents the estimated communication time of selected tensors. We conducted additional experiments simulating four bandwidth tiers (10, 20, 30, 45 Mbps) across 100 clients. We compared: (1) FedAvg (baseline), (2) FedEL (original), and (3) FedEL+Tx (FedEL with communication-aware tensor selection).

Results show FedEL+Tx achieves similar or better accuracy while further reducing total training time, confirming the benefit of modeling communication heterogeneity:

W6: Handling Shared-Weight Layers

We thank the reviewer for this insightful question. As detailed in Section B of our Supplementary Material, our implementation for Albert follows the block design from [39], placing early exits after each encoder-classifier stack. This supports variable-depth execution and is fully compatible with our sliding window mechanism.

Regarding shared-weight layers, such as embeddings and grouped attention heads, we handle them as follows:

(1) Shared Embeddings: We profile the shared tensor as a single unit. Backward time is traced from the output back to the shared embedding, reflecting its impact on both input and output. The weight update time is counted only once since the tensor is updated once per iteration.

(2) Grouped Attention Heads: For attention layers with multiple trainable projections (query, key, value, output), we treat each as a separate tensor and profile their backward time individually. If some projections are shared, we track their computation cost jointly during backpropagation.

W7: Compatibility with DP and Secure Aggregation

Thank you for your comment. Our method is fully compatible with privacy-preserving techniques like Differential Privacy (DP) and Secure Aggregation (SA), as FedEL operates at the system level and does not modify or interfere with model encryption or privatization.

To evaluate this, we extended our experiments by integrating standard DP and SA into FedEL. The table below shows that both FedEL+DP and FedEL+SA achieve similar accuracy to vanilla FedEL. As expected, they incur some extra training time due to the added noise (DP) or encryption (SA).

W8: Jetson Power Modes

We thank the reviewer for this insightful suggestion. In our experiments, Jetson devices ran in MAXN mode to ensure consistency. Inspired by reviewer feedback, we reran experiments under varied power settings (10W, 15W, MAXN). We measured training times under each mode and used this variability in the FedEL scheduling process. As shown in the table below, FedEL still outperforms FedAvg in both accuracy and training time. These results confirm that FedEL remains effective even in more diverse and realistic hardware environments.

W9: Related Works

We appreciate the reviewer’s suggestions and discuss two related works below. We will include both in the final version.

[a] Recurrent Early Exits (ReeFL): ReeFL enables early exits in LLMs during FL. Our method FedEL can work with ReeFL by applying tensor selection inside each early-exit block. We combined both by inserting exits into transformer layers and applying FedEL within each. As shown below, FedEL+ReeFL reduces training time while maintaining accuracy, confirming that the two approaches are complementary.

[b] ScaleFL: ScaleFL adapts model width/depth based on device profiles, but uses static selection via a meta-scheduler. In contrast, FedEL dynamically selects tensors in each round based on importance, allowing finer control and better adaptability.

To validate this, we compared FedEL and ScaleFL across three tasks. FedEL consistently achieves higher accuracy, as shown below.

W10: Result Robustness

Thank you for the comment. We did evaluate robustness and reported it in Section B.5 of the Supplementary Material. As shown in Figure 21 (box plot), each experiment was repeated five times. The results show that FedEL consistently achieves stable and higher accuracy than the baselines, with narrow confidence intervals. This confirms the reliability and robustness of our method across multiple runs.

Q1: Partition rate

We thank the reviewer for the suggestion. To evaluate the impact of lower participation, we conducted experiments with a 10% participation rate. As shown below, all methods experienced slower convergence, leading to longer training time while maintaining similar accuracy. However, FedEL consistently achieves the highest accuracy and best efficiency across all tasks, confirming its advantage even under sparse client participation.

We will include these additional results and discussion on participation rate impact in our final submission.

Partition rate = 25%

Partition rate = 10%

Q2: Scalability Constraint

We thank the reviewer for the helpful comment. While the dynamic programming (DP) algorithm in ElasticTrainer [13] does have quadratic complexity based on the number of tensors $N$, it’s important to understand this in the training context.

The DP algorithm works on tensors, not on each trainable parameter. For example, ResNet has over 23 million trainable parameters, but only 214 tensors ($N$). So, the DP algorithm has a complexity of $O(214^2)$, while training the model takes at least $O(23\text{M})$. The difference becomes even more significant in modern models, for instance, ViT has 327 tensors versus 22 billion parameters. This means that the time spent on tensor selection is very small compared to the time spent on training the model.

In practice, tensor selection only adds about 2.4% to each training round but helps speed up training by as much as 3.87× compared to FedAvg.

We will clearly mention this scalability issue and its impact in our final paper. We truly appreciate the reviewer’s thoughtful feedback and it was a pleasure discussing these research ideas with the reviewer.

Q3: Differential Privacy

We thank the reviewer for the valuable feedback. While differential privacy (DP) is not the central focus of our work, we appreciate the opportunity to evaluate its impact. To examine the effect of the privacy parameter ε, we conducted additional experiments using ε = 30, 5, and 2. As shown in the table below, lowering ε enhances privacy by introducing more noise, but this results in reduced model accuracy. This illustrates the inherent trade-off between privacy and utility. We plan to explore the use of tighter privacy budgets and their associated trade-offs in future work.

W1: Large-Scale Simulation Setup

Thank you for your comment. We acknowledge that the term “random allocation” may have been unclear. In our 100-client simulations, we uniformly assign 25 clients to each of four speed tiers, with clients randomly selected within each tier—consistent with DepthFL, HeteroFL, and FIARSE. We will include the client partitioning scripts in the final submission.

W2: Communication Heterogeneity

We thank the reviewer for this insightful comment. Our study focuses on computational heterogeneity, a dominant bottleneck in mobile edge environments. In practice, training time on mobile devices far exceeds communication time, especially with modern 5G/WiFi networks. For example, ResNet-50 (~97.7 MB) takes 0.28–1.3 minutes to transmit over 10-45 Mbps uplink (e.g., AT&T, 2024), while training on an NVIDIA Jetson Xavier takes ~38.3 minutes, making communication relatively negligible in our setting.

Nonetheless, our method can incorporate communication heterogeneity. During offline profiling, we estimate tensor sizes. If client bandwidth is known, the tensor selection can be extended to:

$\max_{\boldsymbol{A}} \boldsymbol{A} \cdot \boldsymbol{I} ~~ s.t. T_{fw} + T_{bw}(\boldsymbol{A}) + T_{tx}(\boldsymbol{A}) \le T_{th}$

where $T_{tx}(A)$ estimates transmission time. This enables FedEL to jointly optimize computation and communication under a unified latency constraint.

To validate this, we simulated 100 clients across four bandwidth tiers (10, 20, 30, 45 Mbps, 25 clients each). We compared: (1) FedAvg, (2) FedEL, and (3) FedEL+Tx (our method with transmission-aware selection). FedEL+Tx consistently achieved comparable accuracy to FedEL with lower total training time, confirming its effectiveness.

W3: $\beta$ and $T_{th}$

Thank you for your comment. We fixed $\beta$ and $T_{th}$ in our study and performed ablation (Sec. 5.3) to assess their impact. $\beta$ balances global vs. local tensor sensitivity and affects accuracy. $T_{th}$ controls the time budget per round; larger values allow more updates, but increase computation cost.

Deriving explicit analytical relationships between these parameters and accuracy/convergence is challenging due to the coupled effects of training and communication dynamics. That said, we agree adaptive tuning, e.g., via reinforcement learning or bandit control, offers promising potential, which we plan to explore in future work.

W4: Simulation Heterogeneity

We appreciate the reviewer’s thoughtful observation. In this work, we focus on computational heterogeneity, as it is a major bottleneck in mobile FL and allows us to clearly study its impact. To ensure fair comparison, we also follow the setup used in prior work that only considers device speed differences.

We agree that communication, memory, and energy are also important. However, modeling all types of heterogeneity at once would make the problem more complex and shift the focus away from our main contribution.

Our method is flexible and can be extended. As shown in our response to W2, we can include communication costs in the tensor selection process. Similarly, memory and energy limits can be added by updating the profiling step. Early results show our method still works well when communication is considered.

W5: Tasks and Architectures

We thank the reviewer for highlighting this important point. Our current experiments are focused on three representative FL tasks—image classification (VGG16), speech recognition (ResNet50), and lightweight NLP (Albert)—which were selected to align with realistic use cases of FL on mobile and edge devices. These models and tasks are commonly adopted in prior work (e.g., DepthFL, HeteroFL, and FIARSE), ensuring fair benchmarking under resource-constrained settings.

We fully agree that evaluating on larger-scale architectures (e.g., full-size Transformers), graph neural networks (GNNs), and multimodal tasks is a valuable direction. However, such extensions often require considerable training resources and pose challenges for practical FL deployment at the edge. That said, the underlying mechanisms in FedEL—namely sliding-window block scheduling and dynamic tensor selection—are model-agnostic and can be adapted to other architectures. For instance, our supplementary material (Section B) demonstrates how FedEL is applied to Transformer blocks in NLP. We plan to extend our evaluation to more complex tasks and models (e.g., ViTs, GNNs, and vision-language models) in future work.

W6: Communication and System Overhead

Thank you for your comment. To clarify, our method does not add extra communication cost—in fact, it reduces it compared to FedAvg. This is because FedEL only uploads selected important tensors rather than the full model. As shown in the table, FedEL results in: 1. Lower communication time per round than FedAvg. 2. Communication taking up only a small part of the total training time.

We also measured the runtime of FedEL's system modules (e.g., sliding window, tensor importance update, and selection). The added overhead is minimal and has negligible effect on overall training time. We will include these results in the final paper.

W7: Overly Concise Conclusion

Thank you for the helpful feedback. Due to page limits, we kept the Conclusion in the main paper brief to focus on summarizing key contributions, methods, and results. A more detailed discussion is included in Section C of the Supplementary Material, where we further elaborate on the limitation and potential future directions.