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: 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 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.
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 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].
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: In this work, we propose a PT-based model-heterogeneous FL approach named FedRolex to tackle the fundamental issue of existing methods.
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
Revised: The key difference between FedRolex and existing

Original: PT-based methods is how the sub-models are extracted for each client over communication rounds in the federated training process.
Revised: 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 echoes some concurrent efforts in developing FL primitives to support training large server models in cross-device settings, 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
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

Original: 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: 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 (Appendix A.6).
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 to a large model in the server from 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: they 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 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].
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: random sub-model extraction and static 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.

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

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

Original: Dropout becomes less effective when the data heterogeneity is high and the client cohort is small due to its randomness in selecting sub-models.
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: 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.
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 primary 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 would 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 this work, we propose a rolling sub-model extraction scheme that tackles the issues of both random and static 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 local data samples of client n 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 local data samples of client n 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 β = { β 1 , β 2 , ..., β 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]

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

Original: The size of θ n depends on β n , and the parameter θ n is obtained by selecting a sub-model from the global model θ , which can change from one round to another.
Revised: The size of θ n depends on β n , and the parameter θ n is obtained by selecting a sub-model from the global model θ , which can change from one round 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 for all clients and rounds, though they differ between clients and rounds.
Revised: For simplicity, we use the same notation l for the loss function for all clients and rounds, though they differ between clients and rounds.

Original: The key to model-heterogeneous FL is selecting θ n from the global model θ given model capacity β 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 each client,
Revised: As a partial training (PT)-based approach, at each client,

Original: 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.
Revised: 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,
Revised: To help understand how FedRolex works, for simplicity,

Original: Figure 1 illustrates three rounds of federated training of
Revised: Figure 1 illustrates three rounds of federated training of

Original: FedRolex on two participating heterogeneous clients, where one trains a large-capacity sub-model (left) and the other trains a small-capacity one (
Revised: FedRolex on two participating heterogeneous clients, where one trains a large-capacity sub-model (left) and the other trains a small-capacity one (

Original: right
Revised: right

Original: At the
Revised: At the

Original: high
Revised: high

Original: level, at each
Revised: level, at each

Original: round
Revised: round

Original: , the
Revised: , the

Original: server
Revised: server

Original: extracts sub-models of different capacities from the global model
Revised: extracts sub-models of different capacities from the global model

Original: and
Revised: and

Original: sepa-
Revised: sepa-

Original: rately broadcasts them
Revised: rately broadcasts them

Original: to the clients that have the corresponding capabilities.
Revised: to the clients that have the corresponding capabilities.

Original: The clients train the received sub-models on their local data and transmit their heterogeneous sub-model updates to the server.
Revised: The clients train the received sub-models on their local data and transmit their heterogeneous sub-model updates to the server.

Original: Lastly, the server aggregates those updates, and the result of the aggregation is used to update the global model 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: The pseudocode of FedRolex is in Algorithm
Revised: The pseudocode of FedRolex is in Algorithm

Original: The key to the design of FedRolex involves two design choices.
Revised: The key to the design of FedRolex involves two design choices.

Original: In the following, we describe them in detail.
Revised: In the following, we describe them in detail.

Original: (1) What sub-models to be extracted for each client across different rounds?
Revised: (1) What sub-models to be extracted for each client across different rounds?

Original: At the server, FedRolex utilizes a rolling window to extract the sub-model from the global model.
Revised: At the server, FedRolex utilizes a rolling window to extract the sub-model from the global model.

Original: The rolling window advances in each round, and loops over all parts of the global model in sequence across different rounds.
Revised: The rolling window advances in each round, and loops over all parts of the global model in sequence across different rounds.

Original: This process iterates such that the global model is evenly trained until convergence.
Revised: This process iterates such that the global model is evenly trained until convergence.

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

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

Original: 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.
Revised: 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: Such a rolling sub-model extraction scheme can be formalized as follows.
Revised: Such a rolling sub-model extraction scheme can be formalized as follows.

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

Original: [Equation]
Revised: [Equation]

Original: else.
Revised: else.

Original: [Equation]
Revised: [Equation]

Original: where ˆ j = j mod K i .
Revised: where ˆ j = j mod K i .

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

Original: FedRolex employs a straightforward selective averaging scheme with no client weighting to aggregate heterogeneous sub-model updates sent from the clients to update the global model 3 .
Revised: FedRolex employs a straightforward selective averaging scheme with no client weighting to aggregate heterogeneous sub-model updates sent from the clients to update the global model 3 .

Original: Specifically, it computes the average of the updates for each parameter of the global model separately based on how many clients in a round updated that parameter.
Revised: Specifically, it computes the average of the updates for each parameter of the global model separately based on how many clients in a round updated that parameter.

Original: The parameter remains unchanged if no clients updated it.
Revised: The parameter remains unchanged if no clients updated it.

Original: Taking Figure 1 again 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.
Revised: Taking Figure 1 again 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: In contrast, since c and d are part of both models, the update is computed by taking the average 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: Algorithm
Revised: Algorithm

Original: 1 Initialization ; θ (0) , N
Revised: 1 Initialization ; θ (0) , N

Original: Input : D n β n ∀ n ∈ N ,
Revised: Input : D n β n ∀ n ∈ N ,

Original: Output : θ J 2 Server Executes 3 for j ← 0 to J − 1 do
Revised: Output : θ J 2 Server Executes 3 for j ← 0 to J − 1 do

Original: Sample subset M from N
Revised: Sample subset M from N

Original: Broadcast θ ( j ) m, S ( j ) m,i to client m ∈ M ∀ i, S ( j ) m,i from Equation (4) for each client m ∈ M do clientStep
Revised: Broadcast θ ( j ) m, S ( j ) m,i to client m ∈ M ∀ i, S ( j ) m,i from Equation (4) for each client m ∈ M do clientStep

Original: end
Revised: end

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

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

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

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

Original: For comparison purpose, the pseudocodes of both Federated Dropout and HeteroFL are included in Appendix A.8.
Revised: For comparison purpose, the pseudocodes of both Federated Dropout and HeteroFL are included in Appendix A.8.

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

Original: In random sub-model extraction scheme, in each round, the sub-models are extracted from the global model in a random manner.
Revised: In random sub-model extraction scheme, in each round, the sub-models are extracted from the global model in a random manner.

Original: As such, the layer i of the sub-model extracted by the random sub-model extraction scheme for client n in round j is given by:
Revised: As such, the layer i of the sub-model extracted by the random sub-model extraction scheme for client n in round j is given by:

Original: [Equation]
Revised: [Equation]

Original: where a total number of ⌊ β n K i ⌋ nodes are randomly chosen from the global model.
Revised: where a total number of ⌊ β n K i ⌋ nodes are randomly chosen from the global model.

Original: Discussion: As shown in Figure 2(left), similar to the proposed rolling-based scheme, the sub-models extracted across different rounds by the random-based scheme have different architectures.
Revised: Discussion: As shown in Figure 2(left), similar to the proposed rolling-based scheme, the sub-models extracted across different rounds by the random-based scheme have different architectures.

Original: However,
Revised: However,

Original: [Figure]
Revised: [Figure]

Original: due to its randomness in selecting sub-models in each round, the global model is trained less evenly, making it more vulnerable to client drift.
Revised: due to its randomness in selecting sub-models in each round, the global model is trained less evenly, making it more vulnerable to client drift.

Original: In short, although the expected value of the frequency for updating each index is the same for all the indices, their exact frequencies are not the same due to randomness.
Revised: In short, although the expected value of the frequency for updating each index is the same for all the indices, their exact frequencies are not the same due to randomness.

Original: Consequently, the random-based scheme cannot balance the update frequencies of different parts of the global model, and it inevitably takes more rounds to update the whole global model.
Revised: Consequently, the random-based scheme cannot balance the update frequencies of different parts of the global model, and it inevitably takes more rounds to update the whole global model.

Original: Moreover, as we show in Appendix A.1, the expected number of rounds for Federated Dropout selecting all I sub-models at least m times is in the order of I log( I )
Revised: Moreover, as we show in Appendix A.1, the expected number of rounds for Federated Dropout selecting all I sub-models at least m times is in the order of I log( I )

Original: + I ( m − 1)
Revised: + I ( m − 1)

Original: log log I , which is larger than that of FedRolex , mI .
Revised: log log I , which is larger than that of FedRolex , mI .

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

Original: In static sub-model extraction scheme, in each round, the sub-models are always extracted from a designated part of the global model.
Revised: In static sub-model extraction scheme, in each round, the sub-models are always extracted from a designated part of the global model.

Original: As such, the layer i of the sub-model extracted by the static sub-model extraction scheme for client n in round j is given by:
Revised: As such, the layer i of the sub-model extracted by the static sub-model extraction scheme for client n in round j is given by:

Original: [Equation]
Revised: [Equation]

Original: Note that S ( j ) n,i does not depend on j .
Revised: Note that S ( j ) n,i does not depend on j .

Original: In other words, as shown in Figure 2(right), the same sub-model is extracted for each client in every round.
Revised: In other words, as shown in Figure 2(right), the same sub-model is extracted for each client in every round.

Original: Moreover, the client model with smaller capacity and client model with larger capacity are not independent .
Revised: Moreover, the client model with smaller capacity and client model with larger capacity are not independent .

Original: As shown in Figure 2(right), the small-capacity model { a, b, c } is a part of the medium-capacity model { a, b, c, d } , which again, is a part of the large-capacity model { a, b, c, d, e } .
Revised: As shown in Figure 2(right), the small-capacity model { a, b, c } is a part of the large-capacity model { a, b, c, d } , which again, is a part of the global-capacity model { a, b, c, d, e } .

Original: These are the two key differences from both the random-based and the proposed rolling-based scheme.
Revised: These are the two key differences from both the random-based and the proposed rolling-based scheme.

Original: Discussion:
Revised: Discussion:

Original: Given that, static-based scheme, however, has two primary drawbacks.
Revised: Given that, static-based scheme, however, has two primary drawbacks.

Original: First,  the global model is restricted to the same size as the largest client model.There must be clients training with the full global model.
Revised: First, to cover the whole global model, there must be clients to train the full-size global model { a, b, c, d, e } . As such, the global model is restricted to the same size as the largest client model.

Original: Second, depending on their resource demands, different sub-models can only be trained on clients whose on-device resources are matched. As such, 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. We can see this in Figure 2 where a , b and c will see updates from all three sub-models.
Revised: Second, as shown in Figure 2(right), while a , b and c will be trained on data on all three types of clients, d will not be trained on data on small-capacity clients, and e will only be trained on data on global-model-capacity clients.

Original: but e will only see updates from large-capacity clients. This would limit the performance of the global model.
Revised: As a consequence, different parts of the global model are trained on data with different distributions, which inevitably degrades the global model training quality.

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] 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 .
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.7.
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.7.
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 
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. In particular, under high data heterogeneity, Federated Dropout

Original: [Table]
Revised: [Table]

Original: high data heterogeneity scenarios. 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: 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 superiority of FedRolex under both regimes.
Revised: These results together demonstrate the 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 global model accuracy changes when ρ varies from 0 to 1 for CIFAR-10, CIFAR-100 and Stack Overflow.
Revised: Evaluation Results: Figure 3 shows how global model accuracy changes when ρ varies from 0 to 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 
Revised: CIFAR-10, the global model accuracy is bottlenecked by the highest capacity of the models rather than the level of data heterogeneity.

Revised: For both regimes (Figure 3(i)(ii) vs. Figure 3(iii)), we observe

Original: [Figure]
Revised: [Figure]

Original: than the level of data heterogeneity. 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 the accuracy.
Revised: 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 the accuracy.

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]

Revised: [Equation]

Revised: [Figure]

Revised: [Equation]

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

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: iv)
Revised: iv)

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 accuracy, respectively.
Revised: Again, we compare with two model-homogeneous cases where all clients have the smallest and largest model capacities, representing lower and upper-bound 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 thank the reviewers for their helpful comments.
Revised: We thank the reviewers for their helpful comments.

Original: This work was partially supported by NSF
Revised: This work was partially supported by NSF

Original: PFI:BIC-1632051, CNS-1814551, DMS-2012439, and a Google Computing Platform (GCP) grant.
Revised: PFI:BIC-1632051, CNS-1814551, DMS-2012439, and a Google Computing Platform (GCP) grant.