Original: FedRolex: Model-Heterogeneous Federated Learning with Rolling Sub-model Extraction
Revised: FedRolex: Model-Heterogeneous Federated Learning with Rolling Sub-Model Extraction

Original: Anonymous Author(s)

Original: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: Federated learning (FL) is a collaborative machine learning paradigm to trainmodels from decentralized private data.

Original: Most FL research focuses on modelhomogeneous setting wheremodels deployed across all the participating clientsand the  server  arerequired to be identical.
Revised: Most cross-device federated learning (FL) studies focus on the model-homogeneous setting where the global server model and local client models are identical.

Original: However, in real-world scenarios,such a requirement acts as a constraint that restricts the outreach to clients withheterogeneous device resources and unfairly excludesusers with low-end deviceswho would otherwise benefit from FL.
Revised: However, such constraint not only excludes low-end clients who would otherwise make unique contributions to model training but also restrains clients from training large models due to on-device resource bottlenecks.

Original: In this work, we propose  a simple yeteffective model-heterogeneous FL method named FedRolex to tackle this constraint.
Revised: In this work, we propose FedRolex , a partial training (PT)-based approach that enables model-heterogeneous FL and can train a global server model larger than the largest client model.

Original: Unlike the model-homogeneous scenario, the fundamental challenge ofmodel heterogeneity in FL is that different parameters of the global model aretrained on heterogeneous data distributions.
Revised: At its core,

Original: Our method addresses this challengeby rollingthe sub-model in each federated iteration so that the parameters oftheglobal model are evenly trained on the global data distribution across all devices,making it more akin to model-homogeneous training.
Revised: FedRolex employs a rolling sub-model extraction scheme that allows different parts of the global server model to be evenly trained, which mitigates the client drift induced by the inconsistency between individual client models and server model architectures.

Original: Our experiments show that FedRolex outperformsother state-of-the-art  model-heterogeneous FL methods,especially under high data-heterogeneity scenarios. We have conducted ablationstudies to show that submodel rolling is an effective technique to reduce the gapbetween model-heterogeneous andstandard model-homogeneous settings.
Revised: Empirically, we show that FedRolex outperforms state-of-the-art PTbased model-heterogeneous FL methods (e.g. Federated Dropout) and reduces the gap between model-heterogeneous and model-homogeneous FL, especially under the large-model large-dataset regime. In addition, we provide theoretical statistical analysis on its advantage over Federated Dropout.

Original: Lastly,we consider the distribution of client capabilities that is similar to real-world income distribution instead of the uniform distribution used in existing works. Ourresults show a consistent improvement in the accuracies on low-end devices, whichenhances the inclusiveness of federated learning.
Revised: Lastly, we evaluate FedRolex on an emulated real-world device distribution to show that FedRolex can enhance the inclusiveness of FL and boost the performance of low-end devices that would otherwise not benefit from FL. Our code is available at: https://github.com/MSUMLSys-Lab/FedRolex.

Original:  Introduction
Revised:  Introduction

Original: Cross-device federated learning (FL) was originally proposed by McMahan et al. (2017) as aprivacypreserving machine learning paradigm to train a machine learning model on a federation of resourceconstrained edge devices without accessing their data.
Revised: Federated learning (FL) is a machine learning paradigm that trains models from distributed clients with private data under the coordination of a central server [1, 2].

Original: This idea has far-reaching applications, whichis made urgent by modern demands for privacy. Indeed, the goal of strong data privacy is a central motivation for FL.
Revised: In this work, we focus on crossdevice FL where clients are usually resource-constrained edge devices.

Original: By storing data locally, instead of replicating them on a remote server, the attacksurface of the system is decreased (Jere et al., 2020), and by using focused ephemeral updates andearly aggregation, the communication cost is also minimized (Nishio and Yonetani, 2019).

Original: The majority of existing cross-device FL studies focus on  model-homogeneous setting, where modelsdeployed across all the participating clients and the server are required to be identical.
Revised: The majority of existing cross-device FL studies focus on the model-homogeneous setting [3–6], in which the server model and the client models across all the participating client devices are identical .

Original: In practice, device heterogeneity is a more realistic constraint to consider for deploying  in real-world environments
Revised: However, modelhomogeneous FL are confronted with two fundamental constraints: (1) device heterogeneity is a more realistic consideration when deploying FL systems in real-world applications: different client devices could have very diverse on-device resources and are only capable of training models with capacities that match their on-device resources. Having the same model on all the devices would, unfortunately, exclude clients with low-end devices who would otherwise make unique contributions to the model training from their local data; (2) state-of-the-art machine learning has moved towards large models [7] such as Transformer [8].

Original: Different edge devices have different on-device resources and, hence, are only capable of training models with capacities that match their on-device resources. Having the same model on all the devices forces a limitation that restricts the outreach to clients with heterogeneous device resources and unfairly excludes users with low-end devices who would otherwise benefit from FL [Table]
Revised: Restricting server and client models to be the same inevitably causes model-homogeneous FL to fail to train such large models due to the resource constraint of client devices. To relax the fundamental constraints of model-homogeneous FL, model-heterogeneous FL was proposed where heterogeneous models with different capacities across the server and the clients are trained during the federated training process.

Original: To relax this fundamental constraint in standard FL 
Revised: One primary challenge in model-heterogeneous FL is

Original: [Table]
Revised: [Table]

Original: recent works proposed model-heterogeneous
Revised: the aggregation of heterogeneous client models.

Original: FL
Revised: To address this challenge, knowledge distillation (KD)-based approaches have been proposed [9–12]

Original: [Table]
Revised: [Table]

Original: where models with different capacities are trained on edge devices with heterogeneous resourcesduring the federated training process.
Revised: in which the client models serve as teachers and the server ensembles the knowledge distilled from the individual client models.

Original: One fundamental challenge in model-heterogeneous FL isthe aggregation of heterogeneous client models. Lopes et al. ; He et al. ; Cho et al.(2022) proposed to use knowledge distillation (KD) (Hinton et al., 2015)
Revised: However, KD-based approaches in general require public data on the server to achieve competitive model accuracy

Original: [Table]
Revised: [Table]

Original: where the client modelsserve as teachers and the server ensembles the knowledge distilled from the individual client models. However, KD-based methods generally require access to public data toachieve competitive modelaccuracy, where the accessibility of high-quality public data is not always practical.
Revised: whereas the desired public data may not be always available in practice. Moreover, since KD-based approaches need the individual client models (whole models or prediction layers) or their outputs to be sent to the server [Table]

Original: Moreover, KDbecomes computationally expensive when the cohort size in each FL round scales up.
Revised: they are incompatible with secure aggregation protocols [16] [Table] which limits their privacy guarantee.

Original: To remove thedependency on public data, another set of studies focuses on designing partial training (PT) basedtechniques such as federated dropout (Caldas et al., 2018), ordered dropout (Horvath et al., 2021),and static sub-model extraction (Diao et al., 2020), where each client trains a smaller sub-model ofthe larger global model on the server.
Revised: To remove the dependency on public data and ensure compatibility with secure aggregation, partial training (PT)-based approaches such as random sub-model extraction (Federated Dropout [13]) and static sub-model extraction (HeteroFL [14] [Table] FjORD [15]) were proposed.

Original:  However, the fundamental issue of existing PT-based methodsis that the sub-models are extracted in ways  such that the parameters of the global  modelobtained byaggregating the extracted sub-models are not evenly trained, making the global model vulnerable to aphenomenon called client drift (Wang et al., 2021) where it is used to describe the error in gradientestimation due to the differences in data distributions between clients.
Revised: In these approaches, each client trains a smaller sub-model extracted from the larger global server model [Table] and the server model is updated by aggregating those trained sub-models. However, the fundamental issue of existing PT-based methods is that the sub-models are extracted in ways (either random or static) such that the parameters of the global server model are not evenly trained .

Original: However, this estimation errorcan originate from any discrepancy across the federation. For instance, in model heterogeneousfederated learning, client drift will also occur due to the inconsistency between the clients’ individualmodel architectures in addition to the difference in private data distributions. This also becomes moreprominent when the cohort size in each communication round is small or data heterogeneity acrossclients becomes more severe, resulting in lower global model accuracy.
Revised: This makes the server model vulnerable to client drift 1 induced by the inconsistency between individual client model and server model architectures – a unique challenge of model-heterogeneous FL.

Original: In this work, we propose a simple yet effective PT-based model-heterogeneous FL method named FedRolex that tackles the fundamental issue of existing PT-based methods.
Revised: To tackle the fundamental issue of existing PT-based methods, in this work, we propose a simple yet effective PT-based model-heterogeneous FL approach named FedRolex .

Original: Compared to existing PT-based methods 
Revised: The key difference between FedRolex and existing PT-based methods is how the sub-models are extracted for each client over communication rounds in the federated training process.

Original: , FedRolex proposes a rolling sub-model extraction scheme, where the sub-modelis extracted from the global  model using a rolling window that advances in each communicationround. As such, sub-modelsare extracted from different parts of the global model  in differentrounds. Since the extraction window is rolling in each round, parameters of the global model areevenly trained to minimize the client drift.
Revised: Specifically, instead of extracting sub-models in either random or static manner, FedRolex proposes a rolling sub-model extraction scheme, where the sub-model is extracted from the global server model using a rolling window that advances in each communication round. Since the window is rolling, sub-models from different parts of the global model are extracted in sequence in different rounds.

Original: Throughout communication rounds, all the parametersof the global  model are updated over the entire data distribution.
Revised: As a result, all the parameters of the global server model are evenly trained over the local data of client devices.

Original: As we show in §4.2, such arolling sub-model extraction strategy significantly reduces the gap between model-heterogeneous andmodel-homogeneous settings. We evaluate the performance of FedRolex on three datasets – CIFAR-10, CIFAR-100 (Krizhevsky,2009) and StackOverflow (tff) – and compare it against state-of-the-art KD-based modelheterogeneous FL methods FedDF (Lin et al., 2020), DS-FL (Itahara et al., 2020) and Fed-ET (Choet al., 2021) as well as PT-based model-heterogeneous FL methods HeteroFL (Diao et al., 2020) and
Revised: The proposed rolling sub-model extraction scheme, though simple, has equipped FedRolex with multifold merits compared to prior arts (Table 1): (1) FedRolex enables different parts of the global server model to be evenly trained, which mitigates the client drift induced by model heterogeneity. (2) Contrary to static sub-model extraction approaches (HeteroFL, FjORD), FedRolex is able to train a global server model that is larger than the largest client model, enabling FL to benefit from the superior performance brought by large models.

Original:  Federated Dropout (Caldas et al., 2018) 1 .
Revised: It echos some concurrent efforts in developing FL primitives to support training large server models in cross-device setting, e.g. Federated Select [17]. (3) Compared to random sub-model extraction (Federated Dropout), as we show in our theoretical statistical analysis in Section 3 and Appendix A.1, the global server model is trained more evenly by FedRolex as the expected number of rounds for FedRolex going through all the parameters of the global model for at least certain times is smaller than that of Federated Dropout. (4) FedRolex only needs to transmit the sub-model that is needed by a given client instead of the full server model to the client. This allows clients to contribute to federated training under resource constraints and reduces communication overheads.

Original: Our results show that  FedRolex consistently outperformsstate-of-the-art PT-based methods in both high and low data heterogeneity scenarios. FedRolex alsooutperforms state-of-the-art KD-based methods in low data heterogeneity scenarios as well as on themore challenging CIFAR-100 dataset in high data heterogeneity scenarios without public data.
Revised: (5) Lastly, FedRolex is fully compatible with existing secure aggregation protocols that enhance the privacy properties of FL systems. We evaluate the performance of FedRolex under two regimes: i) small-model small-dataset regime (most existing cross-device FL studies use this combination), and ii) large-model large-dataset regime

Original: Furthermore, we show that our proposed rolling sub-model extraction scheme is flexible to train a globalmodel that can be much larger than the largest client model.
Revised: (this combination echos recent efforts on pushing the frontier of cross-device FL towards training large server models on large-scale datasets [18–21]). We highlight five of our findings: (1) FedRolex consistently outperforms state-of-the-art PT-based model-heterogeneous FL methods under both small-model small-dataset and large-model large-dataset regimes (§4.1).

Original: Lastly, we conducted an experiment that uses real-world household income distribution to emulate real-world device distribution.
Revised: (2) FedRolex reduces the gap between model-heterogeneous and model-homogeneous FL, especially under large-model large-dataset regime (§4.2). (3) With FedRolex , under both regimes, having a small fraction of large-capacity models could significantly boost the global model accuracy (§4.3). (4) FedRolex is able to train a global server model that is larger than the largest client model and outperforms Federated Dropout in terms of global model accuracy (§4.4).

Original: Our resultsshow that FedRolex effectively boosts the model accuracy of low-end devices, which enhances theinclusiveness of federated learning. The code of FedRolex will be publicly available on Github.
Revised: (5) Using an emulated real-world device distribution, we show that FedRolex enhances the inclusiveness of FL and boosts the performance of low-end devices that would otherwise not benefit from FL (§4.5).

Original:  Related Work
Revised:  Related Work

Original: Our work focuses on model-heterogeneous FL. Existing methods can be generally grouped into twocategories: knowledge distillation (KD) based and partial training (PT) based methods. Knowledge Distillation (KD) based Model-Heterogeneous FL. 
Revised: Knowledge Distillation (KD)-based Model-Heterogeneous FL. One primary approach for modelheterogeneous FL in cross-device settings is based on knowledge distillation (KD) [22].

Original: The key idea of knowledge distillation (Hinton et al., 2015; Mirzadeh et al., 2020; Liu et al., 2019) is to compress a large pre-trainedmodel by teaching a smaller network, step by step, exactly what to do using this larger pre-trainednetwork. This mechanism has been adopted in model-heterogeneous FL to train the server modelfrom a federation of client models with different architectures.
Revised: In particular,

Original: For example, FedDF (Lin et al., 2020)distills knowledge from a set of classifiers trained with private data from a federation of client devices.
Revised: FedDF [10] distills knowledge from a set of classifiers trained with private data from a federation of client devices.

Original: The logit outputs of each classifier against a public dataset are then used to train a student modelat the server with KD.
Revised: The logit outputs of each classifier against an unlabeled public dataset are then used to train a student model at the server with KD.

Original: Similarly, DS-FL (Itahara et al., 2020) keeps an unlabeled public datasetat the server, aggregates, and averages the logit outputs, and broadcasts the updated model backto the clients. Local models now train on the additional pseudo-labeled data, and performance isenhanced owing to the data augmentation effect.
Revised: Similarly, DS-FL [11] utilized an unlabeled public dataset at the server and proposed a distillation-based semi-supervised FL approach to enhance performance by pseudo-labeling the public data. FedGKT [9] proposed group knowledge transfer in which knowledge is transferred from a large model in the server to clients without public data.

Original: Fed-ET (Cho et al., 2022) considers a weightedconsensus distillation approach with diversity regularization that enables  training of a large servermodel with smaller models at the clients.
Revised: Fed-ET [12] proposed a weighted consensus distillation scheme with diversity regularization that enables the training of a large server model with smaller client models.

Original: These models have several drawbacks. KD is generally acomputationally expensive task and so the costs blow up as the system scales up for larger modelsand more clients.
Revised: KD-based approaches, however, have several limitations: these often require public data to achieve competitive model accuracy. This is because model accuracy is dependent on the size of public data as well as the domain similarity of public data with client data [10, 12, 23].

Original: These are also typically restricted to classification tasks as  theydepend on logitsfor implementing KD. The methods are also likely to face deployment issues, especially for taskswith a large number of possible output classes.
Revised: Furthermore, as KD-based approaches use client model weights partially or entirely as teachers to transfer knowledge to the server, they are incompatible with secure aggregation protocols, making them vulnerable to backdoor attacks [20].

Original: Partial Training (PT) based Model-Heterogeneous FL.
Revised: Partial Training (PT)-based Model-Heterogeneous FL.

Original: There have been several studies employingdifferent methods to train models of different sizes. These generally follow PT. HeteroFL (Diaoet al., 2020) uses this concept to aggregate models of different sizes.
Revised: To address the limitations of KD-based approaches, partial training (PT) has emerged as another solution for model-heterogeneous FL.

Original: Fjord (Horvath et al., 2021) usesordered dropout, which similarly samples submodels from a global  model, trains them, and then doesa weighted aggregation. PruneFL (Jiang et al., 2022) uses model pruning to extract submodels. Thesealgorithms are easily scalable and have minimum overheads on the server.
Revised: Depending on how the sub-models are extracted from the global server model, existing PT-based methods can be in general categorized into two groups: static sub-model extraction and random sub-model extraction. HeteroFL [14] and FjORD [15] proposed static sub-model extraction schemes where sub-models are always extracted from a designated part of the global server model.

Original: However, these methodshave a significant drawback in the sense that different parts of the model are restricted to seeingupdates from a fixed pool of client data.
Revised: However, such a static extraction strategy has two key drawbacks.

Original: This limits the performance as the globalknowledge isconcentrated on a smaller portion of the model. Furthermore, the server modelcapability is restrictedto the same capability as the largest client model. Servers are usually much more computationallycapable and so should be able to use a larger model.
Revised: First, the global server model is restricted to the same size as the largest client model. As such, the size and capability of the global model are implicitly restricted by the resources of client devices, making it not able to train large models due to resource bottlenecks at client devices.

Original: In comparison, Federated Dropout (Caldaset al., 2018; Guliani et al., 2022; Ro et al., 2022) can train onlarger global models and use its entirecapability, especially when the dataset and the number of clients  are large.
Revised: Second and more importantly, under static extraction, depending on their resource demands, different sub-models can only be trained on clients whose on-device resources are matched.

Original: Here, a portion of theparameters in each layer of the global  model is dropped, leaving a smaller submodel. Submodels aretrained on the client devices according to their capacities and the updates are aggregated via weightedaveraging similar to traditional dropout.
Revised: As a consequence, part of the global server model cannot be trained on data at low-end client devices, causing different parts of the global model to be trained on data with different distributions. This would significantly degrade the performance of the global model, especially under high data heterogeneity.

Original: However, as reported by Bouacida et al.
Revised: In contrast, inspired by the dropout technique commonly used in centralized training [24], Federated Dropout [13] proposed a random sub-model extraction method.

Original: (2020), the approachworks well only with less heterogeneous data and large client pools. This isbecause of high variancein heterogeneous data when only a small subset of the client pool is selected.
Revised: Although avoiding the drawbacks of the static sub-model extraction strategy, as reported in [25], it becomes less effective when the data heterogeneity is high and the client cohort is small due to its randomness in selecting sub-models.

Original: Our Method
Revised: In this work, we propose a rolling sub-model extraction scheme that tackles the issues of both static and random sub-model extraction methods.

Original:   3.1Problem Formulation of Model-Heterogeneous FL
Revised:  Methodology 3.1 Formulation of Model-Heterogeneous FL

Original: Consider N different client devices, N , with different compute capacities { C 1 , C 2 , ..., C N } and private data  { D 1 , D 2 , ..., D N } , each with different distributions.
Revised: Let N denote N client devices with non-IID (non-identically and independently distributed) local data D = { D 1 , D 2 , ..., D N } .

Original: The objective is to train a large globalmodel characterized by the set of trainable model parameters θ . Therefore, model-heterogeneous FLcan be formulated as the followingdistributed optimization problem:
Revised: Model-homogeneous FL trains a global model of parameter θ by solving the following optimization problem:

Original: [Equation]
Revised: [Equation]

Original: where F ( θ )
Revised: with

Original: =
Revised: [Equation]

Original: E ξ ∼ D l ( θ ; ξ ) , F n ( · ) is the local objective function on the client dataset D n ={ d n, 1 , d n, 2 , d n, 3 ... d n,m n }with respect to the user defined loss function l ( · ; · ) and p n is the corre- sponding weight of the n th client such that p n ≥ 0 and (cid:80) Nn
Revised: where D n ≜ { d n, 1 , d n, 2 , d n, 3 ...d n,m n } is the set of the n th client’s local data samples and p n is its corresponding weight such that p n ≥ 0 and (cid:80) Nn

Original: =1 p n = 1 .
Revised: =1 p n = 1 .

Original: The n th device contains m ndata points and has model parameters θ n (this can be changed from one iteration to another).

Original: Thesize of θ n depends on the client’s model capacity C n andare extracted from theglobal model i.e.,θ n ⊂ θ . The local objective function F n ( · ) is thereby rewritten as:
Revised: In comparison, in model-heterogeneous FL, clients train local models with heterogeneous capacities { C 1 , C 2 , ..., C N } , and the local objective function of the n th client becomes

Original: [Equation]
Revised: [Equation]

Original:  For simplicity, we abuse the  notationof l by ignoring the other parameters in θ but not in θ n .
Revised: The size of θ n depends on the client’s model capacity C n , and the parameter θ n is obtained by selecting a sub-model from the global model θ , which can change from one iteration to another. If θ n changes, the objective function also changes. For simplicity, we use the same notation l for the loss function at all clients and iterations, though they differ between clients and iterations.

Original: Thecore of the optimization problem is to select subset θ n from the global modelparameters θ based onmodel capacity C n .
Revised: The key to model-heterogeneous FL is selecting θ n from the global model θ given model capacity C n .

Original: 3.2 Partial Training and FedRolex
Revised: 3.2 FedRolex: Model-Heterogeneous FL with Rolling Sub-Model Extraction

Revised: [Figure] As a partial training (PT)-based approach, at each participating client, FedRolex trains only a sub-model extracted from the global server model and sends the corresponding sub-model updates back to the server for update aggregation. To help understand how FedRolex works, for simplicity, Figure 1 illustrates three rounds of federated training of FedRolex on two heterogeneous clients.

Original: As discussed in Section 2, PT-based methods train heterogeneous models by extracting a subset of theglobal model  and training the smaller model on  private data.
Revised: During each iteration, the server selects a subset of the trainable variables from the global model θ ( j ) and broadcasts the selected sub-model to each participating client. The clients train the received sub-models on their private data and post their updates θ ( j ) m to the server. The server then aggregates the variables from the updates and formulates the new global model θ ( j +1) for the next round. Specifically, the design of FedRolex involves the following two primary design decisions.

Original: The primary step in PT-based methodsis this sub-model extraction.
Revised: Equation (1) shows us that we must find θ that minimizes the objective concerning all private datasets. In centralized settings, this is achieved by simply iterating over the entire data. In FL, however, we do not have this option.

Revised: Instead, the parts of the global model are dynamically trained by iterating them over the entire client federation.

Original: Consider θ ( j ) n selected from global model weights, θ for client n at thej th communication round.
Revised: Consider θ ( j ) n selected from global model weights for client n at the j th communication round.

Original: Assume that the proportion of nodes extracted from each layer in θ forclient n be β n , which belongs to the set of unique model capacities β .
Revised: Assume that the proportion of nodes extracted from each layer in θ for client n be β n , which belongs to the set of unique model capacities β .

Original: Let the total number of nodesin the i th layer be K i and the selected nodes are S ( j ) i .
Revised: Let the total number of nodes in the i th layer be K i and the selected nodes are S ( j ) i .

Original: Then the parameters θ ( j ) n, [ i, · ] for the i th layer inextracted model parameters θ ( j ) n consist of { θ ( j ) n, [ i, S ( j ) i, 0 ] , θ ( j ) n, [ i, S ( j ) i, 1 ] , ..., θ ( j ) n, [ i, S ( j ) i,Ki  ] } .
Revised: Then the parameters θ ( j ) n, [ i, · ] for the i th layer in extracted model parameters θ ( j ) n consist of { θ ( j ) n, [ i, S ( j ) i, 0 ] , θ ( j ) n, [ i, S ( j ) i, 1 ] , ..., θ ( j ) n, [ i, S ( j ) i,Ki − 1 ] } .

Revised: Intuitively, we can think of this as iterating over the model instead of the data.

Revised: In FedRolex , the global model is rolled before the subset is extracted.

Original: Let, M ⊂ N be the set of selected clients from the client pool from which the server pulls modelparameters at round j .
Revised: This is shown in Figure 1. Let ˆ j = j mod K i , then the subset is given by: [Equation] The selected nodes start at ˆ j and increase until ⌊ β n K i ⌋ nodes are chosen.

Original: Let θ [ i,k ] be the k th trainable weight of the  i th layer of global model andθ m, [ i,k ] ; m ∈ M be the k th trainable weight of the i th layer of  the client m .
Revised: When the node index reaches K i , it returns to index 0 again. For instance in Figure 1, for the small client model, the extracted sub-model consists of { c, d, e } and is rolled in consecutive rounds to { d, e, a } and { e, a, b } How to aggregate updates from heterogeneous clients to update the global model?

Original: The model parametersare aggregated as follows:
Revised: In FedRolex , we compute the average of the updates separately for each variable based on how many participating clients in the federation updated that variable based on their data. Let M ⊂ N be the set of selected clients from the client pool from which the server pulls model parameters at round j . We use β n to denote the proportion of kernels extracted from the global parameter set θ . β n is chosen according to the client’s capacity C n .

Revised: Let θ [ i,k ] be the k th trainable weight of the i th layer of the global model and θ m, [ i,k ] be the k th trainable weight of the i th layer of the client m .

Revised: Denote M k ⊂ M as the set of clients updating the k th trainable weight, and the model parameters are aggregated as follows:

Original: [Equation]
Revised: [Equation]

Original: Here, p m is an empirical weight given to a client. This can be based ona variety of factors like theclient model capability, the number of data points a client has, etc.
Revised: which is supported by Federated Select [17]. The empirical weight p m is assigned based on factors like the client model capacity, the number of data points a client has, etc.

Original: However, there is no theoreticalfoundation in the literature to quantify the relative importance of a client update nor are thereempirical studies that accurately study this.
Revised: However, there is no theoretical foundation in the literature to quantify the relative importance of a client update nor are there empirical studies that accurately study this.

Original: We consider this tangent to our goals in this study and so,throughout the paper, unless otherwise stated, the weight of all clients is assumed to be the same, i.e,p m = 1  .
Revised: We consider this tangent to our goals in this study and so, throughout the paper, unless otherwise stated, the weight of all clients is assumed to be the same, i.e, p m = 1 /N .

Revised: As an example, we can see in Figure 1, for round 1 , the updates for weights a and b are obtained from the large capacity model and the update for weight e is taken from the small capacity model.

Original:  For different PT-based methods, the method in which the subset of selected parameters, S ischosen varies according to their design. Finally, the optimal sub-model is extracted from the globalmodel for each client according to client capability.
Revised: However, since weights c and d are common to both models, the update for them is computed by taking the average of the updates from both models. For different PT-based methods, how S is chosen varies according to their design which we show in §3.3.1.

Original: 3.3 Comparing Sub-model Extraction Techniques
Revised: The aggregation method is typically similar to these methods.

Original:  FedRolex proposes a new way to extract the sub-models.
Revised: The pseudocode of FedRolex is summarized in Algorithm 1.

Original: In general, sub-models are created bypulling a subset of weights from the global model.
Revised: 3.3 Comparison with Static and Random Sub-model Extraction Schemes

Original: This subset can be dynamic (Federated Dropout, FedRolex ) in each round or remain static for all rounds (HeteroFL).
Revised: Existing sub-model extraction schemes can be grouped as static (HeteroFL, FjORD) and random (Federated Dropout).

Original: In this section, we compare FedRolexwith two existing methods: static sub-model extraction (HeteroFL) and random sub-modelextraction (Federated Dropout).
Revised: We describe the differences between FedRolex and them.

Original: For simplicity, we focus on the i th layer at the n th client during thej th communication round only.
Revised: For simplicity, we focus on the i th layer at the n th client during the j th communication round.

Original: Let K i  denote the number of nodes in the i th layer of the globalmodel and β n denote the capacity of the client model.
Revised: We use K i to denote the number of nodes in the i th layer of the global model and β n denote the proportion of kernels extracted.

Revised: The pseudocodes of both HeteroFL and Federated Dropout are included in Appendix A.7.

Original:  Static Sub-Model Extraction Scheme
Revised: 3.3.1 Comparison with Static Sub-Model Extraction Scheme

Original: In this scheme (used by HeteroFL), the global model remainsstatic and the extracted sub-models are always the same.
Revised: In this scheme (used by HeteroFL), the global model remains static, and the extracted sub-models are always the same.

Original: This is shown in the left column of Figure 1.
Revised: This is shown in the left column of Figure 2.

Original: Unused nodes do not see the full data distribution.
Revised: Unused nodes do not see the full data distribution.

Original: The subset of node indices is given by:
Revised: The subset of node indices is given by:

Original: [Equation]
Revised: [Equation]

Original: K
Revised: i ⌋ − 1 } .

Original: [Equation]
Revised: [Figure]

Original: [Equation]
Revised: Limitations of Static Sub-model Extraction: HeteroFL and FjORD extract the model parameters θ n from designated locations in the global model: the client model size is based on the client capacity, and the nodes are selected starting from the first index. The extracted sub-models are then trained on selected devices from the client pool, and the updates for each parameter are aggregated and averaged across participating clients for that parameter. The nodes in the beginning indices of each layer will therefore see updates that were averaged across more clients than end indices. In contrast, the nodes at the end indices will see updates averaged across much fewer clients, and therefore the updates to these nodes will have more uncertainty.

Original:  Random Sub-Model Extraction Scheme.
Revised: 3.3.2 Comparison with Random Sub-Model Extraction Scheme

Original: In this scheme (used by Federated Dropout), the submodels are extracted by sampling nodes uniformly and randomly.
Revised: In this scheme (used by Federated Dropout), the sub-models are extracted by sampling nodes uniformly and randomly.

Original: This is shown in the middle columnof Figure 1.
Revised: This is shown in the right column of Figure 2.

Original: The subset is given by:
Revised: The subset is given by:

Original: [Equation] [Figure]
Revised: [Equation] [Equation]

Original: where ⌊ β n K i ⌋ nodes are randomly chosen from the global model.
Revised: Here ⌊ β n K i ⌋ nodes are randomly chosen from the global model.

Original: Unlike  static sub-model extraction,since nodes are shuffled, all nodes end up seeing the entire data given enough communication rounds.
Revised: Unlike the static sub-model extraction, all nodes see the entire data given a sufficiently large number of communication rounds.

Original: However, this may require a large number of communication rounds.

Original: Rolling Sub-Model Extraction. In this scheme (used by FedRolex ), the global model is rolledbefore the subset is extracted. This is shown in the right column of Figure 1. Let ˆ j = j mod K i , thenthe subset is given by:
Revised: Limitations of Random Sub-model Extraction: Dropout-based algorithms like Federated

Original: [Equation]

Original: That is, the selected nodes start at ˆ j and increase until ⌊ β n K i ⌋ nodes are chosen.

Original: When the nodeindex reaches K i , it goes back to one again.

Original: Why Rolling Sub-model Extraction is better?

Original: PT-based algorithms HeteroFL and FjORD extractthe model parameters θ n from specific locations in the global model.

Original: Here the client model size isbased on the client capacity, and the nodes are selected starting from the first index.

Original: The extractedsub-models are then trained on selected devices from the client pool, and the updates for eachparameter are aggregated and averaged across participating clients for that parameter.

Original: The nodes inthe beginning indices of each layer will therefore see updates that were averaged across more clientsthan end indices.

Original: In comparison, the nodes at the end indices will see updates averaged across muchfewer clients, and therefore the updates to these nodes will have more uncertainty.

Original: Dropout-basedalgorithms like Federated Dropout Caldas et al.

Original: (2018) propose random dropout to create smallersub-models.
Revised: Dropout [13] use random dropout to create smaller sub-models.

Original: Though the expected value of the frequency for updating each index is the same for allthe indices in Federated Dropout, their exact frequencies are not the same in each experiment becauseof the randomness.
Revised: Though the expected value of the frequency for updating each index is the same for all the indices, their exact frequencies are not the same in each experiment because of the randomness.

Original: This could potentially cause client drift.
Revised: This could potentially cause client drift.

Original: In fact, given I indices, if one index ischosen at each round, theprobability to update all indices once in I rounds is (cid:81) Ii =1 iI . Furthermore,the expected number of rounds to go through all indices at least once is  close to I ln
Revised: Given I components, if one component is chosen at each round, the expected number of rounds to go through all components at least once is (cid:81) Ii =1 iI , which is close to I ln

Original: ( I ) , see Appendix A.1.
Revised: ( I ) ; see Appendix A.1 for more details and the results on updating all components at least more than one time.

Original: It shows that random dropout can not balance the update frequencies of the indices, and it takesa large number of rounds to update all the indices.
Revised: It shows that random dropout can not balance the update frequencies of the components, and it takes too many rounds to update all the components.

Original:  Experimentsand Analysis
Revised:  Experiments

Original: 4.1 Experimental Setup

Original: Datasets and Platform.
Revised: Datasets and Models.

Original: We validate the performance of FedRolex on three datasets – CIFAR-10,
Revised: We evaluate the performance of FedRolex under two regimes.

Revised: Under small-model small-dataset regime, we train pre-activated ResNet18 (PreResNet18) models [26] on

Original: CIFAR-100, and StackOverflow.
Revised: CIFAR-10 and CIFAR-100 [27].

Original: We apply bounding box crop (Zoph et al., 2020) to augment theimages.
Revised: We replace the batch Normalization in PreResNet18 with static batch normalization [14, 28] and add a scalar module after each convolution layer [14]. Under largemodel large-dataset regime, we use Stack Overflow [29] and followed [2] to train a modified 3-layer Transformer [8] with a vocabulary of 10 , 000 words, where the dimension of token embeddings is 128 , and the hidden dimension of the feed-forward network (FFN) block is 2048 .

Original: To simulate the non-iid nature of the data, we use balanced non-iid data distribution following
Revised: We use ReLU activation and use 8 heads for the multi-head attention where each head is based on 12-dimensional (query, key, value) vectors. The statistics of the datasets are listed in Table 2.

Original: HeteroFL.
Revised: [Table]

Original: We skew the data distribution by restricting the clients to having a fixed number of uniquelabels L .
Revised: Data Heterogeneity. For CIFAR-10 and CIFAR-100, we followed HeteroFL [14] to model non-IID distributions by restricting each client to have L labels.

Original: We define L = 2 as high data heterogeneity(i.e., high non-iid) and L = 5 as low dataheterogeneity (i.e., low non-iid) for CIFAR-10.
Revised: In our evaluation, we consider two levels of data heterogeneity. For CIFAR-10, we define L = 2 as high data heterogeneity and L = 5 as low data heterogeneity.

Original: Similarly, we use L = 20 as high non-iid and L = 50as low non-iid for CIFAR-100.
Revised: For CIFAR-100, we use L = 20 as high data heterogeneity and L = 50 as low data heterogeneity.

Original: These  roughly correspond to Dirichlet distribution Dir K ( α ) with αset to 0 .
Revised: These two levels roughly correspond to Dirichlet distribution Dir K ( α ) with α equal to 0 .

Original: 1 and 0 . 5 , respectively.
Revised: 1 and 0 . 5 , respectively.

Original: The data is partitioned into 100 non-iid partitions and assigned to theclient pool.

Original: In each communication round, 10% of the clients are randomly selected from the clientpool of 100 .

Original: For StackOverflow, the data is partitioned over user ids. The total number of clients,in this case, was 342 , 477 and in each round, 200 clients were sampled. We use a vocabulary size
Revised: For Stack Overflow, the dataset is partitioned over user IDs, making the dataset naturally non-IID distributed.

Original: Algorithm 1 Initialization ; θ (0) , N Input : D n β n ∀ n ∈ N , Output : θ J 2 Server Executes 3 for j ← 0 to J do
Revised: Model Heterogeneity.

Original: Broadcast θ ( j ) n, [ i, S ( j ) i,k ] ∀ i, S from Equation 6 , 1 ≤ k ≤ K i and n ∈ N Sample subset M from N for each client m ∈ M do clientStep( θ ( j ) n , D n ) end Aggregate θ ( j +1)
Revised: Without loss of generality, in our evaluation, we consider five different client model capacities β

Original: [ i,k ] =  1 (cid:80) m ∈M 1 (cid:80) m ∈M θ m, [ i,k ] 11 end 12 Subroutine clientStep( θ ( j ) n , D n ) m n ←− len ( D n ) for k ← 0 to m n do θ n ←− θ n − η ∇ l ( θ n ; d n,k ) end return θ n of 10 , 000 words.
Revised: = { 1 , 1 / 2 , 1 / 4 , 1 / 8 , 1 / 16 } where for instance, 1 / 2 means the client model capacity is half of the largest client model capacity (full model).

Original: We run our experiments on 8 Nvidia A 6000 GPUs using the PyTorch library formachine learning and  the Ray library for simulating the distributed system. Models. We use pre-activated ResNet18 (PreResNet18) models (He et al., 2016) for CIFAR-10and CIFAR-100. We replace the batch Normalization in PreResNet18 with static batch normalization (Diao et al., 2020;
Revised: To generate these client models, for ResNet18, we vary the number of kernels in convolution layers and keep the nodes in the output layers the same.

Original: Andreux et al., 2020) and add a scalar module after each convolutionlayer (Diao et al., 2020).

Original: For StackOverflow experiments, we trained a modified 3-layer Transformermodel, where the dimension of the token embeddings was 128 , and the hidden dimension of thefeed-forward network(FFN) block is 2048 .

Original: We use 8 heads for the multi-head attention, where eachhead is based on 12-dimensional (query, key, value) vectors.

Original: We use ReLU activation and set thedropout rate to 0.1.

Original: To simplify our experiments and ablation studies, the model sizes are restricted toa set of 5 , i.e, β n would take any value in { 1 , 1 / 2 , 1 / 4 , 1 / 8 , 1 / 16 } .

Original: Note that we only vary the numberof kernels in convolution layers and keep the nodes in the output layers the same in both our methodand the PT-based baselines to keep the comparison fair.

Original: Similarly, for the Transformer model, weonly vary the number of nodes in the hidden layer of the attention heads.
Revised: For Transformer, we vary the number of nodes in the hidden layer of the attention heads.

Original: We provide pseudocodefor both HeteroFL and our Federated Dropout variant in Appendix A.7.

Original: Unless otherwise stated, theclient capacity distribution is assumed to be uniform over the entire client pool for all methods.

Original: [-Welist the exact hyperparameters and architecture of the models in Appendix A.
-]
Revised: [-Welist the exact hyperparameters and architecture of the models in Appendix A.
-]

Original: Baselines.
Revised: Baselines.

Original: We compare FedRolex against  state-of-the-art PT-based  methods HeteroFL and Federated Dropout  and  state-of-the-art KD-based  methods FedDF, DS-FL, and Fed-ET obtained from Lin et al.(2020).
Revised: We compare FedRolex against both state-of-the-art PT-based model-heterogeneous FL methods including Federated Dropout [13] and HeteroFL [14] 2 as well as state-of-the-art KD-based model-heterogeneous FL methods including FedDF [10], DS-FL [11] and Fed-ET [12] 3 .

Original: For a fair comparison, all thebaselines from the PT-based methods are trained on the samelearning rate, the same  number of communication rounds, andthe same multi-step learning rate decayschedule.
Revised: To ensure a fair comparison, all the PT-based baselines are trained using the same learning rate, number of communication rounds, and multi-step learning rate decay schedule.

Original: The exact schedule for each dataset and experiment is given in Appendix A.6.
Revised: The details of the schedule for each dataset and experiment are described in Appendix A.6.

Revised: Configurations and Platform.

Revised: For CIFAR-10 and CIFAR-100, we apply bounding box crop [30] to augment the images.

Revised: In each communication round, 10% of the clients are randomly selected from a pool of 100 clients.

Original:  Evaluation Metrics.
Revised: For Stack Overflow, we followed [2] to use a 10% dropout rate to prevent overfitting, and 200 clients are randomly selected from a pool of 342 , 477 clients in each communication round. The details of the hyper-parameters for model training are included in Appendix A.6. We implemented FedRolex and PT-based baselines using PyTorch [31] and Ray [32], and conducted our experiments on 8 NVIDIA A 6000 GPUs. Evaluation Metrics.

Original: We uselocal and global  model accuracy as our evaluation metrics.
Revised: We use global and local model accuracy as our evaluation metrics.

Original: Specifically,global model accuracy is defined as the mean accuracy on the test set totally held out during training;local model accuracy is defined as themean accuracy of the server model on each of the clientdatasets.
Revised: Specifically, global model accuracy is defined as the server model accuracy on the test set; and local model accuracy is defined as the accuracy of the server model on each of the client’s local datasets. For

Original: For vision datasets CIFAR-10 and CIFAR-100, we measure classification accuracy andfor the StackOverflow dataset, we measure next word prediction accuracy.
Revised: CIFAR-10 and CIFAR-100, we report the classification accuracy. For Stack Overflow, we report the next word prediction accuracy which includes both out-of-vocabulary (OOV) and end-of-sentence (EOS) tokens.

Original: Since there is somerandomness due to the non-iid partitioning, we run our experiments on 5 different seeds  and list themean andstandard deviation of our metrics. Note that the results of KD-based methods were obtainedfrom Cho et al. (2022) using 3 different seeds.
Revised: We run our experiments using five different seeds for CIFAR-10 and CIFAR-100 and using three different seeds for Stack Overflow.

Original: 4.2 Performance Comparison with State-of-the-Art Model-Heterogeneous FL Methods
Revised: 4.1 Performance Comparison with State-of-the-Art Model-Heterogeneous FL Methods

Original: shows the global model accuracy comparison between FedRolex  and the baselines underboth high and low data heterogeneity settings.
Revised: First, we compare the performance of FedRolex with state-of-the-art PT and KD-based modelheterogeneous FL methods.

Original: We also include the global model accuracy of twomodel-homogeneous cases whereall clients have the highest capacity ( β = 1 . 0 ) and lowest capacity
Revised: For a fair comparison, we followed the experimental settings used in prior arts where the distributions of client model capacities are uniform and the global server model is the same as the largest client model.

Original: [Table] ( β = 0 . representing the upper and lower-bound performance, respectively.
Revised: Evaluation Results: Table 3 summarizes our results.

Original: We have threekey observations.
Revised: We have two observations.

Original:  In comparison with state-of-the-art PT-based methods,  FedRolex consistentlyoutperforms HeteroFL and Federated Dropout in bothhigh and low  data heterogeneity scenarios.
Revised: (1) In comparison with state-of-the-art PT-based methods, under the small-model small-dataset regime, FedRolex consistently outperforms HeteroFL and Federated Dropout under both low and more challenging high data heterogeneity scenarios.

Original: Inparticular, under high data heterogeneity, Federated Dropout that uses random dropout has drasticallyworse performance than FedRolex and HeteroFL which both extract sub-models in a deterministicmanner.
Revised: In particular, under high data heterogeneity, Federated Dropout which extracts sub-model randomly has worse performance than FedRolex and HeteroFL which both extract sub-models in a deterministic manner. Under large-model large-dataset regime, FedRolex also outperforms both HeteroFL and Federated Dropout. These results together demonstrate the

Original:  In comparison with state-of-the-art KD-based methods, FedRolex only performs worsethan Fed-ET and FedDF on CIFAR-10 under high data heterogeneity, but outperforms all the state-ofthe-art in more challenging CIFAR-100 dataset that has many more classes than CIFAR-10 underboth high and low  data heterogeneity scenarios.
Revised: [Table] superiority of FedRolex under both regimes. In comparison with state-of-the-art KD-based methods, FedRolex only performs worse than Fed-ET and FedDF on CIFAR-10 under high data heterogeneity, but outperforms all the KD-based methods on the more challenging CIFAR-100 which has a larger number classes than CIFAR-10 under both low and high data heterogeneity scenarios.

Original: It is important to emphasize that KD-based methodsleverage public data to boost their model accuracy while FedRolex does not.
Revised: It is important to note that KD-based methods leverage public data to boost their model accuracy while FedRolex does not.

Original: In comparison withmodel-homogeneous cases, compared to other PT-based methods, FedRolex considerably reducesthe gap  between model-heterogeneous and model-homogeneous settings, particularly on CIFAR-10under low data heterogeneity.
Revised: We also compare the global model accuracy of FedRolex with two model-homogeneous cases where all the clients have the largest capacity model ( β = { 1 } ) and the smallest capacity model ( β = { 1 / 16 } ), representing the upper and lower-bound performance, respectively. Evaluation Results: As listed in Table 3, compared with other PT-based methods, FedRolex reduces the gap in global model accuracy between model-heterogeneous and upper-bound modelhomogeneous settings. In particular, FedRolex is on par with the upper-bound model-homogeneous case for Stack Overflow, whereas both HeteroFL and Federated Dropout perform even worse than the model homogeneous case using the smallest model.

Revised: This result indicates that with FedRolex , we will not be constrained to only using high-end devices to achieve competitive global model accuracy.

Original: Note that Fed-ET achieves a higher global model accuracy than themodel-homogeneous upper bound on CIFAR-10 under high data heterogeneity, which showcases theadvantage of having access to public data.
Revised: Note that Fed-ET achieves a higher global model accuracy than the model-homogeneous upper bound on CIFAR-10 under high data heterogeneity, which showcases the advantage of using public data.

Original: 4.3 Impact of  Model Capacity Distribution
Revised: 4.3 Impact of Client Model Heterogeneity Distribution

Original: In our previous experiment, in line with existing model-heterogeneous FL studies, we have performedour experiment where the distributions of model capacities  are uniform for a fair comparison. Inpractical scenarios, however, the distributions of model capacities will likely be skewed.
Revised: In our previous experiments, the distributions of model capacities across client devices are set to be uniform.

Original: Therefore,we conduct an experiment to understandhow global model accuracy changes when the proportion ofhigher capacity clients varies.
Revised: In this experiment, we aim to understand the impact of the client model heterogeneity distribution.

Original: To do so, we set the proportion of large to small models to ρ and theset of unique client  capacitiesis β = { 1 , 1 / 16 }  where ρ = 1 represents theedge case in whichthemodels of all  clients have the architecture of the full ResNet18 model, and ρ = 0 is the other edgecase in which all  clients have the smallest model.
Revised: To do so, without loss of generality, we use two client model capacities β = { 1 , 1 / 16 } and vary the distribution ratio between the two (denoted as ρ ) where ρ = 1 represents the case in which all the clients have the largest capacity model ( β = { 1 } ) and ρ = 0 represents the case in which all the clients have the smallest capacity model ( β = { 1 / 16 } ).

Original:  Figure 2 shows how  global model accuracy changes when ρ varies from 0 to 1 for CIFAR-10 and CIFAR-100under both high and low data heterogeneity.
Revised: Evaluation Results: Figure 3 shows how the global model accuracy changes when ρ varies fromto 1 for CIFAR-10, CIFAR-100 and Stack Overflow.

Original: We have threekey observations.
Revised: We have three observations.

Original: (1) For 
Revised: (1) For CIFAR-

Original: CIFAR-10, there is a large gap in global model accuracy between high and low data heterogeneityfor a wide range of ρ (from 0.1 to 1).
Revised: (Figure 3(i)), there is a large gap in global model accuracy between high and low data heterogeneity for a wide range of ρ (from 0.1 to 1).

Original: This is because CIFAR-10 is a relatively simple task andhence the global model accuracy is bottlenecked by the level of data heterogeneity instead of modelcapacity.
Revised: This is because CIFAR-10 is a relatively simple task and hence the global model accuracy is bottlenecked by the level of data heterogeneity instead of model capacity.

Original: This result indicates that having more high-capacity models in the cohort has limitedcontribution to global model accuracy.
Revised: This result indicates that having more high-capacity models in the cohort has only limited contribution to global model accuracy.

Original: For the more challenging CIFAR-100, the gap in globalmodel accuracy is much lower between high and low data heterogeneity.
Revised: For the more challenging CIFAR-100 (Figure 3(ii)), the gap in global model accuracy is much lower between high and low data heterogeneity.

Original: In contrast to CIFAR-10, theglobal model accuracy is bottlenecked by the highest capacity of the models rather than the level ofdata heterogeneity.
Revised: In contrast to CIFAR-10, the global model accuracy is bottlenecked by the highest capacity of the models rather than the level of data heterogeneity.

Original: For both CIFAR-10 and CIFAR-100, we observe that having a small fractionof high-capacity modelsis enough to significantly boost the global model accuracy. 4.4 Scalability One limitation of HeteroFL and FjORD is that the server model size is restricted to the highestcapacity client model. In contrast, FedRolex is flexible totrain a global model that can be much
Revised: For both regimes (Figure 3(i)(ii) vs. Figure 3(iii)), we observe that having a small fraction of large-capacity models significantly boosts the global model accuracy, but keeping increasing the ratio of large-capacity models has limited contribution to global model accuracy.

Original: [Figure]
Revised: [Figure]

Original: larger than the largest client model.
Revised: 4.4 Performance on Training Larger Server Model

Original: This is more akin to real-world settings where the  server has amuch higher capacity.
Revised: Similar to Federated Dropout, one advantage of FedRolex over static sub-model extraction methods (HeteroFL and FjORD) is that FedRolex is able to train a global server model that is larger than the largest client model. In this experiment, we aim to evaluate the performance of FedRolex on training larger server models.

Original: We consider the case where the  global  modelsize is γ  times the size of  the highest capacity clientmodel.
Revised: To do so, we consider the case where the size of the global server model is γ = { 2 , 4 , 8 , 16 } times the size of client models. For simplicity, all client models have the same size.

Original:  Figure 3 shows the global model accuracy of γ = { 2 , 4 , 8 , 16 } for CIFAR-10 andCIFAR-100under highand low data heterogeneity. As shown, for both CIFAR-10 and CIFAR-100, the globalmodel accuracy drastically drops when γ increases from 2 to 4 but stay relatively stable when γincreases from 4 to 8 and from 8 to 16 .
Revised: Evaluation Results: Figure 4(i) and Figure 4(ii) compare FedRolex with Federated Dropout in terms of global model accuracy when γ for CIFAR-10 and CIFAR-100, respectively. As shown, although the global model accuracy drops for both FedRolex and Federated Dropout when γ increases, especially from 1 to 4 , FedRolex consistently achieves higher global model accuracy than Federated Dropout across γ = { 2 , 4 , 8 , 16 } under both low and high data heterogeneity. For Stack Overflow (Figure 4(iii)), the global model accuracy has a much smaller drop when γ increases.

Original: This result indicates that there is a minimum client modelcapacity below which the  server model cannot be trained effectively.
Revised: This demonstrates the superiority of using large models on large-scale datasets for training larger server models.

Original: [Figure]
Revised: [Figure]

Original: 4.5 Enhance Inclusiveness of FL in Real-world Distribution
Revised: 4.5 Enhance Inclusiveness of FL in Real-world Distribution

Original: A primary goal of our method is to enhance the inclusiveness of FL.
Revised: A primary vision of FedRolex is to enhance the inclusiveness of FL.

Original: If the highest model capacity isused for deployment, a lot of data from the low-end devices will be unused.

Original: The exclusion of low-enddevices limits the usefulness of the global model for these low-end devices.

Original: For example, low-enddevices may have more grainy, low-resolution photos which will not be used in training, and so, theserver model cannot make accurate predictions on such images and hence limiting its usefulnessfor these devices.

Original: If we consider homogeneous FL where we want to include all the client data intraining, the smallest model would need to be used.

Original: The accuracy of the model is then limited by thelowest capacity client.

Original: In this experiment, we show that FedRolex is a valid solution that enhances the inclusiveness offederated learning.

Original: To illustrate our point, we conducted an experiment that uses real-world householdincome distribution to emulate real-world device distribution.
Revised: To demonstrate this, in this experiment, we use real-world household income distribution to emulate real-world device distribution.

Original: Specifically, we retrieve householdincome distribution information from Bureau (2021).
Revised: Specifically, we retrieve household income distribution information from Bureau [33].

Original: We map β n = 1 / 16 with the income group withearning less than $ 75000 and assign proportions of remaining groups in $ 25000 increments with [Figure]
Revised: We map β n = 1 / 16 with the income group with earning less than $ 75 , 000 and assign proportions of remaining groups in $ 25 , 000 increments with increasing values of β n .

Original: increasing values of β n in the statistic. Detailed mapping of this distribution to the correspondingincome distribution is given in Appendix A.
Revised: Detailed mapping of this distribution to the corresponding income distribution is provided in Figure 7 in Appendix A.6. Evaluation Results:

Original: [Table]
Revised: Table 4 shows both the global and local model accuracies of FedRolex for CIFAR-10 and CIFAR-100 as well as the global model accuracy on Stack Overflow under the emulated real-world device distribution. Again, we compare with two model-homogeneous cases where all clients have the smallest and largest model capacities, representing lower and upper-bound

Original: [Table]
Revised: [Table]

Original: Table 3 shows the global and local model accuracy of FedRolex on this emulated setting comparedto two model-homogeneous cases where all clients have the smallest and largest model capacities,representing lower and upper bound on accuracy respectively.
Revised: [Figure] accuracy, respectively.

Original: We make twokey observations.
Revised: We make two observations.

Original:  FedRolex consistentlyand significantly outperforms the lower bound model-homogeneous casein terms of local model accuracy for CIFAR-10 and CIFAR-100 under both high and low dataheterogeneity.
Revised: (1) Looking at the global model accuracy, FedRolex consistently outperforms the lower-bound model-homogeneous case across CIFAR-10, CIFAR-100, and Stack Overflow.

Original: This result indicates that FedRolex can effectively boost thelocal model accuracy oflow-end devices, which would otherwise not benefit from FL.
Revised: This result indicates that FedRolex enhances the inclusiveness of FL and improves the accuracy of the global model, which would otherwise not be able to achieve.

Original: (2) FedRolex achieves a local modelaccuracy that is close to the upper bound model-homogeneous case, especially for CIFAR-10underboth high and  low  data heterogeneity.
Revised: Looking at the local model accuracy, FedRolex significantly outperforms the lower-bound modelhomogeneous case on CIFAR-10 and CIFAR-100 under both low and high data heterogeneity.

Original: This result indicates that FedRolex can significantly reducethe gap with the best-performing scenario without being constrained to only using high-end devices.
Revised: This result indicates that FedRolex effectively boosts the performance of low-end devices, which would otherwise not benefit from FL.

Original: Finally, Figure 4 provides a detailed illustration of howthe local model accuracy  shiftsfor CIFAR-10and CIFAR-100 under low and high data heterogeneity when FedRolex is used compared to thesmallest homogeneous case with the same client outreach. As shown, the local model accuracy ofindividual clients is significantly boosted by FedRolex .
Revised: A detailed illustration of how local model accuracy distribution of individual clients shifts when FedRolex is used compared to the smallest model-homogeneous case with the same client outreach is shown in Figure 5.

Original:  Conclusionand Future Directions
Revised:  Conclusion

Original: In this work, we present a partial training (PT)-based model-heterogeneous method named FedRolexthat tackles the fundamental issue of existing PT-based methods.
Revised: We presented FedRolex , a partial training (PT)-based model-heterogeneous FL approach that is able to train a global server model larger than the largest client model.

Original: At the core of FedRolex is asimpleyet effective rolling sub-model extraction scheme that enables parameters of the global  model to beevenly trained to minimize client drift.
Revised: FedRolex proposed a rolling sub-model extraction scheme that enables parameters of the global server model to be evenly trained to mitigate client drift induced by model heterogeneity.

Original:  Our experimental results show that FedRolex consistentlyoutperforms state-of-the-art PT-based methods in both high and low data heterogeneity scenarios. Wedemonstrate its effectiveness in real-world scenarios by showing its performance on a more practicalmodel capacity distribution and show FedRolex contributes to making FL more inclusive. In themeantime, we note that although we provide a statistical analysis of FedRolex , we do not derive arigorous convergence analysis.
Revised: We provided a theoretical statistical analysis on its advantage over Federated Dropout. Our experimental results show that FedRolex consistently outperforms state-of-the-art PT-based methods across models and datasets at both small and large scales.

Original: This is not a trivial task and therefore left for future work.

Original: In addition,our goal in this paper is to train a large global model using a federation of clients with heterogeneouson-device resources. Determining what models to deploy onto each heterogeneous client after theglobal model has finished training during the federated learning process is a separate task. This canbe achieved by a few techniques such as compressing/quantizing/KD the global model into smallermodels to fit the resources of the target client. However, this is not the focus of this work and is aseparate study in itself.
Revised: Moreover, we demonstrated its performance on an emulated real-world device distribution and show FedRolex contributes to making FL more inclusive.

Original: We will pursue it in our future work as well. Checklist For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
Revised: Limitations and Future works.

Original: contributions and scope? [Yes] Our work’ claims are verified through experiment inSection 4.2 and statistical study in Appendix A.1(b) Did you describe the limitations of your work? [Yes] The limitations are mentioned inSection 5(c)
Revised: In this work, we provided a statistical analysis of FedRolex .

Original: Did you discuss any potential negative societal impacts of your work? [No](d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
Revised: A full convergence analysis of FedRolex is not trivial and left for future work.

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A] We do notshow any rigorous theoretical results nor claim as such

Original: (b) Did you include complete proofs of all theoretical results?

Original: [N/A] We do not show any

Original: rigorous theoretical results nor claim as such If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
Revised: In addition, the goal of this work is to train a global server model using a federation of heterogeneous client models.

Original: [Yes] Code is madeavailable to reviewers.

Original: Data is open source. Instructions to run are available in code (b) Did you specify all the training details (e g., data splits, hyperparameters, how theywere chosen)?
Revised: Determining what models to deploy onto each client after the global server model is trained is a separate task, especially when the global server model is large.

Original: [Yes] All major training details are mentioned in Section 4 and fine

Original: details provided in Appendix A.6(c) Did you report error bars (e g., with respect to the random seed after running experi- ments multiple times)?
Revised: We will pursue it as our future work.

Original:  [Yes] All experiments were performed with 5 different seeds and the mean and standard deviations were reported when applicable. (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes]
Revised:  Acknowledgement

Original: If you are using existing assets (e

Original: g., code, data, models) or curating/releasing new assets..

Original: (a) If your work uses existing assets, did you cite the creators? [Yes] (b) Did you mention the license of the assets? [N/A]
Revised: We would like to thank the anonymous reviewers for their helpful comments.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N/A]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
Revised: We also want to thank Jiahui Yu, Zheng Xu, and Chaoyang He for their constructive discussions.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
Revised: This work was partially supported by NSF Awards PFI:BIC-1632051, CNS-1814551, DMS-2012439, and a Google

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
Revised: Computing Platform (GCP) credit grant.

Original: [N/A]