Efficient Split-Mix Federated Learning for On-Demand and In-Situ Customization
Authors: Junyuan Hong
Haotao Wang
Zhangyang Wang
Jiayu Zhou
Venue: ICLR 2022 Poster (2022)
Keywords: federated learning

Abstract: Federated learning (FL) provides a distributed learning framework for multiple participants to collaborate learning without sharing raw data. In many practical FL scenarios, participants have heterogeneous resources due to disparities in hardware and inference dynamics that require quickly loading models of different sizes and levels of robustness. The heterogeneity and dynamics together impose significant challenges to existing FL approaches and thus greatly limit FL's applicability. In this paper, we propose a novel Split-Mix FL strategy for heterogeneous participants that, once training is done, provides in-situ customization of model sizes and robustness. Specifically, we achieve customization by learning a set of base sub-networks of different sizes and robustness levels, which are later aggregated on-demand according to inference requirements. This split-mix strategy achieves customization with high efficiency in communication, storage, and inference. Extensive experiments demonstrate that our method provides better in-situ customization than the existing heterogeneous-architecture FL methods. Codes and pre-trained models are available: https://github.com/illidanlab/SplitMix.

Published as a conference paper at ICLR 2022
EFFICIENT SPLIT-MIX FEDERATED LEARNING FOR ON- DEMAND AND IN-SITU CUSTOMIZATION
Junyuan Hong', Haotao Wang', Zhangyang Wang' and Jiayu Zhou'
Department of Computer Science and Engineering, Michigan State University *Department of Electrical and Computer Engineering, University of Texas at Austin (hongjul2, jiayuz] @msu.edu, (htwang, atlaswang) @ut exas. edu
ABSTRACT
Federated learning (FL) provides a distributed learning framework for multiple participants to collaborate learning without sharing raw data. In many practical FL scenarios, participants have heterogeneous resources due to disparities in hardware and inference dynamics that require quickly loading models of different sizes and levels of robustness. The heterogeneity and dynamics together impose significant challenges to existing FL approaches and thus greatly limit FL's applicability. In this paper, we propose novel Split-Mix FL strategy for heterogeneous participants that, once training done, provides in-situ customization of model sizes and robust- ness. Specifically, we achieve customization by learning set of base sub-networks of different sizes and robustness levels, which are later aggregated on-demand according to inference requirements. This split-mix strategy achieves customiza- tion with high efficiency in communication, storage, and inference. Extensive experiments demonstrate that method provides better in-situ customization than the existing heterogeneous-arcchteecuue FL methods. Codes and pre-trained models are available: https: /
/////gtthuuu thub. com i illl http:/llddaaati waw.dddaaabbb/aaba /ppppppll itMix.
INTRODUCTION
Federated learning (FL. (Konečný et al., 2015) is a distributed learning paradigm that leverages data from remote participants and aggregates their knowledge without requiring their raw data be transferred to a central server, thereby largely reducing the concerns from data security and privacy. FedAvg (McMahan et al., 2017) among the most popular federated instantiations, which aggregates knowledge by averaging models uploaded from different participants.
When deploying federated learning, one challenge in real-world applications is the run-time (i.e., test-time) dynamicx: The requirements on model properties (e.g., inference efficiency, robustness, etc.) can be constantly changing during the run-time, depending on the status the devices or the outside environment. One common and specific type of dynamics is resource dynamics: For each application, the allocated on-device resources (e.g., run-time memory, CPU bandwidth, etc.) may vary drastically during run-time, depending are prioritized on a participant's 's device (Xu et al., 2021). ) Another type of dynamics the robustness dynamics: The constantly changing outside environment safety (or robustness) level of the model (Wang et al., 2020). For instance, the quality of real-time videos captured by autonomous cars suddenly degrade, e.g., on entering poor-lighted alley or tunnel from a well-lighted avenue, on entering a section of bumpy road which leads to a sudden burst of blurring in the videos, etc. In such cases, more robust model should be quickly switch in and replace the one used on benign conditions, in order to prevent catastrophic accidents caused by wrong recognition under poor visual conditions. Such dynamic run-time requirements demand the flexibility to customize the model. However, as illustrated in Fig. 1a, we show that conventional static-model FL methods, represented FedAvg, cannot provide such customization. A naive solution is to train multiple models with different desired properties and keep them all on device. However, this leads to extra training and storage costs proportional to the number of models. Moreover, since is not practical to keep all models simultaneously in run-time memory introduces inference overhead to swap the models into and out of the run-time memory (Dosovitskiy & Djolonga, 2019).
how the resource allocation of the running programs
make different requirements
the
resource- limited device, it also

---

Published as a conference paper at ICLR 2022
FedAvg Trannredett w/ single model
network node
inputiourpu
HeteroFL
TraniPreditt mutt wuth-witthhhhhhhhh masrggaennnl
Distrbuut
Aggregate
Apgregate
Distribute
(a) Illustration of FedAvg (McMahan et al., 2017) with device-incompatine model and heterogenous- architecture variant (HeteroFL) (Diao et al., 2021) with under-trained wide models.
Locally trainipredict with customized models
Apgregate
x0.5
random mix adaptiayy d distibuta
Accurtee medel
Robust mode
Robusi
Stare 90% poomis
(b) The proposed Split-Mix framework provides in-situ customization of widths and adversarial robustness to address heterogeneity and dynamics, enabling efficient training and inference. In this example, we use subnet with 1/4 channels (or widths) per layer as base model for model-width customization. For simplicity, we denote it x0.25 net, and x1 net can he split into 4 x0.25 base models.
Figure 1: Comparison of traditional and proposed methods.
To effectively and efficiently train models for on-demand an in-situ customization, new challenges will be raised by the ubiquitous heterogeneity of federated learning participants. Fist, the participants can have resource heterogeneity; Different participants have different hardware resources available, such as memory, computing power, and network bandwidth (Ignatov et al., 2018). For example, in learning task for face recognition, clients may use different types of devices (e.g., computers, tablets or smartphones) to participate in learning. To accommodate different hardware, one can turn to more resource-flexible architectures trained by distillation from ensemble (Lin et al., 2020), partial model averaging (Diao et al., 2021), or directly combining predictions (Shi et al., 2021). Specifically, (Diao et al., 2021) is the first heterogeneous- width solution (HeteroFL) allowing in- situ mizel-size switching. Nevertheless, it suffers from under- under-training training in its large models due to local budget constraints shown in Fig. la. The degradation could be worsened as facing data heterogemeity: The training datasets from participants are not independent and identically distributed (non-i.i.d)) (Li et al., 2020b; Fallah et al., 2020; Hong et al., 2021c; Zhu et al., 2021). When one device with a unique data distribution cannot afford training large model, the global large model may not transfer to the unseen distribution (Pan & Yang, 2010). Thus, HeteroFL may provide effective customization such that more parameters brings in higher accuracy and how to train an effectively customizable model still remains unknown.
To address the aforementioned challenges from heterogeneity and dynamics, we study novel Split- Mix approach enable FL on heterogeneous devices and achieve in-situ model customization for resource efficiency and robustness: The size and robustness of the resultant model can be efficiently customized run-time. Specifically, we first split the complete knowledge in a large model into several small base sub-networks (shards) according to model widths and robustness levels. To complete the knowledge, we let the base models be fully trained on all clients. To provide customized models, we mix selected base models construct the desired model size and robustness. We illustrate the main idea in Fig. 1b. Overall, our contributions can be summarized in three folds:
Within the domain of heterogeneous federated learning, we are the first study training model with the capability of in-situ customization with heterogeneous local computation budgets, which cannot be resolved by existing methods yet as shown in Fig. 1a.
To address the challenge, propose novel Split-Mix framework that aggregates knowledge from heterogeneous clients into a width- and coousstnes-adjustalll model structure. Remarkably, due to fewer parameters and modular nature, our framework is not only efficient in federated communication and ficxibly adaptable to various client budgets during training, but also efficient and flexible in storage, model loading and execution during inference.
2

---

Published as a conference paper at ICLR 2022
Empirically, demonstrate that the performance of the proposed method is better than other FL baselines under heterogeneous budget constraints. Moreover, we show its effectiveness when facing the challenge of data heterogeneity.
2 RELATED WORK
Heterogeneous Federated Learning. As increasing concerns have been gained on data privacy leakage in machine learning (Dwork, 2008; Weiss & Archick, 2016; Wu et al., 2020; Hong et al., 2021b), federated learning (FI.) protects data privacy by training the model locally on users' own devices without sharing data. In real-world applications, FL with budget-insufficiet devices (e.g., mobile devices) has attracted a great amount of attention. For example, FedDistill (Lin et al., 2020) used logit averages to distill network of the same size or several prototypes, which will be communicated with users. FedDistill made an assumption that the central server has access to a public dataset of the same learning task, which is impractical for many applications. He al. (2020) introduced a distillation-baseee method after aggregating private representations from all participants. The method closely resembles centralized learning because encoded samples are gathered, and however it is less efficient when local clients have large data dimensions or sample sizes. Importantly, the method may not transfer adversarial robustness knowledge through the intermediate representations due to the decoupling of the input and prediction. On the other hand, HeteroFL. (Diao et al., 2021) avoids distillation, allowing participants to train different prototypes and sharing parameters among prototypes to reduce redundant parameters. However, HeteroFL. also reduces the samples available for training each prototype, which leads to degraded performance. Considering the high cost of training robust models, Hong et al. (2021a) proposed efficient way to transfer model robustness from budget-sufficient devices to insufficient ones. FedEnsemble (Shi al., 2021) technically related to the proposed approach, which uses ensemble of diverse models to accommodate non-i.i.d. heterogeneity. The authors showed that combining multiple base models trained simultaneously FL. can outperform single base model in testing. A critical difference between the proposed approach and FedEnsemble comes from the challenging problem setting of constrained heterogencous computation budgets. For the first time, we show that base models can be trained adaptively under budget constraints and used to customize efficient inference networks that can outperform a single model of the same width but trained in a wetprogennouu--rcciteeeee way.
Customizable Models. To our best knowledge, this is the first paper discussing in-situ customization in federated learning and here we review similar concepts in central learning. First, customization of robustness and accuracy was discussed by Wang al. (2020), where together with decoupled bateh-normalization were used. The conditional layer enables continuous trade-off but also brings in more parameters. In comparison, simple weighted combination without additional parameters is used in our method and is very efficient in both communication and inference. In terms of model complexity, a series of work dynamic neural networks were proposed to provide instant, adaptive and efficient trade-off between accuracy and latency (mainly related to model complexity) of neural networks the inference stage. Typically, sub-path networks of different complexity (Liu & Deng, 2018) or sub-depth networks (Huang et al., 2018; Wu et al., 2019; Wang et al., 2018; 8: Zhang et al., 2021) are trained together. However, due to the large memory footprint brought by a constant number of channels per layer, the memory footprint at inference is barely reduced. To address the challenge, slimmable neural network (SNN) (Yu et al., 2018) was proposed to train networks with adaptive layer widths. Distinct from SNN, consider a more challenging scenario with distributed and non-sharable data and heterogeneous client capabilities.
cituustabl-cconitionna layer
PROBLEM SETTING
The goal of this work is to develop a heterogeneous federated learning (FL) strategy, that yields a global model, which can be efficiently customized for dynamic needs. Formally, FL minimizes the objective |.D Σx Σ D. [(f(z; W), y). where L is loss function (c.g., cross-entropy loss), f is a model parameterized by W, and {Dx) are the training datasets on K participants with the image-label pairs (x. y). Following the standard FL setting (McMahan al., 2017), only
3

---

Published as a conference paper at ICLR 2022
model parameters can be shared to protect privacy. We require efficient run time customization model f for resource dynamics (Section 4.2) and robustness dynamics (Section 4.3).
models, significant challenges arise from heterogeneity among clients. In this paper, we consider the following: 1) Heterogeneous computational budgets during training of clients constrain the maximal complexity of local models. The complexity of deep neural networks can be defined in multiple dimensions, like depths or widths (Zagoruyko & Komodakis, 2017: Tan & Le, 2019). In this paper, we consider the width of hidden channels, because it can significantly reduce not only the number of model parameters but also the layer output cache. Thus, we assume clients have confined width capabilities {Rg E (0, 1] K defined by width ratios w.r.t. the original model, i.e., XR net as presented in Fig. 1b. Similar to FedAvg and HeteroFL, the same architecture is used by clients and therefore the model width can be tailored according to local budgets. In many applications, there are usually significant number of devices with insufficient computational budgets (Enge, 2021). For example, we may assume exponentially distributed budgets uniformly [4/K), i.e., the first group with 1/4 clients is capable for 1-width and the rest are for 0.5, 0.25, 0.125-width respectively. The budget distribution simulate real scenario where most federated mobile phones prefer low-energy solution with smaller models for training probably in the background and maintain more resources for major tasks. Other budget distributions, c.g., log-normal distributions, are discussed in Appendix A.5. 2) Heterogeneous data distributions, c.g., non-i.idd features, Dù D, for ordered domain i, induces additional challenges with a skewed budget distribution, since training a single model on multiple domains (Li et al., 2020b) models in single domain (due budget constraints) (Hong et al., 2021c; Dong al., 2021) are known to be suboptimal on transfer performance.
When training customizable/adjustaaae
divided client groups: R&
In summary, training various model sizes has challenges from 1) resource heterogeneity, where disparity in computation budgets will significantly affect the training of large models because they can only be trained on scarce budget-suffccient clients; data heterogeneity, where the under-trained large models may perform poorly on unseen in-domain or out-of- samples.
4 METHOD
To provide efficient in-situ customization, we introduce a simple yet powerful principle, Split-Mix: shatter complete knowledge into smaller pieces, and customize by flexible formations of pieces. Based on the principle, we propose customizations of model size and robustness in the following.
4.1
CASE STUDY: CUSTOMIZABLE NETWORKS FROM BUDGET- BUDGET-CONSTRANN CONSTRAINED FL
For motivation, set up standard non-i.i.d. FL benchmark (Li et al., 2020b) using DomainNet dataset (refer to experimental details in Section 5) and study how the budget constraint impedes effective training of customizable models.
0.70
Case 1: FL without budget constraint. First, we individually train networks of different widths by FedAvg. We see that the slimmer networks converge slower and are less generalizable, even though they are trained all clients (solid lines in Fig. 2). The results justified the motivation of training wider networks than slimmer ones.
0.65
validitoon scureer
0.60
8.55
width
0.125
0.25
Case 2: FL with budget constraint. Following the above budget constraint, i.e., R2 = (1/2) 44/KK1 we deploy HeteroFL to train budget-compatible prototype models locally. Since HeteroFL was not designed for model customization, data are not fully used: each model prototype is only trained by 1/4 of clients, when clients of the R, = 1 budget can actually afford all slimmer models. Therefore, we extend HeteroFl. to a slimmable version (SHeteroFL.) by training all affordable prototypes locally. As shown in Fig. 2, wider models (c.g., x 1 net) converge to validation accuracy lower than not only the FedAvg counterparts,
0.45
mode
0.40
Fedlug
SHieter
100
200
300
communicaaion round
Figure 2: Convergence of different- width models on DomainNet.
CCustomization of other properties can also be applied within our framework. We consider model size and robustness given their practical importance.
4

---

Published as a conference paper at ICLR 2022
but also the slimmer x0.25 net, showing that the widest model may not be good candidate model. Therefore, switching models to wider configurations lowers efficiency but does not improve accuracy, which is nor valid customization.
From the perspective of data allocation, is not surprising that SHeteroFL exhibits a non-monotonone relation between the model width and accuracy. For x1 nets, only 1 1/4 clients and data are accessible for training. In comparison, 3/4 data are accessible by x0.25 nets and the more data empowers x0.25 nets better generalization ability than x1 nets.
4.2 CUSTOMIZE MODEL SIZE
Motivated by the above observations, we propose to increase accessible training data by splitting wide networks into
sub-networks and re-mix afterward. The overall FL algorithm is summarized in Algorithm
More accessible data by splitting wide networks. Since a wide network cannot fit into budget- insufficient clients, we split it into budget-compatiblc sub-networks by channels (width) while maintaining the the total width. In terms of memory limitations, each sub-network can be painlessly and individually trained in all clients. However, sequentially training multiple slim base networks could be much slower than training a single integrated one and increases blocking time on commu- nication. Noticing that all base models can evaluated independently, we instead efficiently train xr base models in parallel (see Algorithm 3). Despite the benefit, client's budget constraint (/t) will limit the number of paralleled base models within [R/r]. as wider channels result larger intermediate layer cache (activations), and excludes the rest base models from the client. Fortunately, we can select different sets of base models for a budget-limine limited client per round, which is inspired by FedEnsemble (Shi et al., 2021) and is presented in Algorithm 2. Hence, all base models can be ever trained on the client for multiple communication rounds, though not continuously every round. Note that since the federated training processes of base models are independent without interference, the training could be as stable as FedAvg with partial participants. In addition, the combination of slimmer base models flexible and can conform a variety of client budgets.
Boost accuracy by mixture of subnet experts. To craft wide model, we combine the outputs of multiple xr base models until the size of the ensemble reaches the same as the number of channels, e.g.. R//r] bases for an x R net. We randomly initialize base models independently such that the diverse bases could extract different features from the same image (Allen-Zhu et al., 2020). Therefore, the ensemble can predict based on a variety of features, resembling an integrated wide network. We follow the common practice, Kaiming's method (He et al., 2015), for initializing base networks with ReLU layers. As Kaiming's 's method is width-dependent, we parameterize the initialization based on the width of x net instcad of the xr one, which leads to smaller initial convolutional weights. In Appendix A.3, intensive ablation studies show that the rescaled initialization is critical for improving test accuracy of wider networks.
Algorithm Federated Split-Mix Learning
Input: Client datasets {De) the number of total communication rounds T, M = [1/r] randomly initialized base xr nets parameterized by (wi, YM client budgets RR ) the number of local epochs E, learning rates {m
1: Initialize [w] M1 model indexes P = Shufffe(l,.... M)) and current index p =
2: for round t E {1, , do
3: Initialize W* +- [wi,- +- O3 and aggregation weights e; = 0 for E {l....,
4: for & E {1...... .K] do
5: Sample hase models W .p= + SampleRaseModels((P,, (P.p. [Rx/r], W'
6: Send W to client and train WE + LrealTrain(W?-- Du, E,nc)
7: Aggregate wi E W to the server: wi,r + wi,r + I
Dal. c: + c: + |Dal
8: Server update V + {wf,- + w. M (Optional) Sort nT = Jim1 by the desconding order of the validation accuracy of wi,r
9:
10: Output the customizable model W = (Mf.
J.
Customize an R-width model by F(x; WT, Ræ IKg f(x; wi, where Kn = [R/r]
5

---

Published as a conference paper at ICLR 2022
Algorithm 2 SumpleBuseModdelsSPPPP
W)
Algorithm 3 LocalTrain( u. DR. E.n)
1: p > |P| then Shuffle P and p t- 1
Initialize models We by Wi. ........ E, do 3: for batch data (Ir, y) in D3 do
2: Initialize W = [WPpp)))
2: for e E
3: if a>11 then
4: Uniformly sample n values into from
4: for WH,r E We in parallel do
PM[[pp]} without replacement
Wi,r + WN. and (a)
5: W + W U (wer, Vi E S}
5:
6:
Return W 11p-1
6: Return We
4.3 EXTENSION TO ADVERSARIAL ROBUSTNESS CUSTOMIZATION
In this section, we extend the customization from one dimension to two dimensions, by jointly customizing model size and model robustness under adversarial attacks (Goodfellow et al., 2014). Model robustness has gained increasing interest (Hendrycks & Dietterich, 2018; Wang et al., 2021), especially in high-stakes federated learning applications (Hong et al., 2021a). Adversarial training (AT) (Madry et al., 2018) is arguably the most popular and effective defense strategy against adversar- ial attacks. Specifically, it uses on-the-fly fly adversarial samples augmentation to improve robustness. Formally, AT minimizes the following augmented loss:
T( (1 (cE)))(((())))) v) An max (Cca(((e 6), y),
(1)
where 5 is subtle e-constrained adversarial noise and transfers clean sample into adversarial sample x + 8. In Eq. (1), LCE is the cross-entropy loss the hyper- paper-parameter and Xn trades off accuracy (the 1st term) and robustness (the 2nd term). When A or 1, the optimization yields standard-training (ST) model or an AT model, respectively. Since an AT model is commonly less accurate in predicting standard images (Tsipras et al., 2019), there is usually no such a sweet point of Xn simultaneously maximizing robustness and accuracy, and one typically needs to carefully gauge the trade-off according to the demand of robustness in specific application context.
Splitting and sharing parameters. Since stan- dard performance and adversarial robustness are irreconcilable, we can directly use two sepa- rated ST and AT models to maximally capture the each property. But do we really need two totally separated models? Intuitively, the two models share some common knowledge, given that an adversarial image share a large part of common features with its original version. As introduced by (Xie & Yuille, 2019), sharing all parameters except the match-normalization (BN) layers can maximize robustness and accuracy by expertised BNs, respectively. Accordingly, we propose to split BN layers (instead of the whole model) into two sub-components: one for standard performance and the other for robustness. At training time, the first loss of Eq. (1) is computed with clean BN (BN,) merely and the second adversarial loss is computed with noised BN (BN,). The FL local training is elaborated in Algorithm 4. As the two BNs are decoupled, there is no more trade-off in loss Eq. (1) and thereby we choose 'n = 0.5 to balance their effects parameter updates.
Training
Loss
Nalaod data Bow Clean data flw Unsoon data Now
Inference
TDBN
TDEN BN.
BN_
BN,
BN,
plean nolsed
Figure 3: Illustration of dual batch-normalization (DBN) in training and inference. The BN, and BN, are for clean and noised samples.
Customizable waser-wise mixing. After training, the problem is how to mix the two models (with different BNs) for prediction. A straightforward solution is averaging their outputs. However, forwarding memory footprint will be doubled in this way, as the two models have to be both executed separately. To avoid the doubled intermediate outputs. gate function can be used adaptively choose BNs like (Liu et al., 2020). Inspired by the method, we further propose simple parameter-free method by weighted-averaging outputs of each BN layer (see Fig. 3):
DBN(2) (1 (z)
(.),
given a BN-layer input x. We note that the averaging strategy is entirely training-fred free and does not use extra parameters, and the customization weight À is intuitive for trade-offs between SA and RAA
Lastly, it is remarkable that the DBN structure is rather lightweight in terms of model complexity. As investigated in (Yu al., 2018) (Table 2), the parameters of BN is more than 1 1% in popular deep

---

Published as a conference paper at ICLR 2022
architectures, e.g., ResNet MobileNet. Therefore, can plug DBN into base models in place of BN, replace Algorithm with Algorithm 4, and jointly customize robustness and model widths.
5 EMPIRICAL STUDIES
We design experiments to compare the proposed method against FL classification benchmarks. For class non-i.id configuration, we use CIFARIO dataset (Krizhevsky, 2009) with preactivated ResNet (PreResNet18 8) (He et al., 2016). CIFARIO contains over 50, 000 32 x 32 images of 10 classes. The CIFARIO data are uniformly split into 100 clients and distribute classes per client. For (feature) non-i.id. configuration, with AlexNet extended with BN layers after each convolutional or linear layer (Li et al., 2020b).
use Digits with CNN defined and DomainNer datasets (Li et al., 2020b)
The first dataset is subset (30%) of Digits, benchmark for domain adaption (Peng et al., 2019b).
Digits has 28 x 28 images and serves as a commonly used benchmark for FL (Caldas et al., 2019;
McMahan et al., 2017: Li et al., 2020a). The dataset includes 5 different domains: MNIST (Lecun et al., 1998), SVHN (Netzer et al., 201 20,1), USPS (Hull, 1994), SynthDigits (Ganin & Lempitsky, 2015), and MNIST-M (Ganin & Lempitsky, 2015). The second dataset is DomainNet (Peng et al., 2019a) processed by (Li et al., 2020b), which contains 6 distinct domains of large-size 256 x 256 real-world images: Clipart, Infograph. Painting, Painting, Quickdraw, Real, Sketch. Each domain of Digits (or DomainNet) are split into 10 (or 5) clients, and therefore 50 (or 30) clients in total. We defer other details such as hyper- paper-parameter to Appendix A, and focus on discussing the results.
5.1 CUSTOMIZE MODEL SIZES
In this section, we evaluate the proposed Split-Mix on tasks of customizing model sizes through adjusting model widths. Recall that in Section 3, we assume a specific heterogeneous training budget to facilitate our discussion, such that one client can only train models within a maximal width, and the resource distribution is imbalance among clients: R& (1/2)[4//11, <k/K1, Following the common FL setting, we do not consider algorithms that use public data or representation sharing in our baselines.
Baselines. As an ideal upper bound but a memory incooppiicc baseline, we (re-)train networks from scratch by FedAvg to obtain individual models with different widths. The srate-of-the-wae veterogeneouu--rccitteee
FL method is HeteroFL (Diao et al., 2021), which trains different slim models in different clients. For fair comparison, we extend HeteroFI. with hounded-slimmabyy training in clients who can afford the larger models, named SHeteroFl.. For example, if a client in HeteroFL can afford x0.5 net, then the client meanwhile trains x 0.225 net and other smaller subnets in slimmable training manner (Yu et al., 2018) by SHeteroFL.
AIIAR10 cless non-i.id
Digits
Domainnet
0.5
0.8
0.5
0.4
0.7
mode
mode
0.4
made
0.3
ShelarofL
fud/vg
0.6
SHateroFf
Pos/vg
0.3
Packig
Sitciaro
Sp 30
Spii-Miic
0.2
0.5
02
communicatinn
Figure 4: Validation accuracy of budget-compatible full-wdith nets by iterations.
Convergence. In Fig. we compare the convergence curves of full-width models on three datasets. For FedAvg, the budget-compatible width is x0.125, Both for SHeteroFL and Split-Mix, the full- width is x1. Compared to baselines, the proposed Split-Mix converges faster, largely due to the splitting strategy. Specifically, because all base models are independently trained, the convergence of each base model mainly depends on how frequently they are trained. When enough clients participate FL. the training frequency of the x net is more than one, all bases will be selected at least once in a communication round."
Performance. In Table 1, we compare the test accuracy of different model widths. We measure the latency in terms of MACs (number of multiplication-and and-adiitionon operations) and model size in method outperforms SHeteroFL, using much fewer
parameter numbers. With the same width, parameters and thus conducts inference in much lower latency. Remarkably, compared to the best model by SHeteroFL., e.g., 81.8% x1 net in CIFARIO, 10, Split-Mix uses only 1.8% parameters and
7

---

Published as a conference paper at ICLR 2022
Table 1: Test results of customizing model width. MACs and the number of parameters are counted inference time. Grey texts indicate that the training cannot conform the predefined budget constraint. The 'M' after metric values means x x10
Individual FedAvg
SHeteroFL
Split-Mix (ours)
width
Acc MACs #Params
Acc
MACs
#Params
Acc MACs #Params
CIFARIO class non-i.i.d FL
x0.125
43,4%
0.9M
0.2M
49.1% %
0.9M
0.2M
48.0%
0.9M
0.2M
x0.25
45.7%
3.5M
0:7M
51.4%
3.5M
0.7M
51.1%
1.8
0.4M
x0.5
50.3%
14-0M
2.8M4
51.5%
14.0M
2.8M
52.1%
3.6M
0.7M
53.3%
55.7M
11.2M
49.9%
55.7M
11.2M
52.7%
7.2M
1.4M
Digits
feature non-i.i.d FL
x0,125
86.1%
0.1M
0.2M
86.8%
0.1M
0.2M
84.6%
0.IM
0.2M
x0.25
87.3%
0:4M
0.9M
87.9%
0.4M
0.9M
87.5%
0.2M
0.4M
x0.5
88.7%
1:M
3.6M
86.9%
1.3M
3.6M
89.0%
0.5M
0.9M
x1
89.6%
4.8M
44.2M
81.3%
4.8M
14.2M
89.8%
0.9M
1.8M
DomainNet feature non-i.id FL
x0.125
67.2%
2.5M
0.9M
66.9%
2.5M
0.9M
68.4%
2.5M
0.9M
x0.25
69.8%
7.5M
3.6M
67.8%
7.5M
3.6M
71.9%
5.0M
1.8M
x0.5
72.6%
25.5M
44.3M
66.9%
25.5M
14.3M
73.0%
9.9M
3.6M
x1
72.9%
92.5M
57.1M
58.7%
92.5M
57.1M
74.2%
19.8M
7.2M
1.6% MACs (x0.125 net) to reach similar level test accuracy. We notice that SHeteroFL has much lower test accuracy with x0.125 net the CIFAR10 dataset. By investigating the loss curves, we find that the inference between parameter-share convergence of the x0.125 5 net. Attributed to the independent splitting, method is stable on convergence with different widths. Remarkably, Split-Mix only requires 12.7% parameters and 19.8% MACs (by x1 ensemble) achieve the comparable accuracy the x individual network on Digits. Results on CIFARIO and DomainNet show a potential limitation of our method. The accuracy of Split-Mix not comparable with the wider individual models due to the locally limited complexity. However, the limitation is more from the problem setting itself, and will not undermine our advantage budget-limited FL. since wide individual models cannot be trained in this case.
different prototypes results unstable
1.0
10tl
*paramutershhunn 100
MACs/batch
207
0.6
mode
reddvg
SHetorsFL
0.2
SplitMic
Ofwrto_mm miidd
Bigits
Domanket
@nhrt0_rnid
Digits
Domaintiet
Oftr30_
Digits
Domaintter
Figure 5: Client-wise statistics of test accuracy, training and communication efficiency by budget constraints. The MACs quantify the complexity of one batch optimization in client, and the number of parameters per round round are the ones uploaded (or downloaded from) server. Test accuracy is by the full-width networks. The results of FedAvg are from codget-compatible x0.125 nets.
Client-wise evaluation. In addition to comparisons of inference on average, we also demonstrate the statistics of test accuracy and the training and communication efficiency in Fig. 5. Conforming the budget constraints, our method outputs more accurate full-width models, transfer fewer parameters round and execute fewer multiplication-ang and-addddd operations for gradient descent than SHeteroFL, either terms of average or variance. Because only the individual x0.125 net can fit into the budget constraint, FedAvg requires the least MACs and parameters, which however significantly sacrifices the final accuracy.
Domain-wise evaluation. To understand why Split-Mix outperforms SHeteroFL, we investigate the total percentage of parameters that can be trained in each domain, in Fig. (Left). count the total parameters that were ever trained in clients of a domain during the learning. Thanks to the base-model sampling strategy (i.e., Algorithm 2). Split-Mix allows all base models, rather than subset, to be trained on all clients. On the other hand, varying client budgets greatly impacted SHeteroFL by limiting the width of models trained in budget-insufficie clients, e.g., in the clipart, infograph and painting domains. Hence, SHeteroFL. leaves large amount of parameters under-trained and suffers
8

---

Published as a conference paper at ICLR 2022
percentage trained parameters
SHeteroFL
SplitMix
100
mode
80
Feddug
0.25 0,125
1.0
-27
-14
-19
08
-1.3
SHeteroFL
60
SpiiMM
width
0.25
10.8
4.6
0.3
-1.5
0.5
-299
witth
48
13
40
15.1
6.7
-11
0.8
-12
15
0.1
0.4
20
1.0
22.3
19.6
26.9
19.4
4.8
1.0
4,0
-.8
0.7
233
cipart
Inlegraph
Painting
Bdrcdraww
real
sketch
dipart
intograph
Painting
quicddraw
reat
dearch
ciport
infograph
painting
qucchrrw
real
stetch
Figure 6: Per domain in DomainNet, the total percentage of parameters that are locally trained (the left figure) and the accuracy (%) drops compared to FedAvg individual models (right two figures). from larger accuracy losses in the three domains and wider models. In comparison, Split-Mix not only has less accuracy drop but also more stable in all domains.
5.2 CUSTOMIZE ROBUSTNESS
Training and evaluation. For local AT, we use an n-step projected gradient descent (PGD) attack (Madry et al., 2018) with a constant noise magnitude Following (Madry et al., 2018), we set ne, = (8/255, 7), and attack inner-loop step size 2/255, for training, validation, and test. For simplicity, we temporarily relax the budger constraint Rg and let the base model be x1 net. For comparison, we extend FedAvg with AT which yields individual models optimized with different trade-off variables, i.e., A, f0, 0.1, 0.2, 0.3. 0.5, 1). Also, we extend FedAvg with state-of-the-ee in-situ trade-off method, OAT (Wang et al., 2020), as baseline. Split-Mix+DAT is an extension of Split-Mix by the proposed DBN-based AT in Algorithm 4. We evaluate and contrast models in two metries: standard accuracy (SA) on the clean test samples and robust accuracy (RA) on adversarial images generated from the clean test set by the PGD attack. Both metric values are averaged by users. We evaluate the trade-off off effectiveness by comparing the RA-SA curves, which is better approaching the upper corner. To plot the curves for Oed and SpTtt-Mix+DAT, we vary their condition variable A in 0.2.. 0.5. 0.8, 1).
0.4
0.3
0.50
0.3
0.55
0.2
0.1
0.50
0.55
Foddugg
FodAwg+AT
Spil-Mis-DAT
mode
5.1
0.25
Tode
0.5
1.
Fostigg
0.0
0.75
0.50
0.25
0.90
0.52
0.54
0.26
0.86
0.70
0.75
0.90
0.50
0.45
0.50
0.55
0.60
SA
SA
SA
SA
(a) CIFARIO
(b) Digits
(c) CIFARIO
(d) DomainNet
Figure 7: Trade-off between robust accuracy (RA) and standard accuracy (SA) with full width (a.b)
and customizable widths (c,d).
Trade-off curves are presented in Fig. (a) and (b). Since the naive extension of OAT with FedAvg adopts heterogeneous objectives and parameterrizationn of conditional layers which suffers from convergence issues, its adversarial training does not converge and RA is incredibly poor in some cases. As result, the trade-off curve is not smooth. Instead, the proposed Split-MixDDAT method has smoother trade-off curve without heavy conditional training or over parameterization. By training in one pass, Split-MixDDAT even outperforms and is more efficient than the FedAvg+AT baselines.
Joint customization of width and robustness under budget constraints. Now consider the width customization and the training budgets as Section 5.1. Due to the constraint, FedAvg can only short name
train x0.125 net. We omit OAT for its unstable convergence and use SplitMixDAT of Split- DATT In Fig. 7 (c) and (d), the trade-off curves with different widths are depicted. As the width increases, both RA and SA of SpTitMixDAT are improved, when they are smoothly traded off. The results demonstrate the flexibility and the modular nature of our method.
6 CONCLUSION
In this paper, we proposed novel federated learning approach for in-situ and on-demand customiza- tion to address challenges arise from resource heterogeneity and inference dynamics. We proposed a Split-Mix strategy that efficiently transfer clients' knowledge to collaboratively learn customizable model. Extensive experiments demonstrate the effectiveness of the principle in adjusting model widths and robustness when much fewer parameters are used compared baselines.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
This material is based part upon work supported by the National Institute of Aging IRF Office of Naval Research 100014-20-1-2388- National Science Foundation under Grant IIS- 114-1749940. Z.W. supported by the U.S. Army Research Laboratory Cooperative Research Agreement W911NF17-2-0196 (IOBT REIGN).
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Zeyuan Allen-Zhu. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning, arXiv:2012.09816 [cs.LG), December 2020.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H. Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv:1812.01097 [es, stat], December 2019.
Enmao Diao, Jie Ding, and Vahid Tarokh. Heterofl: Computation and communication efficient feder- ated learning for heterogeneous clients. In International Conference on Learning Reprexentations, 2021.
Jiahua Dong, Yang Cong, Gan Sun, Zhen Fang, and Zhengming Ding. Where and how to transfer: Knowledge aggregation-induced transferability perception for unsupervised domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence, pp. 1-1, 2021.
Alexey Dosovitskiy and Josip Djolonga. You only train once: Loss-conditional training of deep networks. In International Conference Learning Representations, September 2019.
Cynthia Dwork. Differential privacy: A survey of results. In International conference on theory and applications of models of computation, pp. 1-19. Springer, 2008
Eric Enge. Mobile vs. desktop usage in 2020. http:://www.eeffcccoc..oo////sss
hub/mmbbile-s-ddsstto-uuugg March 2021.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: meta- learning approach. In Advances in Neural Information Processing Systems, June 2020.
Yaroslav Ganin and Victor Lempitsky. Unsupervised domain adaptation by backpropagation. In International Conference Machine Learning, pp- 1180-1189. PMIR, June 2015.
Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv: 141iv:1412.6572. 2014.
Chaoyang He, Murali Annavaram, and Salman Avestimehr. Group knowledge transfer: Federated learning of large cnns at the edge. Advances in Neural Information Processing Systems, November 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1026-1034, 2015.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. Computer Vision and Pattern Recognition, 2016.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness common corruptions and perturbations. In International Conference on Learning Representations, 2018.
Junyuan Hong, Haotao Wang, Zhangyang Wang, and Jiayu Zhou. Federated robustness propagation: Sharing adversarial robustness in federated learning. arXiv preprint arXiv:2106. 10196, 2021a.
Junyuan Hong, Haotao Wang, Zhangyang Wang, and Jiayu Zhou. Learning madel-based privacy pro- tection under budget constraints. Proceedings of the AAAI Conference on Artificial Intelligence, 2021b.
10

---

Published as a conference paper at ICLR 2022
Junyuan Hong, Zhuangdi Zhu, Shuyang Yu, Zhangyang Wang, Hiroko H. Dodge, and Jiayu Zhou. Federated adversarial debiasing for fair and transferable representations. In Proceedings of the 27th ACM SIGKDD Conference Knowledge Discovery & Data Mining, KDD '21, pp. 617-627, New York, NY, USA, August 2021c. Association for Computing Machinery.
Gao Huang, Danlu Chen, Tianhong Li, and Felix Wu. Multi-scale dense networks for resource efficient image classification. Internaiional Conference on Learning Representationss pp- 14, 2018.
Jonathan Hull. A database for handwritten text recognition research. IEEE Transactions on Pattern Analysis and Machine Intelligence, 16(5):550- 554,, May 1994.
Andrey Ignatov, Radu Timofte, William Chou, Ke Wang, Max Wu, Tim Hartley, and Luc Van Gool. Ai benchmark: Running deep neural networks on android smartphones. In Proceedings of the European Conference on Computer Vision (ECCV) Workshops, pp. 0-0, 2018.
Jakub Konečný, Brendan McMahan, and Daniel Ramage. Federated optimizationn:dstiibutc.. opti- mization beyond the datacenter. arXiv:1511.0333 [cs. math], November 2015.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Y. Lecun, L.. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-2324, November 1998.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Conference Systems and Machine Learning Foundation (MLSys), April 2020a.
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. Fedbn: Federated learning on non-iid features via local batch normalization. International Conference on Learning Representations, September 2020b.
Tao Lin, Lingjing Kong, Sebastian U. Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. In Advances in Neural Information Processing Systems, June 2020.
Aishan Liu, Shiyu Tang, Xianglong Liu, Xinyun Chen, Lei Huang, Zhuozhuo Tu, Dawn Song, and Dacheng Tao. Towards defending multiple adversarial perturbations via gated batch normalization. arXiv:2012.01664 [cs], December 2020.
Lanlan Liu and Jia Deng. Dynamic deep neural networks: Optimizing accuracy-efficiency trade-offs by selective execution. arXiv:1701.00 000999999900911111 [cs, stat], March 2018.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. International Conference on Learning Representations, 2018
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication--fffcceett learning of deep networks from decentralized data. In Artificial Intelli- gence and Statistics, pp. 1273-1282, April 2017.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning 20/ 2011.
Sinno Jialin Pan and Qiang Yang. A survey on transfer learning. IEEE Transactions on Knowledge and Data Engineering. 1345513345-339 October 2010.
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1405-1415, 2019a.
Xingchao Peng, Zijun Huang, Yizhe Zhu, and Kate Saenko. Federated adversarial domain adaptation. In International Conference on Learning Representations, September 2019b.
11

---

Published as a conference paper at ICLR 2022
Naichen Shi, Fan Lai, Raed Al Kontar, and Mosharaf Chowdhury. Fed-ensemble: Improving generalization through model ensembling in federated learning. arXiv:2107.1066 [cs, stat], July 2021.
Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In Proceedings of the 36th International Conference on Machine Learning, pp. 6105-61 14. PMLR, May 2019.
Dimitris Tsipras, Shibani Santurkar, Logan Engstrom, Alexander Turner, and Aleksander Madry. Robustness may be at odds with accuracy. International Conference on Learning Representations, September 2019.
Haotao Wang, Tianlong Chen, Shupeng Gui, Ting-Kuei Hu, Ji Liu, and Zhangyang Wang. Once-for- all adversarial training: In-situ tradeoff between robustness and accuracy for free. Advances in Neural Information Processing Systems, November 2020.
Haotao Wang, Chaowei Xiao, Jean Kossaifi, Zhiding Yu, Anima Anandkumar, and Zhangyang Wang. AugMax: Adversarial composition of random augmentations for robust training. Advances in Neural Information Processing Systems, 2021.
Xin Wang, Fisher Yu, Zi-Yi Dou, Trevor Darrell, and Joseph E. Gonzalez. Skipnet: Learning dynamic routing in convolutional networks. ECCV, July 2018.
Martin A Weiss and Kristin Archick. US-EU data privacy: from safe harbor to privacy shield, 2016.
Zhenyu Wu, Haotao Wang, Zhaowen Wang, Hailin Jin, and Zhangyang Wang. Privacy-preserving deep action recognition: An adversarial learning framework and new dataset. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020.
Zuxuan Wu, Tushar Nagarajan, Abhishek Kumar, Steven Rennie, Larry S. Davis, Kristen Grauman, and Rogerio Feris. Blockdrop: Dynamic inference paths in residual networks. arXiv:1711.0853 [cs], January 2019.
Cihang Xie and Alan Yuille. Intriguing properties of adversarial training at scale. International Conference on Learning Representations, December 2019.
Mengwei Xu, Jiawei Liu, Yuanqiang Liu, Felix Xiaozhu Lin, Yunxin Liu, and Xuanzhe Liu. A first look deep learning apps on smartphones. arXiv: /8.Xiv/8/1..0 18/2.05448 [cs], January 2021.
Jiahui Yu, Linjie Yang, Ning Xu, Jianchao Yang, and Thomas Huang. Slimmable neural networks. arXiv:1812.08928 [cs], December 2018.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv:1605.07144 [cs), June 2017.
Linfeng Zhang, Chenglong Bao, and Kaisheng Ma. saff-istillattin: Towards efficient and compact neural networks. IEEE Transactions on Pattern Analysis and Machine Intelligence, pp. 1-1, 2021.
Zhuangdi Zhu, Junyuan Hong, and Jiayu Zhou. Data-frec knowledge distillation for heterogencous federated learning. In International Conference Machine Learning, pp. 12878-12888 PMLR. 2021.
12

---

Published as a conference paper at ICLR 2022
A EXPERIMENTS
In this section, we provide more details about our experiments and additional evaluation results.
A.1 PARALLEL IMPLEMENTATION SPLIT-MIX AND CONVERGENCE
In this section, we elaborate on the implementation and efficiency of Split-Mix. In Fig. 8, we conceptually compare the training by three methods, when two clients capable of training x and x0.5 are considered. FedAvg can train the x1 net on the x -cappabl client but not on the x0.5- capable client. Through parameter sharing among different model widths, SHetcroFL can train multiple widths in sequential manner and can fit into x0.5-capable clients. Unlike SHeteroFL, Split-Mix trains four base models one parallel pass and therefore the most efficient and flexible method. Remarkably, in the worst case, training x1 net of Split-Mix is efficient FedAvg and more efficient than SHeteroFL, if all the base weights, w W3,r E R' (where r = 0.25 here), are embedded into a 4r 4r weight matrix². When base models are embedded into full net, non-trainable parameters can be masked out (grey areas) to avoid interference between base models.
10
LocalTrain
w²
Localiinn
(1)
(4,39
[2]
(a) FedAvg
Locaffrain
(b) SHetcroFL
LocalTrain
AP
(1)
(3)
War
x0.5 Locaffrain
(2) train
(2) train
(6.9)
(4.y)
(c) Split-Mix
Figure 8: Illustration of training weight matrices on x x1-net-capable or x0.5-net-capable client. (1) Download the global weight matrix W of layer I selected subset Wt. (2) Train weights batch data (x, v). (3) Upload trained weight matrix Wi!
With the aforementioned implementation, we compare the convergence versus the wall-clock time in Fig. 9. We implement all algorithms in PyTorch 1. 1.4.
single NVIDIA RTX A5000 GPU and a 104-thread CPU. Fig. 9 shows the elapsed computation time from the initialization to maximal number of iterations. The maximal number of iterations is set be the same for all methods, such that it is easier to compare the stopping time. In Fig. 9, Split-Mix is much more efficient than the FheteroFL. Note that FedAvg can only train the x0. 125 net which includes much fewer parameters, For this reason, Split-Mix is slightly slower than FedAvg. but the degradation of efficiency trades in better accuracy than FedAvg.
For the simplicity of notations, assume the width of layer is though the width could vary by layer in general.
13

---

Published as a conference paper at ICLR 2022
0.8
0.7-
mode
Fedlup
0.6
SHilerrFL
SpllMix
0.5
20
Imiin)
Figure Validation accuracy the bodgett-ommaaiilyy-weess nets by wall-clock time. All algorithms are run with the same number of iterations (200).
A.2 EXPERIMENTAL CONFIGURATIONS
In Algorithm 4, we elaborate the local training of DBN. Specifically, each BN is trained independent with different input samples. The maximization problem in Algorithm can be solved by an n-step projected gradient descent, and is commonly known PGD attack (Madry al., 2018).
Algorithm 4 LocalTrain(W', k, Da, B,y) with DBN and adversarial training
I: Initialize models Wa by Wa
2: for E {1,-... ... E} do
3: for mini-batch B = [(x;, y) } in Dr do
4: for w2.- E W, in paralicl do
5:
Set f use clean BN
6:
[H LenyE LoE(s(z; (1;
7:
Set f
noise BN
8:
=
9:
for E B do
10:
Perturb 3 = x + with 8 +- aag LoE(f
11:
Bu v)}
12:
L 1+
13:
Wy,r + Wsrr 7ow,
Return We
Data. Both CIFARIO and Digits are 10-way classification tasks. We follow the non-i.i.d benchmark of Li et al. (2020b) to extract 10 classes from DomainNet, which is publicly available in FedBN codes. To illustrate the multiple-domate datasets, we sample several images from Digits and DomainNet datasets in Fig. 10.
s/n
UEZE
ty1
9
23
97
9
dpon
kmpa
Figure 10: Sample images from multiple domain datasets.
stochastic gradient descent (SGD) with 0.9 momentum and 5 x 10-4 weight decay. Dataset specific settings are stated as follows. 10: Following HeteroFL. (Diao et al., 2021), train with 5 local epochs and 400 global
Hyper-parameters. In general. for local optimization
14

---

Published as a conference paper at ICLR 2022
Table 2: Network architecture for Digits dataset.
Layer
Details
feature extractor
conv1 bnl
Conv2D(64,, kernel size=5, stride=1, padding=2)
DBN2D, RELU, MaxPool2D(kerrel size=2, stride=2)
conv2
Conv2D(64, kernel size=5, stride=1, padding=2)
bn2
DBN2D, ReLU, (ax size=2, stride=2)
conv3
Conv2D(128, kernel size=5, stride=1, padding=2))
bn.3
DBN2D, ReLU
classifier
fel
FC(2048) DBN2D, ReLU FC(512) DBNID, ReLU FC(10)
bn4
fc2
bn5
fe3
communication rounds. Globally, initialize the learning rate 0.01 and adjust the learning rate 150, 250 communication rounds with a scale rate of 0.1. Locally, we use a larger batch size of 128, to speed up the training in simulation. Digits: We use a cosine annealing learning rate decaying from 0.1 to across 400 global communication rounds. SGD is executed with one epoch for each local client. DomainNet: We use a constant learning rate 0.01 and run 400 communication rounds in total. Similar to Digits, SGD is executed with one epoch for each local client.
Table 3: Network architecture for DomainNet dataset.
Layer
Details
feature extractor
conv bnl
C6ny2D(64, kernel size=11, stride=4, padding-2)
DBN2D. ReLU, MaxPool2d(kernel size=3, stride=2)
conv2
Co9y2D(192, kernel size=5, stride=1, padding=2)
bn2
DBN2D, ReLUU, MaxPow/2d(kermel size=3, stride=2)
conv3
Conv2D(384, kernel size=3, stride=1, padding=1)
bn3
DBN2D, ReLU
conv4
Conv2D(256, kernel size=3, stride=1, padding=))
bn4
DBN2D, ReLU
conv5
Conv2D(256, kernel size=3, stride=1, padding=l) DBN2D, ReLU, MaxPool2ddkeemel sizc=3, stride=2)
bn5
avgpool
AdaptiveAvyPPoo2266.. 6)
classifier
fel bnó
FC(4096) DBNID, ReLU FC(4096) DBNID, ReLU FC(10)
fe2
bn7
Network architectures. Architectures of modified AlexNet (for DomaiNet) and CNN (for Digits) can be found in (Li et al., 2020b) and public codes 3. For reader's 's reference, we provide the layer details in in Tables 2 and 3. For the convolutional layer (Conv2D or Conv ID), the first argument is the number channel. For fully connected layer (FC), we list the number of hidden units as the first argument. The implementation of the preactivated ResNet be found in the repository of HeteroFL (Diao et al., 2021).
natch-normallzztion for eustomizable model sizes. As pointed out by Diao al. (2021), HeteroFL relies on mini-batch batch-normalization to stabilize the training with multiple model widths and
httpps://gg thub. cow.ned-ommdddddeeoeoooo r/FedNN
15

---

Published as a conference paper at ICLR 2022
compute the statistics afterwards. Another advantage of such strategy is federation of multi- clients. Li et al. (2020b) showed that using local natch-normalization helps model personalization for non-i.i. local features. As min-batch BN statistics simulate the local BN idea, it should enjoy the similar benefit. To eliminate the potential biases in personalization or convergence caused by different estimation methods of BN statistic, we let all compared algorithms using the same mini-batch strategy. To reveal the effect of different BN statistic solutions for Split-Mix, we provide a detailed ablation study regarding the estimation of BN statistics in Appendix A.3.
Loss function for class non-i.id FL. As some classes are missing locally in non-i.i.d setting, class- specific parameters in the classifier head may be updated without proper supervision and results in random updates. To mitigate the effect of missing classes locally, we use the same masked cross-entropy loss as introduced by HeteroFI., where absent classes are masked out.
A.3 ABLATION STUDY OF NETWORK SCALING AND BN STATISTICS
In this section, we evaluate how the network rescaling and BN statisties affect the performance. For rescaling, we consider the parameter initialization (rescale init) and layer outputs (rescale layer). For BN statistics, we consider four options. The batch average batch of data. The post average one will use the hatch average strategy during training but re-estimate the statistics using client data afterward, which was adopted by HeteroFL. To gain better BN statistics, we run the model on a training set for 20 epochs. The tracked one will track statistics during training. The locally tracked one will also track statisties but the statistics will not be shared with the for averaging, which can benefit clients' privacy and personalization (Li et al., 2020b).
will estimate the statistics by one
We report full ablation results in Table 4. 1) First we compare the use of BNs without in-training tracking. The post-average BN performs best compared to other BN choices. With similar perfor- mance, the batch average BN does not need multiple rounds of evaluation BN statisties, which more efficient for inference. 2) Then we compare the use of BNs tracked during training. Consistent with the prior study (Li et al., 2020b), locally tracked BN performs better than the globally averaged one. 3) Rescaling layer outputs barely affect the accuracy, but it could be poisonous for tracked BN statistics. 4) The initialization rescaling greatly improves the performance regardless of the choice of BN statistics.
In conclusion, either locally tracked or batch averaged BN statistics can yield both efficient and accurate performance. Post-averaged BN statistics may be preferred if post averaging is bearable for efficiency, especially for large-scale datasets. Rescaled initialization is an essential ingredient for Split-Max perform well. Layer rescale is recommended if batch average tracked BN statistics are utilized.
Table 4: Ablation study of network scaling and BN statistics Digits dataset. Accuracy of different customized widths are presented.
stat
rescale init
rescale layer
x0.125
x0.25
x0.5
x1
X
81.1%
84.3%
86.2%
87.3%
/
81.1%
84.2%
86.2%
87.2%
batch average
X
84.5%
87.5%
88.9%
89.8%
/ 84.6%
87.5%
89.0%
89.8%
X
81.2%
84.5%
86.2%
87.4%
/
81.2%
84.4%
86.2%
87.3%
post average
X
84.5%
87.5%
89.0%
89.9%
V
84.9%
87.8%
89.3%
90.2%
X 79.6%
82.8%
84.8%
85.7%
/
9.4%
10.6%
10.6%
10.6%
tracked
X
83.5%
86.4%
87.9%
88.7%
8.8%
8.8%
8.8%
10.6%
X
X
81.1%
84.3%
86.3%
87.3%
X
/
9.8%
10.6%
10.1%
11.7%
locally tracked
X
84.9%
87.7%
89.1%
90.0%
10.5%
10.6%
12.1%
11.1%
16

---

Published as a conference paper at ICLR 2022
A.4 EXPERIMENTS WITH L.I.D FL
In addition to non-i.i.d FL settings, we experiment with i.i.d FL where each client will own data of 10 classes from the CIFARIO dataset. Results using 100% and 50% training data are included in Table 5. We observe great increase in the accuracy compared to the non-i.i.d experiments, which a common phenomenon that non-iid FI. will perform worse globally. The similar performance degradation observed in (Diao et al., 2021), well. In Table 5, our method performs better in larger widths with fewer training data, whose performance approaches that unconstrained individual FedAvg. Our method provide a monotonous relation between model size and accuracy (larger models are more accurate) and uses fewer parameters and MACs even compared to wider baseline networks, though performs worse in smaller widths because the slimmest networks are updated less frequently in budget-insufficiet clients compared to the SHeteroFL or FedAvg. Worth mention, our method uses much fewer parameters and operation counts for the same accuracy. For example, Split-Mix requires 7.2M MACs and 1.4M parameters for 81.1% accuracy while SHeteroFL needs twice of the complexity, given the 50% CIFARIO i.i.d configuration.
Table 5: Test results of customizing model width on the class non-i.i.d CIFAR 10 dataset.
Individual FedAvg
SHeteroFL
Split-Mix (ours)
width
Ace MACs #Params
Acc MACs #Params
Acc MACs #Params
CIFARIO L.d FL. (100%)
x0.125
82.2%
0.9M
0.2M
81.9%
0.9M
0.2M
80.9%
0.9M
0.2M
x0.25
86.1%
3.5M
0:7M
85.2%
3.5M
0.7M
83.4%
1.8M
0.4M
x0.5
89.8%
14.0M
2.8M
86.5%
14.0M
2.8M
85.2%
3.6M
0.7M
x1
91.0%
55.7M
41-MM
85.9%
55.7M
11.2M
86.0%
7.2M
1.4M
CIFAR1O i.i.d FL (50%)
x0.125
77.3%
0.9M
0.2M
77.2%
0.9M
0.2M
74.2%
0.9M
0.2M
x0.25
79.9%
3.5M
0-7M
79.7%
3.5M
0.7M
77.9%
1.8M
0.4M
x0.5
83.2%
44-0M
2.8M
80.1% %
14.0M
2.8M
79.5%
3.6M
0.7M
x1
84.6%
55.7M
11.2M
75.5%
55.7M
11.2M
81.1 %
7.2M
1.4M
A.5 MORE BUDGET DISTRIBUTIONS
In our experiments, we generally use an exponential budget distribution: R. = (1/2) [4k/k] Though the distribution represents the imbalance between the budget- suffeet-ufficient and gudget-inssfficinn clients, real-world applications may encounter a wider variety of budget distributions. Thus, we extend our problem assumption to budget distributions with more budget-sufficien clients where we let more groups to have x or x0.5 net training capability and with step-increase budgets where we increase budgets by a fixed step (e.g., x0.25). In addition, we consider log normal distribution which concentrate around 0.45 budget with few wider or extremely budget-insufficient clients. We partition the budget distribution into 0. 1.5ddhh bins and each client will only train the maximal compatible width varying from 0.125 to 1, which greatly increases the number of slimmable subnetworks (8 now compared to previous 4). To reduce the overhead of slimmable training, we use HeteroFL instead of SHeteroFL. Fig. 11 reports the per-width accuracy for Split-Mix and SHeteorFL on the Digits dataset. For SHeteroFL, we only report the evaluated performance on trained widths. other words, if the maximal width is x0.5, (Ind. FedAvg), models individually trained for each width are reported, which ignores the width constraints by users and therefore only serves reference upper bounds. Regardless of the budget distributions, for instance, more budget-sufficiee clients (Fig. 1la) or non-exponential distributions (Fig. 1lb), Split-Mix outperforms FL with larger widths.
will not report results of the x 1 net. For Individually-trainn trained FedAvg
A.6 EFFECT OF LOWER CONTACT RATES
Because of varying communication conditions in deployment, the times that clients actively and successfully upload their models could be fewer than expected. To evaluate the robustness of customization federated algorithms, we conduct federated experiments with a varying number of active clients per round in Fig. 12. Because of the limited number of communication rounds (within
17

---

Published as a conference paper at ICLR 2022
300 rounds), the test accuracy decreases by fewer contact clients. This is common phenomenon because lower contact rates requires more communication rounds to reach the same performance as the full-contact competitors do. Though global performance generally decreases, wider networks are more accurate trained by SplitMix. One source for the advantage the modular base models in SplitMix, which can be easily distributed into different rounds for training. Rather, the wider integrated networks in SHeteroFL lower their chance to be aggregated globally and therefore their global performance declines.
find that the
budgeite
12
w0th contrai
0.6
6-133
03220.55.375 05 6.825
bedge imnmral computbe WAR
supppts
We conorra nt
42-079-1
53 DM
Fodivs
Spimtii
10f Henero7
22.5
50.0 widit PM
200.0
12:5230
220.2
125
52.0
300.
wish (S)
with
More budget-sufficient clients
(b) Step-increase budget distributions
(c) Log-normn distributions
Figure 11: Vary the budget distribution. The training budgets, i.e., width constraints, are depicted in the upper figures by group. The budget distribution name. for example, 8-4-2-1, means x 1/8. x1/4, x x1/2 and x1 width constraints for each group, respectively. The lower figures compare the performance of trained models with customized withs.
Dagtt
100%
87.5
1.85
85.0
825
2.25
E-155
80.0
SHotcrofL
SlideofL
Stateof
Figure 12: Vary the number of clients Figure 13: Convergence of different-width models. uploading models per round.
A.7 NEGATIVE IMPACT OF CONSTRAINED BUDGETS ON WIDER NETWORKS
In Fig. 2, we show how the SHeteroFL fails train wider networks with budget-constrained clients. To show the generality of such problem, we extend the experiments to Digits and CIFARIO datasets in Fig. 13. Though wider networks converges faster the beginning. they meanwhile overfit limited data in few clients and therefore their validation accuracy no longer improves.
18


Reviews:
[{"review_id": "dvb9_fuXG0J", "reviewer": "Reviewer_7P9j", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strong and Weak points\n\n- (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. \n\n- (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)\n\n- (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?\n\n- (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 ? \n\n- (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\".\n\n- (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. \n\n- (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.\n\n- (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.\n\n\n- (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.\n\n\n- (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_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_\n", "summary_of_the_review": "### Recommendation\n\nThe 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. \n\nSummary: 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.\n\n\n### Supporting my Recommendation\n\nPlease see points above.\n\n### Minor points\n\n\n- Is it fair to add as cite to \"Federated Learning\" a paper from 2019 (first sentence of the Introduction) ?\n- When describing Fig3 (in Page 4) the sentence starts with \"To validation, we investigate [...]\". Maybe it should be \"To validate this\"?\n- 6th line in Sec5: \"The dataset first is [...]\", should be \"The first dataset is [...]\"\n- 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.\n- 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. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "YjmiMj9169h", "reviewer": "Reviewer_82wJ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n- Addressing data and resource heterogeneity is definitely one of the important topics of FL.\n- 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.\n- 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.\n- Experiments are thorough, demonstrating the effectiveness of the proposed approach on multiple datasets with multiple metrics.\n\n## Weaknesses\nIt'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?\n\nI 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?\n\nAnother 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_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.", "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.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "V77U2dBt1--", "reviewer": "Reviewer_Yvwo", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths:\n1. 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.\n2. In terms of the idea \"Split-Mix\", the authors present empirical evidence of slim models and wide models, including convergence and gradient saliency. \n3. 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.\n\n## Weaknesses:\n1. 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?\n2. 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?\n3. The idea of BN layer-wise mixing is kind of confusing. The reason of sharing all parameters except the batch-normalisation requires further elaboration. \n\n## Detail Comments:\n1. The notations in the three algorithms lack necessary explanation.\n2. 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\n2. 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. \n3. 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. \n4. 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? \n5. 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.\n6. 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? \n7. 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_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.\n\nSplit-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. \n\nExperiments show that Split-Mix achieves better accuracy than naive approaches. Moreover, with customisation, the models are smaller and more robust under budget constraints.", "summary_of_the_review": "\nThis paper proposes to tackle the heterogeneous budget in FL scenario. A weak accept is given based on two reasons. \n(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. \n\n(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. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "87bujcIKtVr", "reviewer": "Reviewer_emNn", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths\n\n[Problem / Motivation] \n\nThe problem of heterogeneous resources in federated training is of utter importance and very challenging both from the systems and algorithms perspective. \n\n[Related Work] \n\nThe authors have a valuable overview of the related work.\n\n[Novelty]\n\nThe paper introduces some interesting ideas based on existing work (HeteroFL and adversarial robustness).\n\n### Weaknesses \n\n[Significance] \n\n- The evaluation results are heavily based on the assumption of exponentially distributed budgets. \n 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.\n\n- 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.\n\n\n[Clarity] \n\nSome arguments of the paper are not explained.\n\n- 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.\n\n- 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.\n\n\n[Evaluation]\n\nSome results are questionable / counter-intuitive.\n\n- Table 1: Why does the accuracy of SHeteroFL not monotonically increase with FLOPS and #Params (as it is for Split-Mix) ?\n\n- Figure 9b: Why λ=0 gives less SA than λ={0.1, 0.2} ? \n\n\n[Presentation]\n\nPaper presentation needs improvement.\n\n- 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.\n\n- Gradient saliency should be defined or described in a few lines.\n\n- Please mention “adversarial robustness” instead of “robustness” at least once in the beginning of the paper.\n\n- There is a large number of typos (mainly after Section 4). A sample:\n  - 4.1 (beginning): “To motivation” -> “For motivation”\n  - “convergence slower” -> “converge slower”\n  - “To validation” -> “For validation”\n  - “Last, it worth note that” -> “Lastly, it is worth noting that”\n\n", "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.", "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).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

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)