Aggregating Capacity in FL through Successive Layer Training for Computationally-Constrained Devices

Summary

The paper focuses on enabling federated learning (FL) among clients that are bounded by the memory capacity. To this end, the work proposes successive layer training (SLT), a technique that sustains the memory usage below a given memory budget throughout the FL process. SLT partitions the model into three parts: *i)* a first part with frozen weights, *ii)* a intermediate part where all weights are trained, and *iii)* a final part (which can stop before the end of the architecture) whose only a fraction of the weights are trained, controlled by a parameter *s* which selects the number of channels to be trained. At training time, the FL process is broken into *N* stages. At each stage, the layers that are included in each of the three parts are changed. Progressively, the final part is moved towards the last layers of the network, until the whole model has been trained. The proposed scheme is applied on CNNs and compared with a number of baselines that employ nested submodels.

Strengths

1) The proposed STL technique is interesting. Its parametrisation with respect to the indices of three-part network partitioning and the number of configuration steps provides a high level of flexibility and can thus be adapted to various settings, *e.g.* different memory constraints and clients with significant heterogeneity in their memory capacity. 2) The paper is well written, the proposed method is clearly described and positioned within the existing FL approaches. 3) The evaluation includes appropriate baselines and covers adequately different setups.

Weaknesses

1) Althought STL seems to outperform the evaluated baselines, the potential reasons behind this performance gain is not discussed. For example, in the Heterogeneous Memory Constraints setup (Section 3.4), despite being the most realistic, the paper does not discuss why STL performs better than HeteroFl and FjORD, *i.e.* the two strongest baselines for this setup. The *Heterogeneity results* paragraph describes the results of Table 3 in a textual form, but does not attempt to explain why these are as they are. 2) The paper does neither discuss the selection of parameter *N*, which determines the number of distinct configurations to be used during the FL process, nor experimentally evaluates its impact. Similarly, there is no investigation on how these configuration steps are distributed across the global rounds. This is briefly touched upon in the last paragraph of Section 3.2, but this is not sufficient. More thorough investigation is needed, for STL to be useful. The key question is how to train using STL, *i.e.* how many configuration steps to use and how much time to spent using each configuration.

Questions

1) Please specify the number of configuration steps *N* and the duration of each step for the presented experiments. 2) See the two points in Weaknesses. 3) As the work focuses only on memory constraints and not computational capabilities, it would be better to modify the title to reflect that, *e.g.* "Aggregating Capacity in FL through Successive Layer Training for Memory-Constrained Devices" instead of "Computationally Constrained".

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The paper adequately covers its limitations and broader impact. All the main points are discussed.

Summary

The paper proposes an FL training methodology to reduce the memory footprint for constrained devices. Specifically, the layers inside the original model are divided into three categories, i.e., frozen, fully trained, and partially trained. The frozen layers only participate in the forward pass but do not require storing the activations or computing gradients for the backward pass. The partially trained layers also only train a fraction of the parameters, thus reducing the memory overhead. The authors further provide heuristics for determining the configuration of layers within these three categories throughout training and the number of training steps allocated to each configuration.

Strengths

- The paper is well-written, easy to follow, and the ideas are clearly presented. - The experiments investigate the method from different aspects like iid versus non-iid data distribution as well as iid versus non-iid compute devices, and show benefits over several baselines.

Weaknesses

- While the proposed method works well on the evaluated benchmarks, it appears to be a straight-forward combination of existing ideas (using s ratio of parameters to partially train layers + progressive freezing). As such, the reviewer finds the originality of the ideas rather limited. - Can the authors provide some intuition or theoretical justification for acceptable bounds on s and the general configuration of F_F, F_T, and F_H? Specifically, what is the bound on these design hyperparameters after which the convergence is impeded? - There seems to be a large gap between the performance of prior work (FedRolex) versus the performance presented in the original paper. Can the authors clarify any difference in the experimental setup that may have led to this? - It would be beneficial to present the corresponding training memory footprint (activation+parameters+grads) for each benchmark to better clarify what constraints each benchmarked s_fd ratio corresponds to, and see e.g., how aggressive is 0.125 ratio.

Questions

Please see weaknesses above.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors do not discuss the potential negative societal impact of the work in the main text.

Summary

This paper introduces an approach known as Successive Layer Training (SLT) for federated learning (FL) on edge devices with limited resources. The core concept of SLT involves training a portion of the neural network (NN) parameters on the devices while keeping the remaining parameters fixed. Subsequently, the training process progressively expands to include larger subsets until all parameters have been trained. The authors assert that SLT offers several advantages, including reduced memory requirements for training, enhanced accuracy and convergence speed in FL, and compatibility with heterogeneous devices that possess varying memory constraints. To validate their claims, the authors conduct evaluations of SLT on multiple datasets and NN architectures. Additionally, they compare SLT against state-of-the-art techniques such as Federated Dropout (FD), FedRolex, HeteroFL, and FjORD.

Strengths

1, This paper is well-motivated. Memory footprint is a key constraint in training models on edge devices. The proposed approach allows for co-adaptation between parameters and reduces the memory footprint. 2, The paper is well-written, offering clarity and organization. 3, This paper conducted a comprehensive experimental analysis on four datasets (CIFAR10, CIFAR100, FEMNIST, and TinyImageNet) and three NN architectures (ResNet20, ResNet32, and DenseNet40). Comparisons are made with several baselines, and the results demonstrate that SLT achieves higher accuracy and faster convergence in both independent and non-independent and identically distributed (iid and non-iid) settings. Additionally, the performance of SLT is investigated in heterogeneous environments, where devices have different memory constraints.

Weaknesses

The paper primarily focuses on addressing memory constraints in federated learning (FL) by introducing the Successive Layer Training (SLT) technique. However, it does not thoroughly consider other types of resource constraints, such as computation or communication, which are also crucial aspects of FL on devices. The paper does not provide insights into how SLT impacts computation time or communication costs in the FL process. It is important to acknowledge that SLT may introduce certain overheads or inefficiencies in these areas, potentially influencing its practical applicability and overall performance. Further investigation and analysis of the computational and communication implications of SLT would provide a more comprehensive understanding of its practicality and effectiveness in resource-constrained scenarios.

Questions

In the evaluation, the used model sizes are relatively large and can be over-parameterized for the dataset. As a result, the method Small Model, which is still sufficient for the task, can achieve better performance than the SOTA methods. If the model size goes smaller (i.e., LeNet with 1x, 2x, or 4x channels), do the proposed method and the method Small Model still outperform the SOTA baselines?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Please see the weakness and questions.

Summary

In this paper, authors study the case where the edge device is not capable of holding the training of the whole network due to limited memory. Authors find that it is necessary to train a whole network rather than train a submodule as previous methods do. Motivated by this observation, authors propose to first scale down the intermediate size of the network and then expanding it as its precedent layers being fully trained. Authors provided experimental results and those results successfully proved that SLT outperforms other algorithms in this scenario.

Strengths

The observation that it is necessary to train a full network, rather than a submodule, sounds reasonable to me. Moreover, the SLT method proposed looks new to me. This idea could be helpful for future studies in memory-constrained FL training.

Weaknesses

The paper is not clearly written in some parts, such as Section 2.2. My major concerns come from two parts: 1) Why SLT can achieve almost the same result w.r.t. the ''small model'' method when the model can accommodate the edge device, as shown in Figure 4? To me, the shallower layers are only trained when the deeper layers get forward partially in the early stage, and they are kept frozen when training the deeper layers. This means the shallower layers are not fully trained with the correct gradient that comes from the whole model, which means SLT is not equivalent to the full training algorithm and it shall get a worse performance than the ''small model'' method. Therefore I doubt why it can get a matched accuracy as shown in Figure4. 2) Actually the footprint of activations can be alleviated by checkpointing or offload, please compare SLT with these methods.

Questions

Please refer to my questions in the weakness part.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

Authors have clarified their limitations.

Summary

This paper proposed a new training method for neural networks in the federated setting, called successive layer training. This method tries to train NN layer by layer by freezing the trained layer and scaling down the NN heads. The authors have a clear presentation for this method and propose detailed experiments to show the performance.

Strengths

The experimental results show good performance when compared with other algorithms for this memory-constrained setting. It is exciting to see that such an (approximate) layer-by-layer training strategy can have an acceptable accuracy. The configuration for this setting is also flexible when adapting to environments with different resource constrained.

Weaknesses

The method itself is somewhat straightforward and the paper makes no explanation about why this strategy can work. Meanwhile, the experiments were only done on some toy neural networks and thus were not very representative. At least testing on Resnet50 on Imagenet is necessary.

Questions

Although this strategy here is proposed under the background of FL, it is a general approach suitable for any distributed environment. So is it effective for applying the method in standard distributed training, while communication is the bottleneck. Meanwhile, is it possible to apply this strategy to attain the accuracy of full training?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Some suggestions are as follows: 1. Give some explanation for this method 2. Do adequate experiments on standard nets and datasets 3. Study this training strategy as a general approach first and then apply it to FL