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

Original: Abstract
Revised: Abstract

Original: Most cross-device federated learning (FL) studies focus on the model-homogeneous setting where the global server model and local client models are 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, 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.
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 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.
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: At its core,
Revised: At its core,

Original: 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.
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: 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.
Revised: 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.

Original: In addition, we provide theoretical statistical analysis on its advantage over Federated Dropout. 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.
Revised: In addition, we provide theoretical statistical analysis on its advantage over Federated Dropout, and 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.

Original: Our code is available at: https://github.com/MSUMLSys-Lab/FedRolex.
Revised: Our code is available at: https://github.com/MSUMLSys-Lab/FedRolex.

Original:  Introduction
Revised:  Introduction

Original: 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].
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: In this work, we focus on crossdevice FL where clients are usually resource-constrained edge devices.
Revised: In this work, we focus on crossdevice FL where clients are usually resource-constrained edge devices.

Original: 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 .
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: 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.
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.

Original: Having the same model on all the devices would, unfortunately, exclude clients with low-end devices who would otherwise make unique contributions tothe model training from their  local data; (2) state-of-the-art machine learning has moved towards large models [7] such as Transformer [8].
Revised: Having the same model on all the devices would, unfortunately, exclude clients with low-end devices who would otherwise make unique contributions to model training from their own local data; (2) state-of-the-art machine learning has moved towards large models [7] such as Transformer [8].

Original: 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.
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.

Original: 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.
Revised: 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: One primary challenge in model-heterogeneous FL is
Revised: One primary challenge in model-heterogeneous FL is

Original: [Table]
Revised: [Table]

Original: the aggregation of heterogeneous client models.
Revised: the aggregation of heterogeneous client models.

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

Original: [Table]
Revised: [Table]

Original: in which the client models serve as teachers and the server ensembles the knowledge distilled from the individual client models.
Revised: in which the client models serve as teachers and the server ensembles the knowledge distilled from the individual client models.

Original: However, KD-based approaches in general require public data on the server to achieve competitive model accuracy
Revised: However, KD-based approaches in general require public data on the server to achieve competitive model accuracy

Original: [Table]
Revised: [Table]

Original: whereas the desired public data may not be always available in practice.
Revised: whereas the desired public data may not be always available in practice.

Original: Moreover, since KD-based approaches need the individual client models (whole models or prediction layers) or their outputs to be sent to the server
Revised: Moreover, since KD-based approaches need the individual client models (whole models or prediction layers) or their outputs to be sent to the server

Original: [Table]
Revised: [Table]

Original: they are incompatible with secure aggregation protocols [16]
Revised: they are incompatible with secure aggregation protocols [16]

Original: [Table]
Revised: [Table]

Original: which limits their privacy guarantee.
Revised: which limits their privacy guarantee.

Original: 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
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

Original: Dropout [13]) and static sub-model extraction (HeteroFL [14]
Revised: Dropout [13]) and static sub-model extraction (HeteroFL [14]

Original: [Table]
Revised: [Table]

Original: FjORD [15]) were proposed.
Revised: FjORD [15]) were proposed.

Original: In these approaches, each client trains a smaller sub-model extracted from the larger global server model
Revised: In these approaches, each client trains a smaller sub-model extracted from the larger global server model

Original: [Table]
Revised: [Table]

Original: and the server model is updated by aggregating those trained sub-models.
Revised: and the server model is updated by aggregating those trained sub-models.

Original: 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 .
Revised: 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: 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.
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: To tackle the fundamental issue of existing PT-based methods, in this work, we propose asimple yet effective PT-based model-heterogeneous FL approach named FedRolex .
Revised: In this work, we propose a PT-based model-heterogeneous FL approach named FedRolex to tackle the fundamental issue of existing methods.

Original: 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.
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: 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.
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.

Original: Since the window is rolling, sub-models from different parts of the global model are extracted in sequence in different rounds.
Revised: Since the window is rolling, sub-models from different parts of the global model are extracted in sequence in different rounds.

Original: As a result, all the parameters of the global server model are evenly trained over the local data of client devices.
Revised: As a result, all the parameters of the global server model are evenly trained over the local data of client devices.

Original: 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.
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.

Original: (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.
Revised: (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: It echos some concurrent efforts in developing FL primitives to support training large server models in cross-device setting, e.g. Federated Select [17].
Revised: It echoes some concurrent efforts in developing FL primitives to support training large server models in cross-device settings, e.g. Federated Select [17].

Original: (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.
Revised: (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.

Original: (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.
Revised: (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.

Original: This allows clients to contribute to federated training under resource constraints and reduces communication overheads.
Revised: This allows clients to contribute to federated training under resource constraints and reduces communication overheads (Appendix A.6).

Original: (5) Lastly, FedRolex is fully compatible with existing secure aggregation protocols that enhance the privacy properties of FL systems.
Revised: (5) Lastly, FedRolex is fully compatible with existing secure aggregation protocols that enhance the privacy properties of FL systems.

Original: 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
Revised: 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: (this combination echos recent efforts on pushing the frontier of cross-device FL towards training large server models on large-scale datasets [18–21]).
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]).

Original: 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).
Revised: 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: (2) FedRolex reduces the gap between model-heterogeneous and model-homogeneous FL, especially under large-model large-dataset regime (§4.2).
Revised: (2) FedRolex reduces the gap between model-heterogeneous and model-homogeneous FL, especially under large-model large-dataset regime (§4.2).

Original: (3) With FedRolex , under both regimes, having a small fraction of large-capacity models could significantly boost the global model accuracy (§4.3).
Revised: (3) With FedRolex , under both regimes, having a small fraction of large-capacity models could significantly boost the global model accuracy (§4.3).

Original: (4) FedRolex is able to train a global server model that is larger than the largest client model and outperforms
Revised: (4) FedRolex is able to train a global server model that is larger than the largest client model and outperforms

Original: Federated Dropout in terms of global model accuracy (§4.4).
Revised: Federated Dropout in terms of global model accuracy (§4.4).

Original: (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).
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: Knowledge Distillation (KD)-based Model-Heterogeneous FL.
Revised: Knowledge Distillation (KD)-based Model-Heterogeneous FL.

Original: One primary approach for modelheterogeneous FL in cross-device settings is based on knowledge distillation (KD) [22].
Revised: One primary approach for modelheterogeneous FL in cross-device settings is based on knowledge distillation (KD) [22].

Original: In particular,
Revised: In particular,

Original: FedDF [10] 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 an unlabeled public dataset are then used to train a student model at 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 [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.
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.

Original: FedGKT [9] proposed group knowledge transfer in which knowledge is transferred from a large model in the server to clients without public data.
Revised: FedGKT [9] proposed group knowledge transfer in which knowledge is transferred to a large model in the server from clients without public data.

Original: 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.
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: KD-based approaches, however, have several limitations: these often require public data to achieve competitive model accuracy.
Revised: KD-based approaches, however, have several limitations: they often require public data to achieve competitive model accuracy.

Original: 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].
Revised: 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: 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].
Revised: Furthermore, as KD-based methods 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: To address the limitations of KD-based approaches, partial training (PT) has emerged as another solution for model-heterogeneous FL.
Revised: To address the limitations of KD-based approaches, partial training (PT) has emerged as another solution for model-heterogeneous FL.

Original: 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.
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: random sub-model extraction and static sub-model extraction.

Revised: Specifically, inspired by the dropout technique commonly used in centralized training [24], Federated Dropout [13] proposed to randomly extract sub-models from the global model.

Revised: Though easy to be integrated into existing FL frameworks, as reported in [25], Federated

Revised: Dropout becomes less effective when the data heterogeneity is high and the client cohort is small due to its randomness in selecting sub-models.

Original:  HeteroFL [14] and FjORD [15] proposed staticsub-model extraction schemes where sub-models are always extracted from a designated part of the global server model.
Revised: In contrast, HeteroFL [14] and FjORD [15] proposed static extraction schemes where sub-models are always extracted from a designated part of the global server model.

Original: However, such a static extraction strategy has two key drawbacks.
Revised: However, such a static extraction strategy has two primary drawbacks.

Original: First, the global server model is restricted to the same size as the largest client model.
Revised: First, the global server model is restricted to the same size as the largest client model.

Original: 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.
Revised: 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: 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.
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: 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.
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.

Original: This wouldsignificantly degrade the performance of the global model, especially under high data heterogeneity.
Revised: This would degrade the performance of the global model, especially under high data heterogeneity.

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

Original: 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: In this work, we propose a rolling sub-model extraction scheme that tackles the issues of bothstatic and random  sub-model extraction methods.
Revised: In this work, we propose a rolling sub-model extraction scheme that tackles the issues of both random and static sub-model extraction methods.

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

Original: Let N denote N client devices with non-IID (non-identically and independently distributed) local data D
Revised: Let N denote N client devices with non-IID (non-identically and independently distributed) local data D

Original: = { D 1 , D 2 , ..., D N } .
Revised: = { D 1 , D 2 , ..., D N } .

Original: Model-homogeneous FL trains a global model of parameter θ by solving the following optimization problem:
Revised: Model-homogeneous FL trains a global model of parameter θ by solving the following optimization problem:

Original: [Equation]
Revised: [Equation]

Original: with
Revised: with

Original: [Equation]
Revised: [Equation]

Original: 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
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: 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
Revised: In comparison, in model-heterogeneous FL, clients train local models with heterogeneous capacities { β 1 , β 2 , ..., β N } , and the local objective function of the n th client becomes

Original: [Equation]
Revised: [Equation]

Revised: Here, β n represents the model’s capacity and we define it as the proportion of nodes extracted from each layer in θ for client n .

Revised: Furthermore, let the set of unique capacities be β .

Original: 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.
Revised: The size of θ n depends on the client’s model capacity β n , and the parameter θ n is obtained by selecting a sub-model from the global model θ , which can change from one iteration to another.

Original: If θ n changes, the objective function also changes.
Revised: If θ n changes, the objective function also changes.

Original: For simplicity, we use the same notation l for the loss function at all clients and iterations, though they differ between clients and iterations.
Revised: For simplicity, we use the same notation l for the loss function for all clients and iterations, though they differ between clients and iterations.

Original: The key to model-heterogeneous FL is selecting θ n from the global model θ given model capacity C n .
Revised: The key to model-heterogeneous FL is selecting θ n from the global model θ given model capacity β n .

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

Original: [Figure]
Revised: [Figure]

Original: As a partial training (PT)-based approach, at eachparticipating 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.
Revised: As a partial training (PT)-based approach, at each client, FedRolex trains only a sub-model extracted from the global server model and sends the corresponding submodel updates back to the server for update aggregation.

Original: To help understand how FedRolex works, for simplicity, Figure 1 illustrates three rounds of federated training of FedRolex on two  heterogeneous clients.
Revised: To help understand how FedRolex works, for simplicity, Figure 1 illustrates three rounds of federated training of FedRolex on two participating heterogeneous clients, where one trains a large-capacity sub-model (left) and the other trains a small-capacity one (right).

Original: 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.
Revised: At the high level, at each round, the server extracts sub -models of different capacities from the global model

Original:  The clients train the received sub-models on their private data and post their  updatesθ ( j ) m to the server.
Revised: and broadcasts them to the clients that have the corresponding capabilities. The clients train the received sub-models on their local data and transmit their heterogeneous submodel updates to the server.

Original: The serverthen aggregates  the variables from the updates and formulates thenew global modelθ ( j +1) for the next round.
Revised: Lastly, the server aggregates those updates, and the result of the aggregation is used to update the global model for the next round.

Original: Specifically, the design of FedRolex involvesthe following twoprimary design decisions.
Revised: The pseudo code of FedRolex is summarized in Algorithm 1. The key to the design of FedRolex involves two design choices.

Original: Equation (1) shows us that we must find θ that minimizes the objective concerning all private datasets.
Revised: In the following, we describe them in detail. (1) What sub-models to be extracted for each client across different rounds?

Original: In centralized settings, this is achieved by simply iterating over the entire data. In FL, however, we do not have this option.
Revised: At the server, FedRolex utilizes a rolling window to extract the sub-model from the global model.

Original: Instead, the parts of the global model are dynamically trained by iterating them over the entire client federation.
Revised: The rolling window advances in each round, and loops over all parts of the global model in sequence across different rounds. This process iterates such that the global model is evenly trained until convergence.

Original: Consider θ ( j ) n selected from  global model weights for client n at the j th communication round. 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 β .
Revised: Taking Figure 1 as an example: in round j , the large-capacity and small-capacity client model extracted from the global model is { a, b, c, d } and { c, d, e } , respectively. In round j + 1 , the rolling window advances one step 2 , the large-capacity and small-capacity client model becomes { b, c, d, e } and { d, e, a } , respectively. Similarly, in round j + 2 , the rolling window advances one step further, and the large-capacity and small-capacity client model becomes { c, d, e, a } and { e, a, b } , respectively.

Original:  Let  the total number of nodes inthe i th layerbe K i and the selected nodes are S ( j )  i . Then the parameters θ ( j ) n, [ i, · ] for thei 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 ] } .Intuitively, we can think of this as iterating over the model instead of the data.
Revised: Such a rolling sub-model extraction scheme can be formalized as follows. Let θ ( j ) n denote the parameters of the sub-model extracted from the global model for client n in round j , K i denote the total number of nodes in layer i of the global model, and S ( j ) n,i denote the node indices of layer i of the global model that belongs to the extracted sub-model for client n in round j .

Original: In FedRolex , the global model is rolled before the subset is extracted. This is shown in Figure 1. Let ˆ j= j mod K i , then the subset is given by:
Revised: Then the layer i of the sub-model extracted by the rolling sub-model extraction scheme for client n in round j is given by:

Original: [Equation]
Revised: [Equation]

Original: The selected nodes start at ˆ j and increase until ⌊ β n K i ⌋ nodes are chosen.
Revised: else

Original: 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 }
Revised: [Equation] 4) where ˆ j = j mod K i .

Original:  How to aggregateupdates from heterogeneous clients to update the global model?
Revised: (2) How to aggregate heterogeneous sub-model updates to update the global model?

Original: In FedRolex , we compute the average of the updatesseparately 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: FedRolex adopts a simple selective averaging scheme to aggregate heterogeneous sub-model updates to update the global model.

Original: 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 . Denote M k ⊂ M as the set of clients updating the k th trainable weight, and the model parameters are aggregated as follows: [Equation]
Revised: Specifically, it computes the average of the updates for each parameter of the global model separately based on how many clients in one round updated that parameter 3 .

Original: 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. However, there is notheoretical foundation in the literature to quantify the relative importance of a client update nor are there empirical studies that accurately study this. 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: The parameter remains unchanged if no clients updated that parameter.

Original: As an example, we can see inFigure 1, for round 1 , the updates forweights a and b are obtained from the large capacity model and the update forweight e istaken from the small capacity model.
Revised: Taking Figure 1 as an example: in round j , the updates for a and b are obtained from the large-capacity model and the update for e is from the small-capacity model only.

Original: However, sinceweights c and d are common to both models, the updatefor them is computed by taking the averageof the updates from both models.
Revised: In contrast, since c and d are part of both models, the update is computed by taking the average from both models.

Original: For different PT-based methods, how S is chosen varies according to their design which we show in §3.3.1.
Revised: Algorithm 1: FedRolex 1 Initialization ; θ (0) , N Input : D n β n ∀ n ∈ N , Output : θ J 2 Server Executes 3 for j ← 0 to J − 1 do Sample subset M from N

Original: The aggregation method is typically similar to these methods.
Revised: Broadcast θ ( j m,

Original: The pseudocode of FedRolex is summarized in Algorithm 1.
Revised: ( j m, to client m ∈ M ∀ i, S m,i from

Original: 3.3 Comparison with Static and Random Sub-model Extraction Schemes
Revised: (4) for each client m M do end

Revised: ( θ ( 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

Revised: 3.3 Comparison with Static and Random Sub-model Extraction Schemes

Original: Existing sub-model extraction schemes can be grouped as static (HeteroFL, FjORD) and random (Federated Dropout).
Revised: Existing sub-model extraction schemes can be grouped as static-based (HeteroFL, FjORD) and random-based (Federated Dropout) methods.

Original: We describe the differences between FedRolex and them.
Revised: In this section, we describe the differences between them and the proposed rolling-based scheme employed in FedRolex .

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

Original: 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.

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

Original: 3.3.1 Comparison with 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 remains static, 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 2.
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: i ⌋ − 1 } .
Revised: Note that here, S ( j ) n,i does not depend on j so in each round, the same sub-model is extracted for each client. This is different from FedRolex .

Original: [Figure]
Revised: [Figure]

Original: 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.
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.

Original: 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.
Revised: 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.

Original: The nodes in the beginning indices of each layer will therefore see updates that were averaged across more clients than end indices.
Revised: The nodes in the beginning indices of each layer will therefore see updates that were averaged across more clients than end indices.

Original: 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.
Revised: 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: 3.3.2 Comparison with 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 sub-models 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 right column of Figure 2.
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]
Revised: [Equation]

Original: [Equation]

Original: Here ⌊ β 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 the static sub-model extraction, all nodes see the entire data given a sufficiently large number of communication rounds.
Revised: Unlike the static sub-model extraction, all nodes see the entire data given a sufficiently large number of communication rounds.

Original: Limitations of Random Sub-model Extraction: Dropout-based algorithms like Federated
Revised: Limitations of Random Sub-model Extraction: Dropout-based algorithms like Federated

Original: Dropout [13] use random dropout to create smaller sub-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 all the indices, their exact frequencies are not the same in each experiment because of 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: 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
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 for more details and the results on updating all components at least more than one time.
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 components, and it takes too many rounds to update all the components.
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:  Experiments
Revised:  Experiments

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

Original: We evaluate the performance of FedRolex under two regimes.
Revised: We evaluate the performance of FedRolex under two regimes.

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

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

Original: We replace the batch Normalization in PreResNet18 with static batch normalization [14, 28] and add a scalar module after each convolution layer [14].
Revised: We replace the batch Normalization in PreResNet18 with static batch normalization [14, 28] and add a scalar module after each convolution layer [14].

Original: Under largemodel large-dataset regime, we use Stack Overflow [29] and followed [2] to train a modified 3-layer
Revised: Under largemodel large-dataset regime, we use Stack Overflow [29] and followed [2] to train a modified 3-layer

Original: 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 .
Revised: 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: 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.
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.

Original: The statistics of the datasets are listed in Table 2.
Revised: The statistics of the datasets are listed in Table 2.

Original: [Table]
Revised: [Table]

Original: Data Heterogeneity.
Revised: Data Heterogeneity.

Original: For CIFAR-10 and CIFAR-100, we followed HeteroFL [14] to model non-IID distributions by restricting each client to have L labels.
Revised: For CIFAR-10 and CIFAR-100, we followed HeteroFL [14] to model non-IID distributions by restricting each client to have L labels.

Original: In our evaluation, we consider two levels of data heterogeneity.
Revised: In our evaluation, we consider two levels of data heterogeneity.

Original: For CIFAR-10, we define L = 2 as high data heterogeneity and L = 5 as low data heterogeneity.
Revised: For CIFAR-10, we define L = 2 as high data heterogeneity and L = 5 as low data heterogeneity.

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

Original: These two levels roughly correspond to Dirichlet distribution Dir K ( α ) with α equal 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: For Stack Overflow, the dataset is partitioned over user IDs, making the dataset naturally non-IID distributed.
Revised: For Stack Overflow, the dataset is partitioned over user IDs, making the dataset naturally non-IID distributed.

Original: Model Heterogeneity.
Revised: Model Heterogeneity.

Original: Without loss of generality, in our evaluation, we consider five different client model capacities β
Revised: Without loss of generality, in our evaluation, we consider five different client model capacities β

Original: = { 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).
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: 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.
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: For Transformer, we 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: Baselines.
Revised: Baselines.

Original: 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 .
Revised: We compare FedRolex against both state-of-the-art PT-based model-heterogeneous FL methods including Federated Dropout [13] and HeteroFL [14] 4 as well as state-of-the-art KD-based model-heterogeneous FL methods including FedDF [10], DS-FL [11] and Fed-ET [12] 5 .

Original: 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.
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 details of the schedule for each dataset and experiment are described in Appendix A.6.
Revised: The details of the schedule for each dataset and experiment are described in Appendix A.7.

Original: Configurations and Platform.
Revised: Configurations and Platform.

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

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

Original: 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.
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.

Original: The details of the hyper-parameters for model training are included in Appendix A.6.
Revised: The details of the hyper-parameters for model training are included in Appendix A.7.

Original: We implemented FedRolex and PT-based baselines using PyTorch [31] and Ray [32], and conducted our experiments on 8 NVIDIA A 6000 GPUs.
Revised: We implemented FedRolex and PT-based baselines using PyTorch [31] and Ray [32], and conducted our experiments on 8 NVIDIA A 6000 GPUs.

Original: Evaluation Metrics.
Revised: Evaluation Metrics.

Original: We use global and local 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 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.
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.

Original: For
Revised: For

Original: CIFAR-10 and CIFAR-100, we report the classification accuracy.
Revised: CIFAR-10 and CIFAR-100, we report the classification accuracy.

Original: For Stack Overflow, we report the next word prediction accuracy which includes both out-of-vocabulary (OOV) and end-of-sentence (EOS) tokens.
Revised: For Stack Overflow, we report the next word prediction accuracy which includes both out-of-vocabulary (OOV) and end-of-sentence (EOS) tokens.

Original: We run our experiments using five different seeds for CIFAR-10 and CIFAR-100 and using three different seeds for Stack Overflow.
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.1 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: First, we compare the performance of FedRolex with state-of-the-art PT and KD-based modelheterogeneous FL methods.
Revised: First, we compare the performance of FedRolex with state-of-the-art PT and KD-based modelheterogeneous FL methods.

Original: 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.
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: Evaluation Results: Table 3 summarizes our results.
Revised: Evaluation Results: Table 3 summarizes our results.

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

Original: (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.
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: 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.
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.

Original: Under large-model large-dataset regime, FedRolex also outperforms both HeteroFL and Federated Dropout.
Revised: Under large-model large-dataset regime, FedRolex also outperforms both HeteroFL and Federated Dropout.

Original: These results together demonstrate the
Revised: These results together demonstrate the

Original: [Table]
Revised: [Table]

Original: superiority of FedRolex under both regimes.
Revised: superiority of FedRolex under both regimes.

Original: 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.
Revised: 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 note that KD-based methods leverage public data to boost their model accuracy while
Revised: It is important to note that KD-based methods leverage public data to boost their model accuracy while

Original: FedRolex does not.
Revised: FedRolex does not.

Original: 4.2 Performance Comparison with Model-Homogeneous FL Methods
Revised: 4.2 Performance Comparison with Model-Homogeneous FL Methods

Original: 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.
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.

Original: 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.
Revised: 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.

Original: 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: 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.

Original: This result indicates that with FedRolex , we will not be constrained to only using high-end devices to achieve competitive global model accuracy.
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 the model-homogeneous upper bound on CIFAR-10 under high data heterogeneity, which showcases the advantage of using 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 Client Model Heterogeneity Distribution
Revised: 4.3 Impact of Client Model Heterogeneity Distribution

Original: In our previous experiments, the distributions of model capacities across client devices are set to be uniform.
Revised: In our previous experiments, the distributions of model capacities across client devices are set to be uniform.

Original: In this experiment, we aim to understand the impact of the client model heterogeneity distribution.
Revised: In this experiment, we aim to understand the impact of the client model heterogeneity distribution.

Original: 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 } ).
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: Evaluation Results: Figure 3 shows how the global model accuracy changes when ρ varies fromto 1 for CIFAR-10, CIFAR-100 and Stack Overflow.
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 three observations.
Revised: We have three observations.

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

Original: (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).
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 and hence the global model accuracy is bottlenecked by the level of data heterogeneity instead of model capacity.
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 only limited contribution 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 (Figure 3(ii)), the gap in global model 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
Revised: In contrast to

Original: CIFAR-10, the global model accuracy is bottlenecked by the highest capacity of the models rather than the level of data heterogeneity.
Revised: 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 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.
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: 4.4 Performance on Training Larger Server Model
Revised: 4.4 Performance on Training Larger Server Model

Original: 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.
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.

Original: In this experiment, we aim to evaluate the performance of FedRolex on training larger server models.
Revised: In this experiment, we aim to evaluate the performance of FedRolex on training larger server models.

Original: 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.
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.

Original: For simplicity, all client models have the same size.
Revised: For simplicity, all client models have the same size.

Original: 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.
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.

Original: 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.
Revised: 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.

Original: For Stack Overflow (Figure 4(iii)), the global model accuracy has a much smaller drop when γ increases.
Revised: For Stack Overflow (Figure 4(iii)), the global model accuracy has a much smaller drop when γ increases.

Original: This demonstrates the superiority of using large models on large-scale datasets for training larger server models.
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 vision of FedRolex is to enhance the inclusiveness of FL.
Revised: A primary vision of FedRolex is to enhance the inclusiveness of FL.

Original: To demonstrate this, in this experiment, we use real-world household income 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 household income distribution information from Bureau [33].
Revised: Specifically, we retrieve household income distribution information from Bureau [33].

Original: 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 .
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: Detailed mapping of this distribution to the corresponding income distribution is provided in Figure 7 in Appendix A.6.
Revised: Detailed mapping of this distribution to the corresponding income distribution is provided in Figure 7 in Appendix A.7.

Original: Evaluation Results:
Revised: Evaluation Results:

Original: Table 4 shows both the global and local model accuracies of FedRolex for
Revised: Table 4 shows both the global and local model accuracies of FedRolex for

Original: CIFAR-10 and CIFAR-100 as well as the global model accuracy on Stack Overflow under the emulated real-world device distribution.
Revised: CIFAR-10 and CIFAR-100 as well as the global model accuracy on Stack Overflow under the emulated real-world device distribution.

Original: Again, we compare with two model-homogeneous cases where all clients have the smallest and largest model capacities, representing lower and upper-bound
Revised: 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: [Figure]
Revised: [Figure]

Original: accuracy, respectively.
Revised: accuracy, respectively.

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

Original: (1) Looking at the global model accuracy,
Revised: (1) Looking at the global model accuracy,

Original: FedRolex consistently outperforms the lower-bound model-homogeneous case across CIFAR-10,
Revised: FedRolex consistently outperforms the lower-bound model-homogeneous case across CIFAR-10,

Original: CIFAR-100, and Stack Overflow.
Revised: CIFAR-100, and Stack Overflow.

Original: This result indicates that FedRolex enhances the inclusiveness of
Revised: This result indicates that FedRolex enhances the inclusiveness of

Original: FL and improves the accuracy of the global model, which would otherwise not be able to achieve.
Revised: FL and improves the accuracy of the global model, which would otherwise not be able to achieve.

Original: 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.
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 effectively boosts the performance of low-end devices, which would otherwise not benefit from FL.
Revised: This result indicates that FedRolex effectively boosts the performance of low-end devices, which would otherwise not benefit from FL.

Original: 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.
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:  Conclusion
Revised:  Conclusion

Original: 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.
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: 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.
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: We provided a theoretical statistical analysis on its advantage over Federated Dropout.
Revised: We provided a theoretical statistical analysis on its advantage over Federated Dropout.

Original: 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.
Revised: 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: Moreover, we demonstrated its performance on an emulated real-world device distribution and show FedRolex contributes to making FL more inclusive.
Revised: Moreover, we demonstrated its performance on an emulated real-world device distribution and show FedRolex contributes to making FL more inclusive.

Original: Limitations and Future works.
Revised: Limitations and Future works.

Original: In this work, we provided a statistical analysis of FedRolex .
Revised: In this work, we provided a statistical analysis of FedRolex .

Original: A full convergence analysis of FedRolex is not trivial and left for future work.
Revised: A full convergence analysis of FedRolex is not trivial and left for future work.

Original: In addition, the goal of this work is to train a global server model using a federation of heterogeneous client models.
Revised: In addition, the goal of this work is to train a global server model using a federation of heterogeneous client models.

Original: 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.
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: We will pursue it as our future work.
Revised: We will pursue it as our future work.

Original:  Acknowledgement
Revised:  Acknowledgement

Original: We would like to thank the anonymous reviewers for their helpful comments.
Revised: We would like to thank the anonymous reviewers for their helpful comments.

Original: We also want to thank Jiahui Yu, Zheng Xu, and Chaoyang He for their constructive discussions.
Revised: We also want to thank Jiahui Yu, Zheng Xu, and Chaoyang He for their constructive discussions.

Original: This work was partially supported by NSF Awards PFI:BIC-1632051, CNS-1814551, DMS-2012439, and a Google
Revised: This work was partially supported by NSF Awards PFI:BIC-1632051, CNS-1814551, DMS-2012439, and a Google

Original: Computing Platform (GCP) credit grant.
Revised: Computing Platform (GCP) credit grant.