Year: 2022

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.

heading: 1 INTRODUCTION
text: Federated learning (Konečný et al., 2015) is a distributed learning paradigm that leverages data from remote participants and aggregates their knowledge without requiring their raw data to be transferred to a central server, thereby largely reducing the concerns from data security and privacy. FedAvg (McMahan et al., 2017) is among the most popular federated instantiations, which aggregates knowledge by averaging homogeneous model parameters.
In many practical federated learning scenarios, the participants usually exhibit a variety of heterogeneity. The training datasets from participants can be non-i.i.d. and therefore lead to data heterogeneity, which has received extensive attention in the community (Li et al., 2020b; Fallah et al., 2020; Hong et al., 2021; Li et al., 2021). In addition, the participants can also have resource heterogeneity, i.e., participants have different hardware resources available, such as memory, computing power, and network bandwidth (Ignatov et al., 2018). For example, in a learning task for face recognition, clients may use different types of devices (e.g., desktop computers, different generations of tablets or smartphones) to participate in learning, which thereby leads to drastic differences in computational resources for training models, and in sizes of models that can be loaded to the memory and transferred through the network. The disparity of hardware resources results in varying budgets in both training and inference stages (Xu et al., 2021). To accommodate different model sizes, one can turn to more 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).
Yet besides heterogeneity, the deployment of models usually has a great amount of dynamics. One type of dynamics is directly related to the resources of participants’ devices. For example, memory consumptions and computational capabilities may vary drastically between training and inference stages, and the same dynamics in hardware can also happen during inference (Yu et al., 2018; Xu et al., 2021), depending on how the resource allocation of the running programs is prioritized on a participant’s device. Another example of dynamics is the requirements of properties of the model, such as adversarial robustness. In autonomous driving, for instance, cloudy and rainy weather can
(a) Illustration of FedAvg (McMahan et al., 2017) with a device-incompatible model and a heterogenousarchitecture variant (HeteroFL) (Diao et al., 2021) with under-trained wide models.
(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 a subnet with 1/4 channels (or widths) per layer as a base model for model-width customization. For simplicity, we denote it ×0.25 net, and a ×1 net can be split into 4 ×0.25 base models.
Figure 1: Comparison of traditional and proposed methods.
lower the quality of camera and other sensors, and therefore one would like to quickly switch to a more robust prediction model, as compared to the one used on sunny days. The dynamics along with the low latency requirements for the inference, altogether demand a model that can be customized in situ, meaning that we can instantly tailor our predictive model to a proper size that fits our available budget, and promptly adjust its robustness. Most existing federate learning models used a rather static model architecture that prevents such customization. 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 in a budget-limited device. Moreover, the trade-off between model adversarial robustness and generalization performance can hardly be reconcilable (Tsipras et al., 2019; Merchant et al., 2020), which further complicates the customization.
To address the aforementioned challenges from heterogeneity and dynamics, we study a novel SplitMix federated learning approach for in-situ customization, which follows a communication-efficient paradigm and yields a model whose size and robustness can be efficiently customized, without additional training. 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 to construct desired widths and required 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 to study training a model with the capability of in-situ customization with heterogeneous local computation budgets. Illustrated in Fig. 1a, we show that conventional static-model solutions, such as FedAvg, cannot provide customized models for dynamic client environments, and that its heterogeneous-architecture variant (HeteroFL) suffers from under-training in its large models due to local budget constraints.
• To address the challenge, we propose a novel Split-Mix framework that aggregates knowledge from heterogeneous clients into a width- and robustness-adjustable model structure. Remarkably, due to fewer parameters and modular nature, our framework is not only efficient in federated communication and flexibly adaptable to various client budgets during training, but also efficient and flexible in storage, model loading and execution during inference.
• Empirically, we 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, i.e., clients’ training datasets suffer from distributional shifts.

heading: 2 RELATED WORK
text: Heterogeneous Federated Learning. Federated learning with budget-insufficient devices has attracted a great amount of attention. For example, FedDistill (Lin et al., 2020) used logit averages to distill a 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 et al. (2020) introduced a distillation-based method after aggregating private representations from all participants. The method closely resembles centralized learning because all 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. FedEnsemble (Shi et al., 2021) is 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 in FL can outperform a single base model in testing. A critical difference between the proposed approach and FedEnsemble comes from the challenging problem setting of constrained heterogeneous 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 heterogeneous-architecture 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 et al. (2020), where an adjustable-conditional layer together with decoupled batch-normalization were used. The conditional layer enables continuous trade-off but also brings in over 50% more parameters. In comparison, a 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 on 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 at 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; 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, we consider a more challenging scenario with distributed and non-sharable data and heterogeneous client capabilities.

heading: 3 PROBLEM SETTING
text: The goal of this work is to develop a federated learning (FL) strategy, that yields a global model, which can be instantly customized for dynamic needs. Formally, FL minimizes the objective 1∑ k |Dk| ∑K k=1 ∑ (x,y)∈Dk L(f(x;W ), y) with distributed data {Dk} K k=1 across K participants/clients, where L is loss function, e.g., cross-entropy loss, f is a model parameterized by W , and x, y are a sample and a label, respectively. In this paper, we consider two kinds of customization: adjustable model sizes to trade off between execution complexity and accuracy, adjustable robustness for dynamic security sensitivity. Though other types of customization could be considered, for example, neural-architecture search or fairness, we believe the two types considered are quite representative for both architecture and non-architecture knowledge (Wang et al., 2020).
When training customizable/adjustable models, significant challengs 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 {Rk ∈ (0, 1]}Kk=1, defined by width ratios w.r.t. the original model, i.e., ×Rk 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. 2) Imbalance computation budgets. In many applications, there are usually a significant number of devices with insufficient computational budgets (Enge, 2021). Without loss of generality, we assume exponentially distributed budgets in uniformly divided client groups: Rk = (1/2)⌈4k/K⌉, i.e., the first group with 1/4 clients is capable for a 1-width and the rest are for 0.5, 0.25, 0.125-widths respectively. Such a budget distribution could result in additional challenges in the presence of data heterogeneity, e.g., non-i.i.d. features, Dk ∼ Di for ordered domain i, since training a single model on multiple domains (Li et al., 2020b) or models in a single domain (due to budget constraints) (Hong et al., 2021) are known to be suboptimal on average performance. Finally, following the standard FL setting (McMahan et al., 2017), only model parameters can be shared to protect privacy.
In summary, the problem setting 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 scarce budget-sufficient clients; 2) data heterogeneity, where the under-trained large models may perform poorly on unseen in-domain or out-of-domain samples.

heading: 4 METHOD
text: 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.

heading: 4.1 CUSTOMIZE MODEL SIZE
text: For motivation, we first set up a non-i.i.d. FL experiment to compare model performance with/without budget constraint, where the same 6-domain DomainNet images are uniformly distributed into 50 clients for both cases (refer to experimental details in Section 5).
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 on all clients (solid lines in Fig. 2). The results justified the motivation of training wider networks than slimmer ones.
Case 2: FL with budget constraint. Following the above budget constraint, i.e., Rk = (1/2)⌈4k/K⌉, we deploy HeteroFL to train budget-compatible prototype models locally. We note that data are not fully used by HeteroFL: each model prototype is only trained by 1/4 of clients, when clients of the Rk = 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 (e.g., ×1 net) converge to validation accuracy lower than not only the FedAvg counterparts, but also the slimmer ×0.25 net, showing that the widest model may not be a good candidate model. Intuitively, as a consequence of the budget constraint, fewer training clients are able to participate in the learning of wider networks, where the reduced sample size or model size greatly impacted the generalization. Since FL does not allow global data sharing to mitigate the issue, the key is to study how the scarcity of data or smaller width affects model generalization in FL. We hypothesize that a model may pick up signals from more features for reasoning a class if trained with more data. For validation, we investigate the features that draw significant attention of different networks in Fig. 3. We found that both FedAvg and SHeteroFL ×1 nets extract richer features than the slimmer ×0.125 net. This may partially explain why the slimmer network generalizes poorly. On the other side, though the SHeteroFL×1 net has the same capacity as the FedAvg×1 net and thereby extracted similarly-rich features, its accuracy is lower than FedAvg. This suggests that the features learned from limited data could be less generalizable. For example, the SHeteroFL loses attention of edges of the feather in the fifth image. Thus, more diverse or higher-quality features are desired. The importance of features for generalization was also echoed by a recent study (Allen-Zhu et al., 2020), The study showed that encouraging neural networks to extract diverse features can boost testing accuracy.
Figure 4: Different gradient saliency generated by 6 ×0.125 base networks on the same image. More images in Fig. 14.
Diverse bases from random initializations. Motivated by the above observations, we propose to train multiple base models from different random initializations to encourage their diversity. The diverse bases likely extract different features from the same image (see Fig. 4). The overall FL algorithm is summarized in Algorithm 1. To craft a wide model, we combine the output of a number of base models until the size of the ensemble reaches the same as the number of channels. Therefore, the ensemble predicts based on a variety of features, resembling an integrated wide network. To facilitate these base models to extract generalizable features, we want to train base models in all clients which are slim and therefore adaptable to all clients. Inspired by FedEnsemble (Shi et al., 2021), we randomly sample a set of bases for each client in Algorithm 2. In each client, we speed up the local training of multiple independent base models by direct parallelization (see Algorithm 3). Note however that the parallelism is limited by the maximal width in local, as wider channels result in larger intermediate layer cache (activations). We thus limit the total channels of ensemble by client’s budgets. 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. Different from sequential training, the parallelism speeds up the local training and reduces the possible client blocking time on communication. In addition, slimmer base models can conform a variety of client budgets while maintaining the same model width.
Algorithm 1 Federated Split-Mix Learning Input: Client datasets {Dk}Kk=1, the number of total communication rounds T , M = ⌊1/r⌋ randomly initialized base ×r nets parameterized by {wi,r}Mi=1, client budgets {Rk}Kk=1, the number of local epochs E, learning rates {ηt}Tt=1 1: Initialize {w0i,r}Mi=1, model indexes P = Shuffle([1, · · · ,M ]) and current index p = 1 2: for round t ∈ {1, · · · , T} do 3: Initialize W t ← {wti,r ← 0}Mi=1 and aggregation weights ci = 0 for i ∈ {1, . . . ,M} 4: for k ∈ {1, · · · ,K} do 5: Sample base models W t−1k , p← SampleBaseModels(P, p, ⌊Rk/r⌋,W t−1)
6: Send W t−1k to client k and train Ŵ t k ← LocalTrain(W t−1k , Dk, E, ηt) 7: Aggregate ŵt,ki,r ∈ Ŵ t k to the server: w t i,r ← wti,r + ŵt,ki,r |Dk|, ci ← ci + |Dk| 8: Server update W t ← {wti,r ← wti,r/ci}Mi=1 9: (Optional) Sort WT = {wTi,r}Mi=1 by the descending order of the validation accuracy of wTi,r
10: Output the customizable model WT = {wTi,r}Mi=1 11: Customize an R-width model by F (x;WT ) = 1
KR
∑KR i=1 f(x;w T i,r) where KR = ⌊R/r⌋
Algorithm 2 SampleBaseModels(P, p, n,W ) 1: if p > |P | then Shuffle P and p← 1 2: Initialize W = {wP [p],r} 3: if n>1 then 4: Uniformly sample n − 1 values into S from
P\{P [p]} without replacement 5: W ←W ∪ {wi,r,∀i ∈ S} 6: Return W, p+ 1
Algorithm 3 LocalTrain(Wk, Dk, E, η)
1: Initialize models Ŵk by Wk 2: for e ∈ 1, · · · , E do 3: for batch data (x, y) in Dk do 4: for ŵi,r ∈ Ŵk in parallel do 5: ŵi,r ← ŵi,r − η ∂L(f(x;ŵi,r),y)∂ŵi,r 6: Return Ŵk

heading: 4.2 CUSTOMIZE ADVERSARIAL ROBUSTNESS
text: In many existing robustness learning approaches, model robustness and generalization are typically irreconcilable, and a trade-off between the two has to be determined when applying robust models. In this section, we study the customization of model robustness, and specifically adversarial robustness. The adversarial robustness of a model reflects its prediction accuracy when subtle and predictionflipping adversarial noise presents. Recently, adversarial training (AT) (Madry et al., 2018) becomes a popular approach that uses on-the-fly adversarial samples as augmentation to improve adversarial robustness. Formally, AT minimizes the following augmented loss:
L(f) = (1− λn)LCE(f(x), y) + λn max∥δ∥∞≤ϵ LCE(f(x+ δ), y), (1)
where δ is a subtle adversarial noise which maximizes the loss in a given constraint ϵ. The addition of δ transfers a clean sample x into an adversarial sample. In Eq. (1), the first term characterizes the cross-entropy loss (LCE) on clean samples while the later one on adversarial samples. The hyper-parameter λn trades off robustness and accuracy. When λn = 0 or 1, the optimization yields a standard-training (ST) model or an AT model, respectively. Since an adversarially-trained model is commonly less accurate in predicting standard images (Tsipras et al., 2019), there is usually no such a sweet point of λn 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 standard performance and adversarial robustness are irreconcilable, we can directly use two separated 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; Xie et al., 2020), sharing all parameters except the batch-normalization (BN) layers can maximize robustness and accu-
racy 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 cross-entropy loss of Eq. (1) is computed with clean BN (BNc) merely and the second adversarial loss is computed with noised BN (BNn). 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 on parameter updates.
Layer-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 cache of double size of intermediate outputs, a gate function can be used to adaptively switch between two BNs like (Liu et al., 2020). Inspired by the method, we further propose a simple parameter-free method by weighted-averaging outputs of each BN layer (see Fig. 5):
DBN(x) = (1− λ)BNc(x) + λBNn(x), (2)
given a BN-layer input x. We note that the averaging strategy is entirely training-free and does not use extra parameters, and the customization weight λ is intuitive for trade-offs between SA and RA.
Lastly, it is remarkable that the DBN structure is rather lightweight in terms of model complexity. As investigated in (Yu et al., 2018) (Table 2), the parameters of BN is no more than 1% in popular deep architectures, e.g., ResNet or MobileNet. Therefore, we can plug DBN into base models in place of BN, replace Algorithm 3 with Algorithm 4, and jointly customize robustness and model widths.

heading: 5 EMPIRICAL STUDIES
text: We design experiments to compare the proposed method against FL classification benchmarks. Configurations. For i.i.d configuration, we use CIFAR10 dataset (Krizhevsky, 2009) with preactivated
ResNet (PreResNet18) (He et al., 2016). CIFAR10 contains over 50, 000 32 × 32 images of 10 classes. The CIFAR10 data are uniformly split into 100 clients. For non-i.i.d configuration, we use Digits with a CNN defined and DomainNet datasets (Li et al., 2020b) with AlexNet extended with BN layers after each convolutional or linear layer (Li et al., 2020b). Data. The first dataset is a subset (30%) of Digits, a benchmark for domain adaption (Peng et al., 2019b). Digits has 28× 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., 2011), 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 × 256 real-world images: Clipart, Infograph, 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-parameters to Appendix A, and focus on discussing the results.

heading: 5.1 CUSTOMIZE MODEL SIZES
text: 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: Rk = (1/2)⌈4k/K⌉. 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-incompatible baseline, we (re-)train networks from scratch by FedAvg to obtain individual models with different widths. The state-of-the-art heterogeneous-architecture FL method is HeteroFL (Diao et al., 2021), which trains different slim models in different clients. For a fair comparison, we extend HeteroFL with bounded-slimmable training in clients who can afford the larger models, named SHeteroFL. For example, if a client in HeteroFL can afford ×0.5 net, then the client meanwhile trains ×0.25 net and other smaller subnets in slimmable training manner (Yu et al., 2018) by SHeteroFL.
Convergence. In Fig. 6, we compare the convergence curves of full-width models on three datasets. For FedAvg, the budget-compatible width is ×0.125. Both for SHeteroFL and Split-Mix, the fullwidth is ×1. 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 ×1 net is more than one, as 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-addition operations) and model size in parameter numbers. With the same width, our method outperforms SHeteroFL, using much fewer parameters and thus conducts inference in much lower latency. Remarkably, compared to the best model by SHeteroFL, e.g., 81.8% ×1 net in CIFAR10, Split-Mix uses only 1.8% parameters and 1.6% MACs (×0.125 net) to reach a similar level of test accuracy. We notice that SHeteroFL has a much lower test accuracy with ×0.125 net on the CIFAR10 dataset. By investigating the loss curves, we find that the inference between parameter-shared different prototypes results in unstable convergence of the ×0.125 net. Attributed to the independent splitting, our method is more stable on convergence with different widths. Remarkably, Split-Mix only requires 12.7% parameters and 19.8% MACs (by ×1 ensemble) to achieve the comparable accuracy as the ×1 individual network on Digits. Results on CIFAR10 and DomainNet show a potential limitation of our method. The accuracy of Split-Mix is 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 in budget-limited FL, since wide individual models cannot be trained in this case.
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. 7. Conforming the budget constraints, our method outputs more accurate full-width models, transfer fewer parameters per round and execute fewer multiplication-and-add operations for gradient descent than SHeteroFL, either in terms of average or variance. Because only the individual ×0.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. 8 (Left). We 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 a 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-insufficient clients, e.g., in the clipart, infograph and painting domains. Hence, SHeteroFL leaves a large amount of parameters under-trained and suffers from larger accuracy losses in the three domains and wider models. In comparison, Split-Mix not only has less accuracy drop but is also more stable in all domains.

heading: 5.2 CUSTOMIZE ROBUSTNESS
text: 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 (ϵ, n) = (8/255, 7), and attack inner-loop step size 2/255, for training, validation, and test. For simplicity, we temporarily relax the budget constraint Rk and let the base model be ×1 net. For comparison, we extend FedAvg with AT which yields individual models optimized with different trade-off variables, i.e., λn ∈ {0, 0.1, 0.2, 0.3, 0.5, 1}. Also, we extend FedAvg with state-of-the-art in-situ trade-off method, OAT (Wang et al., 2020), as a 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 metrics: 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 effectiveness by comparing the RA-SA curves, which is better approaching the right-upper corner. To plot the curves for FedAvg+OAT and Split-Mix+DAT, we vary their condition variable λ in {0, 0.2, 0.5, 0.8, 1}. Trade-off curves are presented in Fig. 9 (a) and (b). Since the naive extension of OAT with FedAvg adopts heterogeneous objectives and over-parameterization of conditional layers which suffers from convergence issues, its adversarial training does not converge and RA is incredibly poor in some cases. As a result, the trade-off curve is not smooth. Instead, the proposed Split-Mix+DAT method has a smoother trade-off curve without heavy conditional training or over-parameterization. By training in one pass, Split-Mix+DAT even outperforms and is more efficient than the FedAvg+AT baselines.
Joint customization of width and robustness under budget constraints. Now we consider the width customization and the training budgets as Section 5.1. Due to the constraint, FedAvg can only train ×0.125 net. We omit OAT for its unstable convergence and use SplitMixDAT as a short name of Split-Mix+DAT. In Fig. 9 (c) and (d), the trade-off curves with different widths are depicted. As the width increases, both RA and SA of SplitMixDAT are improved, when they are smoothly traded off. The results demonstrate the flexibility and the modular nature of our method.

heading: 6 CONCLUSION
text: In this paper, we proposed a novel federated learning approach for in-situ and on-demand customization to address challenges arise from resource heterogeneity and inference dyanmics. We proposed a Split-Mix strategy that efficiently transfer clients’ knowledge to collaboratively learn a customizable model. Extensive experiments demonstrate the effectiveness of the principle in adjusting model widths and robustness when much fewer parameters are used compared to baselines.

heading: ACKNOWLEDGMENTS
text: This material is based in part upon work supported by the National Institute of Aging 1RF1AG072449, Office of Naval Research N00014-20-1-2382, National Science Foundation under Grant IIS-1749940. Z.W. is supported by the U.S. Army Research Laboratory Cooperative Research Agreement W911NF17-2-0196 (IOBT REIGN).

heading: A EXPERIMENTS
text: In this section, we provide more details about our experiments and additional evaluation results.

heading: A.1 PARALLEL IMPLEMENTATION OF SPLIT-MIX AND CONVERGENCE
text: In this section, we elaborate on the implementation and efficiency of Split-Mix. In Fig. 10, we conceptually compare the training by three methods, when two clients capable of training ×1 and ×0.5 are considered. FedAvg can train the ×1 net on the ×1-capable client but not on the ×0.5- capable client. Through parameter sharing among different model widths, SHeteroFL can train multiple widths in a sequential manner and can fit into ×0.5-capable clients. Unlike SHeteroFL, Split-Mix trains four base models in one parallel pass and therefore is the most efficient and flexible method. Remarkably, in the worst case, training ×1 net of Split-Mix is as efficient as FedAvg and more efficient than SHeteroFL, if all the base weights, wl0,r, . . . , w l 3,r ∈ Rr×r (where r = 0.25 here), are embedded into a 4r × 4r weight matrix1. When base models are embedded into a full net, non-trainable parameters can be masked out (grey areas) to avoid interference between base models.
With the aforementioned implementation, we compare the convergence versus the wall-clock time in Fig. 11. We implement all algorithms in PyTorch 1.4.1 run on a single NVIDIA RTX A5000 GPU and a 104-thread CPU. Fig. 11 shows the elapsed computation time from the initialization to a maximal number of iterations. The maximal number of iterations is set to be the same for all methods, such that it is easier to compare the stopping time. In Fig. 11, Split-Mix is much more efficient than the SheteroFL. Note that FedAvg can only train the ×0.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.
1For the simplicity of notations, we assume the width of layer l is r, though the width could vary by layer in general.

heading: A.2 EXPERIMENTAL CONFIGURATIONS
text: 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 4 can be solved by an n-step projected gradient descent, and is commonly known as PGD attack (Madry et al., 2018).
Algorithm 4 LocalTrain(Wk, Dk, E, η) with DBN and adversarial training
1: Initialize models Ŵk by Wk 2: for e ∈ {1, · · · , E} do 3: for mini-batch B = {(x, y)} in Dk do 4: for ŵi,r ∈ Ŵk in parallel do 5: Set f to use clean BN 6: L← 1|B| ∑ (x,y)∈B LCE(f(x; ŵi,r), y) 7: Set f to use noise BN 8: B̃ = ∅ 9: for x ∈ B do
10: Perturb x̃ = x+ δ with δ ← argmax∥δ∥∞≤ϵ LCE(f(x+ δ; ŵi,r)), y) 11: B̃ ← B̃ ∪ {(x̃, y)} 12: L← 1
2 { L+ 1|B̃| ∑ (x̃,y)∈B̃ [LCE(f(x̃; ŵi,r), y)] } 13: ŵi,r ← ŵi,r − η ∂L∂ŵi,r 14: Return Ŵk
Data. Both CIFAR10 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-domain datasets, we sample several images from Digits and DomainNet datasets in Fig. 12.
Hyper-parameters. In general, for local optimization we use stochastic gradient descent (SGD) with 0.9 momentum and 5 × 10−4 weight decay. Dataset specific settings are stated as follows. CIFAR10: Following HeteroFL (Diao et al., 2021), we train with 5 local epochs and 400 global
communication rounds. Globally, we initialize the learning rate as 0.01 and adjust the learning rate at 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 0 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.
Network architectures. Architectures of modified AlexNet (for DomaiNet) and CNN (for Digits) can be found in (Li et al., 2020b) and public codes 2. For reader’s reference, we provide the layer details in in Tables 2 and 3. For the convolutional layer (Conv2D or Conv1D), the first argument is the number channel. For a fully connected layer (FC), we list the number of hidden units as the first argument. The implementation of the preactivated ResNet can be found in the repository of HeteroFL (Diao et al., 2021).
2https://github.com/med-air/FedBN

heading: A.3 CUSTOMIZE MODEL SIZES
text: Batch-normalization. As pointed out by Diao et al. (2021), HeteroFL relies on mini-batch batchnormalization to stabilize the training and compute the statistics afterwards. To eliminate the impact of biased mini-batch BN statistics, we let all compared algorithms using the same strategy. Another advantage of such strategy is on federation of multi-domain clients. Li et al. (2020b) showed that using local batch-normalization helps model personalization for local features. As min-batch BN statistics simulate the local BN idea, which should enjoy the similar benefit.
Saliency is extracted using the captum package3. We provide more visualization of the saliency of networks of varying widths.

heading: A.4 EXPERIMENTS WITH CLASS NON-I.I.D FL
text: In addition to feature non-i.i.d FL settings, we experiment with class non-i.i.d FL where each client will randomly select 3 out of 10 classes from the CIFAR10 dataset. To mitigate the effect of missing classes locally, we use the same masked cross-entropy as introduced by HeteroFL (Diao et al., 2021). Results are included in Table 4. We observe a great drop in the accuracy compared to the i.i.d experiments, which is a common phenomenon that non-iid FL will perform worse globally. The similar performance degradation was observed in (Diao et al., 2021), as well. Besides the degradation, our method outperforms baselines consistently. Furthermore, our method performs better in larger widths, whose performance approaches that of unconstrained individual FedAvg.

heading: A.5 MORE BUDGET DISTRIBUTIONS
text: In our experiments, we generally use an exponential budget distribution: Rk = (1/2)⌈4k/K⌉. Though the distribution represents the imbalance between the budget-sufficient and budget-insufficient clients, real-world applications may encounter a wider variety of budget distributions. Thus, we extend our problem assumption to budget distributions with more budget-sufficient clients where we let
3https://github.com/pytorch/captum/tree/master
more groups to have ×1 or ×0.5 net training capability and with step-increase budgets where we increase budgets by a fixed step (e.g., ×0.25). Fig. 16 reports the per-width accuracy for Split-Mix and SHeteorFL on the Digits dataset. For SHeteroFL, we only report the evaluated performance on trained withs. In other words, if the maximal width is ×0.5, we will not report results of the ×1 net. For Individually-trained FedAvg (Ind. FedAvg), models individually trained for each width are reported, which ignores the width constraints by users and therefore only serves as reference upper bounds. Regardless of the budget distributions, for instance, more budget-sufficient clients (Fig. 16a) or non-exponential distributions (Fig. 16b), Split-Mix outperforms SHeteroFL with larger widths.

heading: A.6 EFFECT OF LOWER CONTACT RATES
text: 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. 17. Because of the limited number of communication rounds (within 300 rounds), the test accuracy decreases by fewer contact clients. This is a 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, we find that the wider networks are more accurate if trained by SplitMix. One source for the advantage is 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.

heading: A.7 NEGATIVE IMPACT OF CONSTRAINED BUDGETS ON WIDER NETWORKS
text: In Fig. 2, we show how the SHeteroFL fails to train wider networks with budget-constrained clients. To show the generality of such a problem, we extend the experiments to Digits and CIFAR10 datasets in Fig. 18. Though wider networks converges faster at the beginning, they meanwhile overfit limited data in a few clients and therefore their validation accuracy no longer improves.

Review summaries:
This paper proposes a customisation strategy named "Split-Max" for federated learning. The authors identify the heterogeneity of devices and data in FL scenarios. They present the importance of considering devices' budgets and dynamics when dispatching training models. Split-max can adjust the model size according to the devices' budget while maintaining good accuracy and robustness.  Split-max works in steps. First, multiple base models from different initialisations are trained to improve diversity. These base models are randomly given to clients to extract generalisable features. Then, base models are aggregated to the server. Secondly, to provide devices with models with different robustness, it trains two similar models together to capture both the standard-training accuracy and adversarial-training accuracy. Then layer-wise mixing is conducted to achieve both standard accuracy and adversarial accuracy.   Experiments show that Split-Mix achieves better accuracy than naive approaches. Moreover, with customisation, the models are smaller and more robust under budget constraints.

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_ 

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.

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.