Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes a customisation strategy named "Split-Max" for federated learning. The authors identify the heterogeneity of devices and data in FL scenarios. They present the importance of considering devices' budgets and dynamics when dispatching training models. Split-max can adjust the model size according to the devices' budget while maintaining good accuracy and robustness.

Split-max works in steps. First, multiple base models from different initialisations are trained to improve diversity. These base models are randomly given to clients to extract generalisable features. Then, base models are aggregated to the server. Secondly, to provide devices with models with different robustness, it trains two similar models together to capture both the standard-training accuracy and adversarial-training accuracy. Then layer-wise mixing is conducted to achieve both standard accuracy and adversarial accuracy. 

Experiments show that Split-Mix achieves better accuracy than naive approaches. Moreover, with customisation, the models are smaller and more robust under budget constraints.
main_review: ## Strengths:
1. The author clarifies the heterogeneity under FL scenario into two aspects, resource and data. Then, the importance of model customisation is intuitive and motivates this work.
2. In terms of the idea "Split-Mix", the authors present empirical evidence of slim models and wide models, including convergence and gradient saliency. 
3. I like the deep dive in the evaluation section, both the client-wise statistics in Figure 7 and the percentage difference of parameters locally trained between SHeteroFL and Split-Mix. It vividly shows why these approaches perform as presented and how models with different width perform on different categories.

## Weaknesses:
1. The difference between customisable models in a non-FL setting is somehow vague. It is true that the training data are non-sharable. But the primary question of why the mentioned approaches cannot be adapted to the FL scenario could be further clarified. Are there any technical challenges? Or are these approaches tightly coupled with shareable training data?
2. The evaluation could include training time benchmarks of different approaches. As said that Split-Mix allows all base models to be trained on all clients, what is the cost of time in terms of such training method?
3. The idea of BN layer-wise mixing is kind of confusing. The reason of sharing all parameters except the batch-normalisation requires further elaboration. 

## Detail Comments:
1. The notations in the three algorithms lack necessary explanation.
2. Regarding the second contribution stated in Sec. 1, it is better to present direct evidence of efficiency in storage, model loading and inference. For example, the time to load the model, the inference time and latency. The existing evaluations like #Params and FLOPs are indirect metrics
2. The two considerations listed in Sec. 3 are somehow similar. It is better to clarify the difference between heterogeneous computational budgets and imbalance computation budgets. 
3. The authors discuss their solutions on customising robustness as it is quite essential in federated learning. I would like to know how this idea can be adapted to other aspects like as mentioned in Sec. 3 neural -architecture search and fairness. 
4. I am a little confused by the illustration of Fig. 5. So there are two BN in the training process, each of which handles noised/clean data. And in the inference process, the results are aggregated based on the coefficient. How are the noised and clean data differentiated during training? Are the label input together with the training data? 
5. Regarding the evaluation setup, it is recommended to have more details on the clients' budget, e.g. the distribution of client budgets. Are there more clients with a large budget? I suppose the distribution will impact the ultimate accuracy performance.
6. Some minor questions. In Fig. 9(a), why is there only one data point of FedAvg+OAT? Is it due to the divergence in training? In Fig. 7, why are there outliers? It is better to have some general description the baseline methods (i.e. OAT in-situ method) in the evaluation section. I am curious about the condition variable. So what is the ultimate setting in the evaluation section? 
7. Some minor typos/grammar mistakes. In the **Layer-wise mixing" of Sec. 4.2, "a straightforward solution is averaging their outputs". In the ** Training and valuation** of Sec. 5.2, the abbreviation of"robust accuracy" should be "RA" instead of "SA". 
summary_of_the_review: 
This paper proposes to tackle the heterogeneous budget in FL scenario. A weak accept is given based on two reasons. 
(1) The Split-Mix not only considers how to customise models on different clients but also includes how to maintain high robustness of client models where they are trained over non iid data. I think the paper is well motivated and the solution proposed is intuitive and practical. Though the technique used to improve robustness (ie.e ST and AT, parameter sharing and layer-wise mixing) is not new, it well suits the problems and the authors have successfully integrated them with Split-Mix. Moreover, the evaluation shows promising results that Split-Mix could train client models with much fewer parameters but have similar accuracy performance as FedAvg under the budget constraints. 

(2) As mentioned in the Main Review, more clarification could be provided to improve the claims in the paper. Why can't conventional model customisation apply to FL scenario? How to handle customisation on other aspect? And the evaluation could include more details regarding training setups and efficiency.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: Split-Mix is a Federated Learning strategy aimed at easing the problems that arise when a heterogeneous pool of devices/clients (i.e. some with more compute/memory capabilities than others, different data distributions) collaboratively train a global model. Split-Mix trains a model that can later on be customized in terms of model size and robustness. As the name suggests, Split-Mix has two stages: During `split`, a large model is split into smaller base models. Base models are constructed by discarding channels but maintaining the the number and type of layers in the network. During `mix`, the server samples a fraction of the base models (depending on a given client's compute capabilities) and fuses them into a single one that is send to a client to train. All base models are trained on all clients (should these meet the compute requirements of a sub-model) in a federated manner. This means that the more capable devices, train all base models. This is envisioned to happen in a parallel fashion in a given device. Once FL training is completed, a customized model can be deployed to a device/client in hardware-aware and robustness-ware fashion. The Authors refer to this as _in-situ customization_

main_review: ### Strong and Weak points

- (strong) while most works aim at adapting the compute/memory of a network to match the capabilities of the target client, Split-Mix also considers model robustness. This is an important consideration. 

- (strong) the Authors do a good job with Figure 2, showing that the widest model might not always be a good fit (i.e. it might overfit if the data is skewed or contains very few training examples)

- (doubt) the Authors stress about the `in-situ customization` of their proposed framework. As they state, customization is possible "once training is done". My doubt is, in what context is customization useful after training? is it do do inference using the learned based models? is customization available if another FL task (of the same type -- e.g. image classification) is going to start after Split-Mix is done?

- (comment) the Authors seem to state that locally training multiple models in parallel is feasible. Is this a realistic assumption? Can a smartphone or a NVIDIA Jetson device really train (which comes with a relatively high memory peak requirement and energy consumption) several models in parallel ? 

- (doubt) if all base models need to be trained on all clients, wouldn't this translate into a slow FL training process? It would be interesting to see the line plots in Figure 6 showing wall clock time. In 5.1.convergence the Authors state "all base models are independently trained".

- (weak) the `split` stage is not properly explained: how are the base models derived? Even if fixing the ratio to, for example, 0.125x there are many ways of making a model 8x smaller by discarding channels. This is one of the main ideas in the paper but is not covered. Some approaches are better than others, for example, accuracy would be better retained if more channels are discarded towards the last layers (which are also the layers with more parameters) instead of doing so uniformly. 

- (weak) the `mix` stage is not properly explained: From the sentence in the introduction "we mix selected base models to construct desired widths and required robustness", it gives the impression that base models might be merged before sending them to a client but the Authors do not describe how this happens. For example, if a device can afford training a 0.25x model, how are 0.125x base models combined into a 0.25x model? I imagine there are different ways to do this depending on whether the base models are "disjoint" (in the full model representation) or overlapped. This is one of the main ideas in the paper but is not covered.

- (weak) "we measure latency in terms of FLOPS". This is wrong. Latency can be estimated using FLOPS (although is known not to be an accurate representation) but Authors never report any latency estimates (there are no training "wall time" or inference "time") in the paper.


- (weak) The number of parameters for ResNet18 x1 is ~11M (not the reported 5.6M). The number of FLOPS for a ResNet18 using as input a 1x3x32x32 tensor (just a single CIFAR-10 image) is 1.1 Billion FLOPs (I believe Authors are reporting MACs, which, in general are = 0.5*FLOPs). The Authors should revise all FLOPs and #params reported in the paper.


- (weak) For non-IID experiments, the Authors split Digits and DomainNet datasets into 10 and 5 clients, respectively. Given such a small number of partitions, does it mean all clients are sampled at all rounds? I do not think this can be considered a realistic FL setting. This conflicts with the argumentation in the paper, aiming at considering more realistic FL settings accounting for the `different dynamics` and heterogeneity.
summary_of_the_review: ### Recommendation

The main weaknesses of this work is that is that the contributions are not properly explained, Split-Mix requires training the base models in all clients (which is unrealistic for large FL settings); for ResNet18 the FLOPs and number of parameters reported are not computed correctly for the x1 baseline (the others might be inaccurate too). The Authors, without argumentation why, directly state that FedAvg baselines can only afford training 0.125x models, why is this? Which device cannot afford training a 0.25x ResNet18 ? Even a RaspberryPi 4 (i.e. CPU-only) can train a x1 ResNet18 using batchsize=32, which takes ~15seconds per batch. As the Authors state themselves, the advantages of Split-Mix for CIFAR10 and DomainNet (the two most challenging task tested in this work) at somewhat limited. What the results evidence is that only models that are somewhat over-parameterized for a given task (as seems to be the case for Digits -- only a 4% acc drop when going from x1 to x0.125) work well with the proposed Split-Mix strategy. In addition, the Authors should evaluate other types of client heterogeneity distributions. 

Summary: The Authors have motivated well the paper but need to improve the presentation of the technical contribution. A good way of achieve things would be to present Algorithm 3 sooner and then explain its main stages. To me it does not sound natural to give such brief explanations about `split` and `mix` only at the end of Page 6, just before the experiments section.


### Supporting my Recommendation

Please see points above.

### Minor points


- Is it fair to add as cite to "Federated Learning" a paper from 2019 (first sentence of the Introduction) ?
- When describing Fig3 (in Page 4) the sentence starts with "To validation, we investigate [...]". Maybe it should be "To validate this"?
- 6th line in Sec5: "The dataset first is [...]", should be "The first dataset is [...]"
- When Authors mention "layer output cache", I have assumed this is equal to layer activations or "activations buffer". Please correct me if this means something else.
- It is not clear which architecture was used for Digits. Adding a 1-line description of what the task is for Digits and DomainNet would be useful for readers not used to work with these datasets.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper presents a new federated learning approach named Split-Mix FL that allows clients to train customized models efficiently while considering heterogeneity in data and computation resources. The key idea is threefold: 1) the global model is first split into several sub-networks called base models that each have different sizes, thus requiring a different amount of computational resources; 2) Each selected client trains a random subset of base models, under its computational resource constraints; 3) updated base models are aggregated at the server-side and distributed again. Furthermore, the proposed approach can train both accurate and robust models in a joint fashion, where all but batch-norm layers are shared for efficiency. Experimental results on multiple datasets (CIFAR10, Digits, DomainNet) demonstrate the effectiveness of the proposed approach compared to FedAvg and HeteroFL.
main_review: ## Strengths
- Addressing data and resource heterogeneity is definitely one of the important topics of FL.
- The main technical idea, which splits the global network into a collection of smaller base models of various sizes (though I do not perfectly understand how it was done actually) and asks clients to train a random subset of them under computational resource constraints, seems like a straightforward and reasonable solution to the problem.
- Also, jointly training both accurate and robust models while keeping all but BN layers shared sounds like an interesting idea, though it could be useful in contexts other than federated learning.
- Experiments are thorough, demonstrating the effectiveness of the proposed approach on multiple datasets with multiple metrics.

## Weaknesses
It's not perfectly clear how the split and aggregation are done in the proposed approach. On the one hand, Figure 1 illustrates that multiple base models are aggregated into literally a single global model, which is split into random base models that may or may not have identical architectures over rounds. On the other hand, Algorithm 3 describes that W^t is given by a set of base models, which are not aggregated into a single model but kept as a set of models throughout the training procedure. Which is the one that the proposed approach is actually doing?

I think that clarifying the above point is important as this would be one of the major differences between the proposed approach and HeteroFL. HeteroFL aggregates multiple models with different sizes into a single one, and SHeteroFL asks clients to train all affordable base models. If the SplitMix does not actually aggregate multiple base models, does this the main reason to bring performance improvements in the experimental results?

Another concern is the problem setting. The proposed work argues that "Without loss of generality, we assume exponentially distributed budgets in uniformly divided client groups:..." I wonder how much this assumption is reasonable and how dividing clients into four groups in this way is optimal for various situations. My simple question is: what happens if we change the number of groups to better deal with the computational resource heterogeneity of clients, and what happens when clients are not distributed in an assumed way regarding the computational resource (e.g., many clients have sufficient computational budgets)? Further clarification on this point would make the experimental results stronger.
summary_of_the_review: Overall, I think that this is a nice contribution with compelling empirical results. Although the manuscript is largely well written, it was not clear to me how model split and aggregation are done in the proposed approach, and HeteroFL, SHeteroFL, and Split-Mix are different and how these differences bring performance improvements. Also, the assumption about the computational resources of clients would be better clarified.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: The paper proposes a new federated learning scheme that is suitable for devices with heterogeneous resources. The proposal, namely Split-Mix, trains multiple models of different sizes and adversarial robustness levels, tailored to the budget of each device. Empirical results demonstrate the efficacy of the method against the main competitor.
main_review: ### Strengths

[Problem / Motivation] 

The problem of heterogeneous resources in federated training is of utter importance and very challenging both from the systems and algorithms perspective. 

[Related Work] 

The authors have a valuable overview of the related work.

[Novelty]

The paper introduces some interesting ideas based on existing work (HeteroFL and adversarial robustness).

### Weaknesses 

[Significance] 

- The evaluation results are heavily based on the assumption of exponentially distributed budgets. 
 Justification is needed for why this distribution is the only one to consider. Otherwise other distributions need to be evaluated to show the generality of the proposal.

- The notion of budget is not clearly defined and potentially problematic. Each client has a budget in terms of physical resources (e.g., latency, memory, CPU time, disk). This budget can translate into the resources necessary for training a network of a specific size (e.g., one x0.5 net). According to line 5 of Algorithm 3, the authors indicate that the same client can train a linear combination of smaller networks that sum to the budget (e.g., x0.5 -> one x0.25, one x0.2, one x0.05). This is not valid as the resource consumption of a network does not scale linearly with its size (at least for some of the important physical resources). This assumption on the relation between network size and budget essentially permits Split-Mix to train all base models and outperform SHeteroFL which can train only a subset.


[Clarity] 

Some arguments of the paper are not explained.

- The authors mention: "A naive solution is to simultaneously train multiple models with wanted architectures and properties, which however may lead to prohibitive training cost and latency of switching models." The procedure of switching and why it incurs prohibitive latency costs is not clear.

- In Figure 2: x1 net SHeteroFL seems to converge faster than slimmer nets till around round 50. What happens afterwards ? The authors conclude that the widest model may not be a good candidate model but don’t provide any intuition behind this. After reading further the reason, behind this seems to be the less amount of data.


[Evaluation]

Some results are questionable / counter-intuitive.

- Table 1: Why does the accuracy of SHeteroFL not monotonically increase with FLOPS and #Params (as it is for Split-Mix) ?

- Figure 9b: Why λ=0 gives less SA than λ={0.1, 0.2} ? 


[Presentation]

Paper presentation needs improvement.

- The setup for Figure 2 is not clear enough for understanding the demonstrated behaviour. The reader has to read Section 5 to understand the setup and then continue to Section 4. This breaks the flow of the paper.

- Gradient saliency should be defined or described in a few lines.

- Please mention “adversarial robustness” instead of “robustness” at least once in the beginning of the paper.

- There is a large number of typos (mainly after Section 4). A sample:
  - 4.1 (beginning): “To motivation” -> “For motivation”
  - “convergence slower” -> “converge slower”
  - “To validation” -> “For validation”
  - “Last, it worth note that” -> “Lastly, it is worth noting that”


summary_of_the_review: I recommend this paper to be rejected mainly given the issues regarding significance and lack of proper explanations for various claims throughout the paper (see Main Review section).
review_title: null

Meta-review: This paper proposes a federated learning (FL) scheme that is suitable for clients/devices with heterogeneous resources. The scheme Split-Mix trains multiple models of different sizes and adversarial-robustness levels, which are tailored to the budgets of the individual device. Empirical results show encouraging results.

It is clear that FL will have to work with clients with diverse resources, a point that is appreciated. Indeed, it is anticipated that widely-dispersed inference will have to deal with a highly-heterogeneous mix of clients. The study is quite thorough. One aspect that is not convincing in the experiments is the budgets being exponentially distributed: having a strong concentration around a mean (with something like a Gaussian tail), or a power-law distribution, would be more suitable.
Decision: Accept (Poster)