> $\|\mathcal{M}(H)\|_{L^{2}(\mu)}\leq(1-\varepsilon)\|H\|_{L^{2}(\mu)}$
>
>

>
> The first point is done using the so called Lasota-Yorke inequalities, see
> the lemma 5.1 in [OW17]. The second point goes, loosely, as follow.
>
> We write, using Cauchy-Schwartz:
>
>
> $(\mathcal{M}H)^{2}=\mathcal{L}_{\varphi}^{N}(H\beta)^{2}\leq\mathcal{L}^{N}_{\varphi}(H^{2})\mathcal{L}^{N}_{\varphi}(\beta^{2})$
>
> On $\widehat{S}$, the cancellation represented in the function $\beta$
> spreads, thanks to the fact that $\varphi$ is normalized, as follow:
>
> $\mathcal{L}_{\varphi}^{N}(\beta^{2})=\sum_{\mathbf{a}}e^{S_{N}\varphi\circ
> g_{\mathbf{a}}}\beta^{2}\circ g_{\mathbf{a}}$ $=\sum_{\mathbf{a}\text{ where
> }\beta=1}e^{S_{N}\varphi\circ g_{\mathbf{a}}}\beta^{2}\circ
> g_{\mathbf{a}}+\sum_{\mathbf{a}\text{ where }\beta\text{ is
> smaller}}e^{S_{N}\varphi\circ g_{\mathbf{a}}}\beta^{2}\circ g_{\mathbf{a}}$
> $\leq\sum_{\mathbf{a}\text{ where }\beta=1}e^{S_{N}\varphi\circ
> g_{\mathbf{a}}}+\sum_{\mathbf{a}\text{ where }\beta\text{ is
> smaller}}e^{S_{N}\varphi\circ g_{\mathbf{a}}}(1-\eta)$ $=1-\eta
> e^{-N\|\varphi\|_{\infty}}.$
>
> Then, we use the doubling property of $\mu=\sum_{a}\mu_{a}$ and the control
> given by the fact that $H\in K_{R}(\mathfrak{U})$ to bound the integral on
> all $J$ by the integral on $\widehat{S}$:
>
> $\int_{J}\mathcal{L}^{N}_{\varphi}(H^{2})d\mu\leq
> C_{0}\int_{\widehat{S}}\mathcal{L}^{N}_{\varphi}(H^{2})d\mu.$
>
> Hence, we can write, using the fact that $\mathcal{L}_{\varphi}$ preserves
> $\mu$ and the previously mentioned Cauchy-Schwartz inequality:
>
>
> $\|H\|_{L^{2}(\mu)}-\|\mathcal{M}H\|_{L^{2}(\mu)}=\int_{J}\left(\mathcal{L}_{\varphi}^{N}(H^{2})-\mathcal{L}_{\varphi}(H^{2})\mathcal{L}_{\varphi}(\beta^{2})\right)d\mu$
> $\geq\int_{\widehat{S}}\left(\mathcal{L}_{\varphi}^{N}(H^{2})-\mathcal{L}_{\varphi}(H^{2})\mathcal{L}_{\varphi}(\beta^{2})\right)d\mu$
> $\geq\eta
> e^{-N\|\varphi\|_{\infty}}\int_{\widehat{S}}\mathcal{L}_{\varphi}^{N}(H^{2})d\mu$
> $\geq\eta
> C_{0}^{-1}e^{-N\|\varphi\|_{\infty}}\|H\|_{L^{2}(\mu)}:=\varepsilon\|H\|_{L^{2}(\mu)}$
>
> Hence,
>
> $\|\mathcal{M}H\|_{L^{2}(\mu)}\leq(1-\varepsilon)\|H\|_{L^{2}(\mu)}.$

Step 4: We conclude by an iterative argument.

> To conclude, we need to see that we may bound $h$ by some $H\in
> K_{R}(\mathfrak{U})$, and also that the contraction property is true for all
> $n$, not just $N$.
>
> For any $n=kN$, we can inductively prove our bound. If $k=1$, we can choose
> $H_{0}:=\|h\|_{|t|+|l|}$. Then, $|h|\leq H_{0}$ and so
>
>
> $\|\mathcal{L}_{\varphi-i(t\tau+l\theta)}^{N}h\|_{L^{2}(\mu)}\leq\|\mathcal{M}H_{0}\|_{L^{2}(\mu)}\leq(1-\varepsilon)\|h\|_{(|t|+|l|)}$
>
> Then, choosing $H_{k+1}:=\mathcal{M}H_{k}\in K_{R}(\mathfrak{U})$, we can
> proceed to the next step of the induction and get
>
>
> $\|\mathcal{L}_{\varphi-i(t\tau+l\theta)}^{kN}h\|_{L^{2}(\mu)}\leq\|\mathcal{M}H_{k-1}\|_{L^{2}(\mu)}\leq(1-\varepsilon)^{k}\|h\|_{(|t|+|l|)}$
>
> Finally, if $n=kN+r$, with $0\leq j\leq N-1$, we write
>
>
> $\|\mathcal{L}_{\varphi-i(t\tau+l\theta)}^{kN+r}h\|_{L^{2}(\mu)}\leq(1-\varepsilon)^{k}\|\mathcal{L}_{\varphi}^{r}h\|_{(|t|+|l|)}\lesssim(1-\varepsilon)^{n/N}\|h\|_{(|t|+|l|)}$
>
> And the proof is done.

## References

  * [Be91] A.F. Beardon, _Iteration of Rational Functions: Complex Analytic Dynamical Systems_ , Graduate Texts in Mathematics, 1991, Springer New York
  * [Bl96] C. Bluhm, _Random recursive construction of Salem sets_ Ark. Mat. 34 (1996), 51-63.
  * [Bo10] J. Bourgain, _The discretized sum-product and projection theorems_ JAMA 112, 193-236 (2010) https://doi.org/10.1007/s11854-010-0028-x
  * [BD17] J. Bourgain, S. Dyatlov, _Fourier dimension and spectral gaps for hyperbolic surfaces_ , Geom. Funct. Anal. 27, 744–771 (2017), https://arxiv.org/abs/1704.02909
  * [Do98] D. Dolgopyat, _On decay of correlations in Anosov flows_ ,
Ann. of Math. (2) 147 (2) (1998) 357–390.

  * [ES11] A. Eremenko and S. Van Strien. _Rational maps with real multipliers_ , Trans. AMS, 363, (2011) p. 6453-6463.
  * [Fr35] O. Frostman, _Potentiel d’équilibre et capacité des ensembles avec quelques applications à la théorie des fonctions._ Meddel. Lunds Univ. Math. Sem. 3 (1935), 1-118
  * [Gr09] B. Green, _Sum-product phenomena in $\mathbb{F}_{p}$: a brief introduction_
Notes written from a Cambridge course on Additive Combinatorics, 2009.
https://arxiv.org/pdf/0904.2075.pdf

  * [Ha17] K. Hambrook, _Explicit Salem sets in $\mathbb{ℝ}^{2}$._ Adv. Math. 311 (2017), 634–648.
  * [JS16] T. Jordan, T. Sahlsten _Fourier transforms of Gibbs measures for the Gauss map_
Math. Ann., 364(3-4), 983-1023, 2016 https://arxiv.org/abs/1312.3619

  * [Ka66] J. P. Kahane, _Images d’ensembles parfaits par des séries de Fourier Gaussiennes_ Acad. Sci. Paris 263 (1966), 678-681.
  * [Kau81] R. Kaufman, _On the theorem of Jarnik and Besicovitch_ , Acta Arith. 39 (1981), 265-267.
  * [Li18] J. Li, _Discretized Sum-product and Fourier decay in_ $\mathbb{R}^{n}$, JAMA 143, 763–800 (2021).
https://doi.org/10.1007/s11854-021-0169-0, https://arxiv.org/abs/1811.06852v2

  * [Li18b] J. Li, _Fourier decay, Renewal theorem and Spectral gaps for random walks on split semisimple Lie groups_ preprint 2018, https://arxiv.org/abs/1811.06852v2
  * [LNP19] Li, Naud, Pan, _Kleinian Schottky groups, Patterson-Sullivan measures and Fourier decay_
Duke Math. J. 170(4): 775-825 (15 March 2021). DOI: 10.1215/00127094-2020-0058
https://arxiv.org/abs/1902.01103

  * [Ma15] P. Mattila, _Fourier analysis and Hausdorff dimension_ , Cambridge University Press, 2015.
  * [Mi90] J. Milnor, _Dynamics in one complex variable_
Stony Brook IMS Preprint 1990/5, https://arxiv.org/abs/math/9201272

  * [MR92] R. Mane, L.F. Da Rocha, _Julia sets are uniformly perfect_
Proc. Amer. Math. Soc. 116 (1992), no. 1, 251–257.

  * [Na05] F. Naud, _Expanding maps on Cantor sets and analytic continuation of zeta functions_ Volume 38, Issue 1, 2005, Pages 116-153, ISSN 0012-9593, https://doi.org/10.1016/j.ansens.2004.11.002.
  * [OW17] H. Oh, D. Winter, _Prime number theorem and holonomies for hyperbolic rational maps_
Inventiones mathematicae, 2017. 208. 10.1007/s00222-016-0693-1.
https://arxiv.org/abs/1603.00107

  * [PP90] W. Parry, M. Pollicott, _Zeta Functions and the Periodic Orbit Structure of Hyperbolic Dynamics_ Asterisque; 187-188, société mathématique de France, 1990.
  * [PS16] V. Petkov, L. Stoyanov, _Ruelle transfer operators with two complex parameters and applications_ Discrete and Continuous Dynamical Systems, 2016, 36 (11) : 6413-6451. doi: 10.3934/dcds.2016077
  * [PU09] F. Przytycki, M. Urbański, _Conformal fractals: ergodic theory methods_
Cambridge university, 2009.

  * [PU17] M. Pollicott and M. Urbański, _Asymptotic Counting in Conformal Dynamical Systems._ Memoirs of the American Mathematical Society. 271. 10.1090/memo/1327. (2017).
  * [PW97] Y. Pesin and H. Weiss, _A multifractal analysis of equilibrium measures for conformal expanding maps and Moranlike geometric constructions_ ,
Journal of Statistical Physics 86, 233—275, (1997)

  * [QR03] M. Queffélec and O. Ramaré, _Analyse de Fourier des fractions continues à quotients restreints_ , Enseign. Math. (2), 49(3-4):335–356, 2003
  * [Ru78] D. Ruelle, _Thermodynamic Formalism_ , Second edition, Cambridge University Press 2004
  * [Ru89] D. Ruelle, _The thermodynamic formalism for expanding maps_ ,
Commun. Math. Phys. 125, 239-262 (1989)

  * [Sa51] R. Salem, _On singular monotonic functions whose spectrum has a given Hausdorff dimension._ Ark. Mat., 1:353–365, 1951.
  * [ShSt20] R. Sharp, A. Stylianou, _Statistics of multipliers for hyperbolic rational maps_ ,
Preprint, 2020. https://arxiv.org/abs/2010.15646

  * [SS20] T. Sahlsten, C. Stevens, _Fourier transform and expanding maps on Cantor sets_
preprint, 2020. https://arxiv.org/abs/2009.01703v4

  * [St11] L. Stoyanov, _Spectra of Ruelle transfer operators for Axiom A flows._ Nonlinearity, 24(4):1089, 2011.


# FedRolex: Model-Heterogeneous Federated Learning with Rolling Sub-Model
Extraction

Samiul Alam1,2, Luyang Liu3, Ming Yan4,1, Mi Zhang2,1
1Michigan State University, 2The Ohio State University,
3Google Research, 4The Chinese University of Hong Kong, Shenzhen
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Most cross-device federated learning (FL) studies focus on the model-
homogeneous setting where the global server model and local client models are
identical. 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. 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. At its core, 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. We show that
FedRolex outperforms state-of-the-art PT-based model-heterogeneous FL methods
(e.g. Federated Dropout) and reduces the gap between model-heterogeneous and
model-homogeneous FL, especially under the large-model large-dataset regime.
In addition, we provide theoretical statistical analysis on its advantage over
Federated Dropout 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. Our code is available at: https://github.com/AIoT-MLSys-Lab/FedRolex.

##  1 Introduction

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]. In this work, we focus on cross-device FL where clients are
usually resource-constrained edge devices. The majority of existing cross-
device FL studies focus on the model-homogeneous setting [3, 4, 5, 6], in
which the server model and the client models across all the participating
client devices are identical. However, model-homogeneous FL are confronted
with two fundamental constraints: (1) device heterogeneity is a more realistic
consideration when deploying FL systems in real-world applications: different
client devices could have very diverse on-device resources and are only
capable of training models with capacities that match their on-device
resources. Having the same model on all the devices would, unfortunately,
exclude clients with low-end devices who would otherwise make unique
contributions to model training from their own local data; (2) state-of-the-
art machine learning has moved towards large models [7] such as Transformer
[8]. 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.

Table 1: Comparison of FedRolex with model-homogeneous and model-heterogeneous
FL methods.

|  | Model
---
Heterogeneity
| Aggregation
---
Scheme
| Sub-model
---
Extraction Scheme
| Need of
---
Public Data
| Server Model
---
Size
| Compatibility with
---
Secure Aggregation
FedAvg [3] | No | - | - | No |  $=$ Client Model | Yes
FedProx [4] | No |  $=$ Client Model | Yes
SCAFFOLD [5] | No |  $=$ Client Model | Yes
FedBE [6] | Unlabeled |  $=$ Client Model | No
FedGKT [9] | Yes |  Knowledge Distillation | - | No |  $\geq$ Largest Client Model | No
FedDF [10] | Unlabeled |  $=$ Largest Client Model | No
DS-FL [11] | Unlabeled |  $=$ Largest Client Model | No
Fed-ET [12] | Unlabeled |  $\geq$ Largest Client Model | No
Federated Dropout [13] | Yes |  Partial Training | Random | No |  $\geq$ Largest Client Model | Yes
HeteroFL [14] | Static | No |  $=$ Largest Client Model | Yes
FjORD [15] | Static | No |  $=$ Largest Client Model | Yes
FedRolex (Our Approach) | Rolling | No |  $\geq$ Largest Client Model | Yes

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. One primary challenge in model-heterogeneous FL is the
aggregation of heterogeneous client models. To address this challenge,
knowledge distillation (KD)-based approaches have been proposed [9, 10, 11,
12], in which the client models serve as teachers and the server ensembles the
knowledge distilled from the individual client models. However, KD-based
approaches in general require public data on the server to achieve competitive
model accuracy, whereas the desired public data may not be always available in
practice. Moreover, since KD-based approaches need the individual client
models (whole models or prediction layers) or their outputs to be sent to the
server, they are incompatible with secure aggregation protocols [16], which
limits their privacy guarantee. To remove the dependency on public data and
ensure compatibility with secure aggregation, partial training (PT)-based
approaches such as random sub-model extraction (Federated Dropout [13]) and
static sub-model extraction (HeteroFL [14], FjORD [15]) were proposed. In
these approaches, each client trains a smaller sub-model extracted from the
larger global server model, and the server model is updated by aggregating
those trained sub-models. However, the fundamental issue of existing PT-based
methods is that the sub-models are extracted in ways (either random or static)
such that the parameters of the global server model are not evenly trained.
This makes the server model vulnerable to client drift111In model-homogeneous
FL, client drift is primarily induced by data heterogeneity across clients. In
model-heterogeneous FL, model heterogeneity across clients is another critical
source that induces client drift. induced by the inconsistency between
individual client model and server model architectures – a unique challenge of
model-heterogeneous FL.

In this work, we propose a PT-based model-heterogeneous FL approach named
FedRolex to tackle the fundamental issue of existing methods. 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. Specifically, instead of extracting sub-models in
either random or static manner, FedRolex proposes a rolling sub-model
extraction scheme, where the sub-model is extracted from the global server
model using a rolling window that advances in each communication round. Since
the window is rolling, sub-models from different parts of the global model are
extracted in sequence in different rounds. As a result, all the parameters of
the global server model are evenly trained over the local data of client
devices.

The proposed rolling sub-model extraction scheme, though simple, has equipped
FedRolex with multifold merits compared to prior arts (Table 1): (1) FedRolex
enables different parts of the global server model to be evenly trained, which
mitigates the client drift induced by model heterogeneity. (2) Contrary to
static sub-model extraction approaches (HeteroFL, FjORD), FedRolex is able to
train a global server model that is larger than the largest client model,
enabling FL to benefit from the superior performance brought by large models.
It echoes some concurrent efforts in developing FL primitives to support
training large server models in cross-device settings, e.g. Federated Select
[17]. (3) Compared to random sub-model extraction (Federated Dropout), as we
show in our theoretical statistical analysis in Section 3 and Appendix A.1,
the global server model is trained more evenly by FedRolex as the expected
number of rounds for FedRolex going through all the parameters of the global
model for at least certain times is smaller than that of Federated Dropout.
(4) FedRolex only needs to transmit the sub-model that is needed by a given
client instead of the full server model to the client. This allows clients to
contribute to federated training under resource constraints and reduces
communication overheads (Appendix A.6). (5) Lastly, FedRolex is fully
compatible with existing secure aggregation protocols that enhance the privacy
properties of FL systems.

We evaluate the performance of FedRolex under two regimes: i) small-model
small-dataset regime (most existing cross-device FL studies use this
combination), and ii) large-model large-dataset regime (this combination echos
recent efforts on pushing the frontier of cross-device FL towards training
large server models on large-scale datasets [18, 19, 20, 21]). We highlight
five of our findings: (1) FedRolex consistently outperforms state-of-the-art
PT-based model-heterogeneous FL methods under both small-model small-dataset
and large-model large-dataset regimes (§4.1). (2) FedRolex reduces the gap
between model-heterogeneous and model-homogeneous FL, especially under large-
model large-dataset regime (§4.2). (3) With FedRolex, under both regimes,
having a small fraction of large-capacity models could significantly boost the
global model accuracy (§4.3). (4) FedRolex is able to train a global server
model that is larger than the largest client model and outperforms Federated
Dropout in terms of global model accuracy (§4.4). (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).

##  2 Related Work

Knowledge Distillation (KD)-based Model-Heterogeneous FL. One primary approach
for model-heterogeneous FL in cross-device settings is based on knowledge
distillation (KD) [22]. In particular, FedDF [10] distills knowledge from a
set of classifiers trained with private data from a federation of client
devices. The logit outputs of each classifier against an unlabeled public
dataset are then used to train a student model at the server with KD.
Similarly, DS-FL [11] utilized an unlabeled public dataset at the server and
proposed a distillation-based semi-supervised FL approach to enhance
performance by pseudo-labeling the public data. FedGKT [9] proposed group
knowledge transfer in which knowledge is transferred to a large model in the
server from clients without public data. 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. KD-based
approaches, however, have several limitations: they often require public data
to achieve competitive model accuracy. This is because model accuracy is
dependent on the size of public data as well as the domain similarity of
public data with client data [10, 12, 23]. 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].

Partial Training (PT)-based Model-Heterogeneous FL. To address the limitations
of KD-based approaches, partial training (PT) has emerged as another solution
for model-heterogeneous FL. 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. 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. Though easy to be integrated into existing
FL frameworks, as reported in [25], Federated Dropout becomes less effective
when the data heterogeneity is high and the client cohort is small due to its
randomness in selecting sub-models. 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. However, such a static
extraction strategy has two primary drawbacks. First, the global server model
is restricted to the same size as the largest client model. As such, the size
and capability of the global model are implicitly restricted by the resources
of client devices, making it not able to train large models due to resource
bottlenecks at client devices. 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. As a consequence,
part of the global server model cannot be trained on data at low-end client
devices, causing different parts of the global model to be trained on data
with different distributions. This would degrade the performance of the global
model, especially under high data heterogeneity. In this work, we propose a
rolling sub-model extraction scheme that tackles the issues of both random and
static sub-model extraction methods.

##  3 Methodology

###  3.1 Formulation of Model-Heterogeneous FL

Let $\mathcal{N}$ denote $N$ client devices with non-IID (non-identically and
independently distributed) local data $D=\\{D_{1},D_{2},...,D_{N}\\}$. Model-
homogeneous FL trains a global model of parameter $\theta$ by solving the
following optimization problem:

$\displaystyle\min_{\theta}F\left(\theta\right)\triangleq\sum_{n=1}^{N}p_{n}F_{n}(\theta)$
(1)

with

$\displaystyle
F_{n}(\theta)\triangleq\frac{1}{m_{n}}\sum_{k=1}^{m_{n}}l(\theta;d_{n,k}),$
(2)

where $D_{n}\triangleq\\{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}\geq 0$ and $\sum_{n=1}^{N}p_{n}=1$.

In comparison, in model-heterogeneous FL, clients train local models with
heterogeneous capacities
$\boldsymbol{{\beta}}=\\{\beta_{1},\beta_{2},...,\beta_{N}\\}$, and the local
objective function of the $n^{th}$ client becomes

$\displaystyle
F^{\prime}_{n}(\theta_{n})\triangleq\frac{1}{m_{n}}\sum_{k=1}^{m_{n}}l(\theta_{n};d_{n,k}).$
(3)

Here, $\beta_{n}$ denote the model capacity of client $n$, and we define it as
the proportion of nodes extracted from each layer in $\theta$ for client $n$.
The size of $\theta_{n}$ depends on $\beta_{n}$, and the parameter
$\theta_{n}$ is obtained by selecting a sub-model from the global model
$\theta$, which can change from one round to another. If $\theta_{n}$ changes,
the objective function also changes. For simplicity, we use the same notation
$l$ for the loss function for all clients and rounds, though they differ
between clients and rounds. The key to model-heterogeneous FL is selecting
$\theta_{n}$ from the global model $\theta$ given model capacity $\beta_{n}$.

###  3.2 FedRolex: Model-Heterogeneous FL with Rolling Sub-Model Extraction

Figure 1: Overview of the rolling sub-model extraction scheme in FedRolex.

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 sub-model updates back to the server for update aggregation. To
help understand how FedRolex works, for simplicity, Figure 1 illustrates three
rounds of federated training of FedRolex on two participating heterogeneous
clients, where one trains a large-capacity sub-model (left) and the other
trains a small-capacity one (right). At the high level, at each round, the
server extracts sub-models of different capacities from the global model and
separately 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 sub-model updates to the server. Lastly, the
server aggregates those updates, and the result of the aggregation is used to
update the global model for the next round. The pseudocode of FedRolex is in
Algorithm 1.

The key to the design of FedRolex involves two design choices. In the
following, we describe them in detail.

(1) What sub-models to be extracted for each client across different rounds?
At the server, FedRolex utilizes a rolling window to extract the sub-model
from the global model. 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.

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
step222The step size is a hyperparameter of FedRolex. Please refer to Appendix
A.4 for our ablation study on it., 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.

Such a rolling sub-model extraction scheme can be formalized as follows. Let
$\theta^{(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 $\mathcal{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$. 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:

$\displaystyle\mathcal{S}^{(j)}_{n,i}$
$\displaystyle=\left\\{\begin{array}[]{ll}\\{\hat{j},\hat{j}+1,\dots,\hat{j}+\lfloor\beta_{n}K_{i}\rfloor-1\\}&\mbox{
if }\hat{j}+\lfloor\beta_{n}K_{i}\rfloor\leq K_{i},\\\
\\{\hat{j},\hat{j}+1,\dots,K_{i}-1\\}\cup\\{0,1,\dots,\hat{j}+\lfloor\beta_{n}K_{i}\rfloor-1-K_{i}\\}&\mbox{
else.}\end{array}\right.$ (6)

where $\hat{j}=j\bmod{K_{i}}$.

(2) How to aggregate heterogeneous sub-model updates to update the global
model? 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 model333In Appendix A.3, we did an ablation study
on three client weighting schemes and compared them with the non-weighting
(selective averaging) scheme. We find that the performance of the three
weighting schemes is not significantly better than the non-weighting scheme.
Please refer to Appendix A.3 for details.. 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. The parameter remains
unchanged if no clients updated it.

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. In contrast, since $c$ and $d$ are part of both
models, the update is computed by taking the average from both models.

1

Initialization ; $\theta^{(0)}$, $\mathcal{N}$

Input : $D_{n}\;\beta_{n}\;\forall n\in\mathcal{N}$,

Output : $\theta^{J}$

2 Server Executes

3 for  _$j\leftarrow 0$ to $J-1$_ do

4 Sample subset $\mathcal{M}$ from $\mathcal{N}$

5 Broadcast $\theta^{(j)}_{m,\mathcal{S}^{(j)}_{m,i}}$ to client
$m\in\mathcal{M}$

6 $\qquad\forall i,\mathcal{S}^{(j)}_{m,i}\;\text{from
Equation}\;\eqref{eqn:subset_fedrolex}$

7 for  _each client $m\in\mathcal{M}$_ do

8 clientStep(_$\theta^{(j)}_{m}$ , $D_{m}$_)

9 end for

10 Aggregate $\theta^{(j+1)}_{[i,k]}$ according to Equation (12)

11 end for

12 Subroutine  _clientStep(_$\theta^{(j)}_{n}$ , $D_{n}$_)_

13 $m_{n}\longleftarrow len(D_{n})$

14 for  _$k\leftarrow 0$ to $m_{n}$_ do

15 $\theta_{n}\longleftarrow\theta_{n}-\eta\nabla l(\theta_{n};d_{n,k})$

16 end for

17 return $\theta_{n}$

18

Algorithm 1 FedRolex Figure 2: Illustration of how sub-models are extracted by
random sub-model extraction scheme (Left) and static sub-model extraction
scheme (Right) over two rounds.

###  3.3 Comparison with Random and Static Sub-model Extraction Schemes

Existing sub-model extraction schemes can be grouped as random-based
(Federated Dropout) and static-based (HeteroFL, FjORD) methods. In this
section, we describe the differences between them and the proposed rolling-
based scheme employed in FedRolex. For comparison purpose, the pseudocodes of
both Federated Dropout and HeteroFL are included in Appendix A.8.

####  3.3.1 Comparison with Random Sub-Model Extraction Scheme

In random sub-model extraction scheme, in each round, the sub-models are
extracted from the global model in a random manner. 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:

$\displaystyle\mathcal{S}^{(j)}_{n,i}=\\{k_{c}\;|\;\mbox{integer
}k_{c}\in[0,K_{i}-1]~{}\text{for}~{}1\leq
c\leq\lfloor\beta_{n}K_{i}\rfloor\\},$ (7)

where a total number of $\lfloor\beta_{n}K_{i}\rfloor$ nodes are randomly
chosen from the global model.

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. However, 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. 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.
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. 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)+I(m-1)\log\log I$, which is
larger than that of FedRolex, $mI$.

####  3.3.2 Comparison with Static Sub-Model Extraction Scheme

In static sub-model extraction scheme, in each round, the sub-models are
always extracted from a designated part of the global model. 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:

$\displaystyle\mathcal{S}^{(j)}_{n,i}=\\{0,1,2,\dots,\lfloor\beta_{n}K_{i}\rfloor-1\\}.$
(8)

Note that $\mathcal{S}^{(j)}_{n,i}$ does not depend on $j$. In other words, as
shown in Figure 2(right), the same sub-model is extracted for each client in
every round. Moreover, the client model with smaller capacity and client model
with larger capacity are not independent. 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\\}$. These are the two key differences from both the random-
based and the proposed rolling-based scheme.

Discussion: Given that, static-based scheme, however, has two primary
drawbacks. 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. 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. As a consequence, different parts of the global model are trained on
data with different distributions, which inevitably degrades the global model
training quality.

##  4 Experiments

Datasets and Models. We evaluate the performance of FedRolex under two
regimes. Under small-model small-dataset regime, we train pre-activated
ResNet18 (PreResNet18) models [26] on CIFAR-10 and CIFAR-100 [27]. We replace
the batch Normalization in PreResNet18 with static batch normalization [14,
28] and add a scalar module after each convolution layer [14]. Under large-
model large-dataset regime, we use Stack Overflow [29] and followed [2] to
train a modified 3-layer Transformer [8] with a vocabulary of $10,000$ words,
where the dimension of token embeddings is $128$, and the hidden dimension of
the feed-forward network (FFN) block is $2048$. We use ReLU activation and use
$8$ heads for the multi-head attention where each head is based on
12-dimensional (query, key, value) vectors. The statistics of the datasets are
listed in Table 2.

Table 2: Dataset statistics.

Dataset | Train Clients | Train Examples | Validation Clients | Validation Examples | Test Clients | Test Examples
---|---|---|---|---|---|---
CIFAR-10 | $100$ | 50,000 | N/A | N/A | N/A | 10,000
CIFAR-100 | $100$ | 50,000 | N/A | N/A | N/A | 10,000
Stack Overflow | 342,477 | 135,818,730 | 38,758 | 16,491,230 | 204,088 | 16,586,035

Data Heterogeneity. For CIFAR-10 and CIFAR-100, we followed HeteroFL [14] to
model non-IID distributions by restricting each client to have $L$ labels. In
our evaluation, we consider two levels of data heterogeneity. For CIFAR-10, we
define $L=2$ as high data heterogeneity and $L=5$ as low data heterogeneity.
For CIFAR-100, we use $L=20$ as high data heterogeneity and $L=50$ as low data
heterogeneity. These two levels roughly correspond to Dirichlet distribution
$Dir_{K}(\alpha)$ with $\alpha$ equal to $0.1$ and $0.5$, respectively. For
Stack Overflow, the dataset is partitioned over user IDs, making the dataset
naturally non-IID distributed.

Model Heterogeneity. Without loss of generality, in our evaluation, we
consider five different client model capacities
$\boldsymbol{\beta}=\\{1,\nicefrac{{1}}{{2}},\nicefrac{{1}}{{4}},\nicefrac{{1}}{{8}},\nicefrac{{1}}{{16}}\\}$
where for instance, $\nicefrac{{1}}{{2}}$ means the client model capacity is
half of the largest client model capacity (full model). 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. For Transformer, we
vary the number of nodes in the hidden layer of the attention heads.

Baselines. We compare FedRolex against both state-of-the-art PT-based model-
heterogeneous FL methods including Federated Dropout [13] and HeteroFL
[14]444We did not compare with FjORD because its code is not open-source and
we could not reproduce their results following the paper. as well as state-of-
the-art KD-based model-heterogeneous FL methods including FedDF [10], DS-FL
[11] and Fed-ET [12] 555We did not compare with FedGKT [9] as it is only
compatible with CNN models.. 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. The details of the
schedule for each dataset and experiment are described in Appendix A.7.

Configurations and Platform. For CIFAR-10 and CIFAR-100, we apply bounding box
crop [30] to augment the images. In each communication round, $10\%$ of the
clients are randomly selected from a pool of $100$ clients. For Stack
Overflow, we followed [2] to use a $10\%$ dropout rate to prevent over-
fitting, and $200$ clients are randomly selected from a pool of $342,477$
clients in each communication round. The details of the hyper-parameters for
model training are included in Appendix A.7. We implemented FedRolex and PT-
based baselines using PyTorch [31] and Ray [32], and conducted our experiments
on $8$ NVIDIA A$6000$ GPUs.

Evaluation Metrics. We use global and local model accuracy as our evaluation
metrics. Specifically, global model accuracy is defined as the server model
accuracy on the test set; and local model accuracy is defined as the accuracy
of the server model on each of the client’s local datasets. For CIFAR-10 and
CIFAR-100, we report the classification accuracy. For Stack Overflow, we
report the next word prediction accuracy which includes both out-of-vocabulary
(OOV) and end-of-sentence (EOS) tokens. We run our experiments using five
different seeds for CIFAR-10 and CIFAR-100 and using three different seeds for
Stack Overflow.

###  4.1 Performance Comparison with State-of-the-Art Model-Heterogeneous FL
Methods

First, we compare the performance of FedRolex with state-of-the-art PT and KD-
based model-heterogeneous FL methods. 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.

Evaluation Results: Table 3 summarizes our results. We have two observations.
(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
which extracts sub-model randomly has worse performance than FedRolex and
HeteroFL which both extract sub-models in a deterministic manner. Under large-
model large-dataset regime, FedRolex also outperforms both HeteroFL and
Federated Dropout. These results together demonstrate the superiority of
FedRolex under both regimes. (2) 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. It is important to note that
KD-based methods leverage public data to boost their model accuracy while
FedRolex does not.

Table 3: Global model accuracy comparison between FedRolex, PT and KD-based
model-heterogeneous FL methods, and model-homogeneous FL methods. Note that
the results of KD-based methods were obtained from [12]. For Stack Overflow,
since KD-based methods cannot be directly used for language modeling tasks,
their results are marked as N/A.

| Method |  | High Data Heterogeneity |  | Low Data Heterogeneity | Stack Overflow
---|---|---|---|---|---|---
|  | CIFAR-10 | CIFAR-100 |  | CIFAR-10 | CIFAR-100
KD-based | FedDF |  | 73.81 (± 0.42) | 31.87 (± 0.46) |  | 76.55 (± 0.32) | 37.87 (± 0.31) | N/A
DS-FL |  | 65.27 (± 0.53) | 29.12 (± 0.51) |  | 68.44 (± 0.47) | 33.56 (± 0.55) | N/A
Fed-ET |  | 78.66 (± 0.31) | 35.78 (± 0.45) |  | 81.13 (± 0.28) | 41.58 (± 0.36) | N/A
PT-based | HeteroFL |  | 63.90 (± 2.74) | 52.38 (± 0.80) |  | 73.19 (± 1.71) | 57.44 (± 0.42) | 27.21 (± 0.22)
Federated Dropout |  | 46.64 (± 3.05) | 45.07 (± 0.07) |  | 76.20 (± 2.53) | 46.40 (± 0.21) | 23.46 (± 0.12)
FedRolex |  | 69.44 (± 1.50) | 56.57 (± 0.15) |  | 84.45 (± 0.36) | 58.73 (± 0.33) | 29.22 (± 0.24)
| Homogeneous (smallest) |  | 38.82 (± 0.88) | 12.69 (± 0.50) |  | 46.86 (± 0.54) | 19.70 (± 0.34) | 27.32 (± 0.12)
| Homogeneous (largest) |  | 75.74 (± 0.42) | 60.89 (± 0.60) |  | 84.48 (± 0.58) | 62.51 (± 0.20) | 29.79 (± 0.32)

###  4.2 Performance Comparison with Model-Homogeneous FL Methods

We also compare the global model accuracy of FedRolex with two model-
homogeneous cases where all the clients have the largest capacity model
($\boldsymbol{\beta}=\\{1\\}$) and the smallest capacity model
($\boldsymbol{\beta}=\\{\nicefrac{{1}}{{16}}\\}$), representing the upper and
lower-bound performance, respectively.

Evaluation Results: As listed in Table 3, compared with other PT-based
methods, FedRolex reduces the gap in global model accuracy between model-
heterogeneous and upper-bound model-homogeneous settings. In particular,
FedRolex is on par with the upper-bound model-homogeneous case for Stack
Overflow, whereas both HeteroFL and Federated Dropout perform even worse than
the model homogeneous case using the smallest model. This result indicates
that with FedRolex, we will not be constrained to only using high-end devices
to achieve competitive global model accuracy. 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.

###  4.3 Impact of Client Model Heterogeneity Distribution

In our previous experiments, the distributions of model capacities across
client devices are set to be uniform. In this experiment, we aim to understand
the impact of the client model heterogeneity distribution. To do so, without
loss of generality, we use two client model capacities
$\boldsymbol{\beta}=\\{1,\nicefrac{{1}}{{16}}\\}$ and vary the distribution
ratio between the two (denoted as $\rho$) where $\rho=1$ represents the case
in which all the clients have the largest capacity model
($\boldsymbol{\beta}=\\{1\\}$) and $\rho=0$ represents the case in which all
the clients have the smallest capacity model
($\boldsymbol{\beta}=\\{\nicefrac{{1}}{{16}}\\}$).

Evaluation Results: Figure 3 shows how global model accuracy changes when
$\rho$ varies from $0$ to $1$ for CIFAR-10, CIFAR-100 and Stack Overflow. We
have three observations. (1) For CIFAR-10 (Figure 3(i)), there is a large gap
in global model accuracy between high and low data heterogeneity for a wide
range of $\rho$ (from 0.1 to 1). 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. This result indicates that
having more high-capacity models in the cohort has only limited contribution
to global model accuracy. (2) For the more challenging CIFAR-100 (Figure
3(ii)), the gap in global model accuracy is much lower between high and low
data heterogeneity. In contrast to CIFAR-10, the global model accuracy is
bottlenecked by the highest capacity of the models rather than the level of
data heterogeneity. (3) 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.

|  |
---|---|---
(i) |  (ii) |  (iii)

Figure 3:  Impact of client model heterogeneity distribution on global model
accuracy for (i) CIFAR-10, (ii) CIFAR-100, and (iii) Stack Overflow.

###  4.4 Performance on Training Larger Server Model

Similar to Federated Dropout, one advantage of FedRolex over static sub-model
extraction methods (HeteroFL and FjORD) is that FedRolex is able to train a
global server model that is larger than the largest client model. In this
experiment, we aim to evaluate the performance of FedRolex on training larger
server models. To do so, we consider the case where the size of the global
server model is $\boldsymbol{\gamma}=\\{2,4,8,16\\}$ times the size of client
models. For simplicity, all client models have the same size.

Evaluation Results: Figure 4(i) and Figure 4(ii) compare FedRolex with
Federated Dropout in terms of global model accuracy when $\boldsymbol{\gamma}$
for CIFAR-10 and CIFAR-100, respectively. As shown, although the global model
accuracy drops for both FedRolex and Federated Dropout when $\gamma$
increases, especially from $1$ to $4$, FedRolex consistently achieves higher
global model accuracy than Federated Dropout across
$\boldsymbol{\gamma}=\\{2,4,8,16\\}$ under both low and high data
heterogeneity. For Stack Overflow (Figure 4(iii)), the global model accuracy
has a much smaller drop when $\gamma$ increases. This demonstrates the
superiority of using large models on large-scale datasets for training larger
server models.

|  |
---|---|---
(i) |  (ii) |  (iii)

Figure 4:  Performance on training larger server model when the server model
is $\gamma$ times the size of the client model for (i) CIFAR-10, (ii)
CIFAR-100, and (iii) Stack Overflow.

###  4.5 Enhance Inclusiveness of FL in Real-world Distribution

A primary vision of FedRolex is to enhance the inclusiveness of FL. To
demonstrate this, in this experiment, we use real-world household income
distribution to emulate real-world device distribution. Specifically, we
retrieve household income distribution information from Bureau [33]. We map
$\beta_{n}=\nicefrac{{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 $\beta_{n}$. Detailed mapping of this distribution
to the corresponding income distribution is provided in Figure 7 in Appendix
A.7.

Table 4: Performance of FedRolex under emulated real-world device
distribution.

Dataset | Method |  | High Data Heterogeneity |  | Low Data Heterogeneity
---|---|---|---|---|---
| Local Accuracy | Global Accuracy |  | Local Accuracy | Global Accuracy
CIFAR-10 | Homogeneous (smallest) |  | 85.90 (± 0.46) | 38.82 (± 0.88) |  | 66.02 (± 0.52) | 46.86 (± 0.54)
Homogeneous (largest) |  | 95.54 (± 0.26) | 75.74 (± 0.41) |  | 93.54 (± 0.44) | 84.48 (± 0.58)
FedRolex |  | 94.05 (± 1.01) | 63.17 (± 1.45) |  | 91.03 (± 0.36) | 80.14 ± 0.52)
CIFAR-100 | Homogeneous (smallest) |  | 34.51 (± 0.56) | 12.69 (± 0.50) |  | 33.22 (± 0.10) | 19.70 (± 0.34)
Homogeneous (largest) |  | 81.99 (± 0.78) | 60.89 (± 0.60) |  | 76.43 (± 0.54) | 62.51 (± 0.20)
FedRolex |  | 73.33 (± 0.96) | 45.78 (± 1.71) |  | 66.31 (± 0.34) | 48.44 (± 0.51)
Stack Overflow | Homogeneous (smallest) |  | 27.32 (± 0.12)
Homogeneous (largest) |  | 29.79 (± 0.32)
FedRolex |  | 29.55 (± 0.41)

Evaluation Results: Table 4 shows both the global and local model accuracies
of FedRolex for CIFAR-10 and CIFAR-100 as well as the global model accuracy on
Stack Overflow under the emulated real-world device distribution. Again, we
compare with two model-homogeneous cases where all clients have the smallest
and largest model capacities, representing lower and upper-bound accuracy,
respectively. We make two observations. (1) Looking at the global model
accuracy, FedRolex consistently outperforms the lower-bound model-homogeneous
case across CIFAR-10, CIFAR-100, and Stack Overflow. This result indicates
that FedRolex enhances the inclusiveness of FL and improves the accuracy of
the global model, which would otherwise not be able to achieve. (2) Looking at
the local model accuracy, FedRolex significantly outperforms the lower-bound
model-homogeneous case on CIFAR-10 and CIFAR-100 under both low and high data
heterogeneity. This result indicates that FedRolex effectively boosts the
performance of low-end devices, which would otherwise not benefit from FL. 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.

|  |  |
---|---|---|---
(i) |  (ii) |  (iii) |  (iv)

Figure 5:  Local model accuracy distribution of FedRolex (orange color) vs.
the smallest model-homogeneous case (blue color) for CIFAR-10 and CIFAR-100
under low and high data heterogeneity.

##  5 Conclusion

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. 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. We provided a theoretical
statistical analysis on its advantage over Federated Dropout. Our experimental
results show that FedRolex consistently outperforms state-of-the-art PT-based
methods across models and datasets at both small and large scales. Moreover,
we demonstrated its performance on an emulated real-world device distribution
and show FedRolex contributes to making FL more inclusive.

Limitations and Future works. In this work, we provided a statistical analysis
of FedRolex. A full convergence analysis of FedRolex is not trivial and left
for future work. In addition, the goal of this work is to train a global
server model using a federation of heterogeneous client models. 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.
We will pursue it as our future work.

##  6 Acknowledgement

We thank the reviewers for their helpful comments. This work was partially
supported by NSF PFI:BIC-1632051, CNS-1814551, DMS-2012439, and a Google
Computing Platform (GCP) grant.

## References

  * Kairouz et al. [2021] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al.  Advances and open problems in federated learning.  _Foundations and Trends® in Machine Learning_ , 14(1–2):1–210, 2021.
  * Wang et al. [2021] Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H. Brendan McMahan, Blaise Aguera y Arcas, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, Suhas Diggavi, Hubert Eichner, Advait Gadhikar, Zachary Garrett, Antonious M. Girgis, Filip Hanzely, Andrew Hard, Chaoyang He, Samuel Horvath, Zhouyuan Huo, Alex Ingerman, Martin Jaggi, Tara Javidi, Peter Kairouz, Satyen Kale, Sai Praneeth Karimireddy, Jakub Konecny, Sanmi Koyejo, Tian Li, Luyang Liu, Mehryar Mohri, Hang Qi, Sashank J. Reddi, Peter Richtarik, Karan Singhal, Virginia Smith, Mahdi Soltanolkotabi, Weikang Song, Ananda Theertha Suresh, Sebastian U. Stich, Ameet Talwalkar, Hongyi Wang, Blake Woodworth, Shanshan Wu, Felix X. Yu, Honglin Yuan, Manzil Zaheer, Mi Zhang, Tong Zhang, Chunxiang Zheng, Chen Zhu, and Wennan Zhu.  A field guide to federated optimization, 2021.
  * McMahan et al. [2017] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In _Artificial intelligence and statistics_ , pages 1273–1282. PMLR, 2017.
  * Li et al. [2020] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith.  Federated optimization in heterogeneous networks.  _Proceedings of Machine Learning and Systems_ , 2:429–450, 2020.
  * Karimireddy et al. [2020] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh.  Scaffold: Stochastic controlled averaging for federated learning.  In _International Conference on Machine Learning_ , pages 5132–5143. PMLR, 2020.
  * Chen and Chao [2020] Hong-You Chen and Wei-Lun Chao.  Fedbe: Making bayesian model ensemble applicable to federated learning.  _arXiv preprint arXiv:2009.01974_ , 2020.
  * Bommasani et al. [2021] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al.  On the opportunities and risks of foundation models.  _arXiv preprint arXiv:2108.07258_ , 2021.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * He et al. [2020] Chaoyang He, Murali Annavaram, and Salman Avestimehr.  Group knowledge transfer: Federated learning of large cnns at the edge.  _Advances in Neural Information Processing Systems_ , 33:14068–14080, 2020.
  * Lin et al. [2020] Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi.  Ensemble distillation for robust model fusion in federated learning.  _Advances in Neural Information Processing Systems_ , 33:2351–2363, 2020.
  * Itahara et al. [2020] Sohei Itahara, Takayuki Nishio, Yusuke Koda, Masahiro Morikura, and Koji Yamamoto.  Distillation-based semi-supervised federated learning for communication-efficient collaborative training with non-iid private data.  _arXiv preprint arXiv:2008.06180_ , 2020.
  * Cho et al. [2022] Yae Jee Cho, Andre Manoel, Gauri Joshi, Robert Sim, and Dimitrios Dimitriadis.  Heterogeneous ensemble knowledge transfer for training large models in federated learning.  _International Joint Conference on Artificial Intelligence (IJCAI)_ , 2022.
  * Caldas et al. [2018] Sebastian Caldas, Jakub Konečny, H Brendan McMahan, and Ameet Talwalkar.  Expanding the reach of federated learning by reducing client resource requirements.  _arXiv preprint arXiv:1812.07210_ , 2018.
  * Diao et al. [2020] Enmao Diao, Jie Ding, and Vahid Tarokh.  Heterofl: Computation and communication efficient federated learning for heterogeneous clients.  _arXiv preprint arXiv:2010.01264_ , 2020.
  * Horvath et al. [2021] Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos Venieris, and Nicholas Lane.  Fjord: Fair and accurate federated learning under heterogeneous targets with ordered dropout.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Bonawitz et al. [2016] Keith Bonawitz, Vladimir Ivanov, Ben Kreuter, Antonio Marcedone, H Brendan McMahan, Sarvar Patel, Daniel Ramage, Aaron Segal, and Karn Seth.  Practical secure aggregation for federated learning on user-held data.  _arXiv preprint arXiv:1611.04482_ , 2016.
  * Charles et al. [2022a] Zachary Charles, Kallista Bonawitz, Stanislav Chiknavaryan, Brendan McMahan, et al.  Federated select: A primitive for communication-and memory-efficient federated learning.  _arXiv preprint arXiv:2208.09432_ , 2022a.
  * Ro et al. [2022] Jae Hun Ro, Theresa Breiner, Lara McConnaughey, Mingqing Chen, Ananda Theertha Suresh, Shankar Kumar, and Rajiv Mathews.  Scaling language model size in cross-device federated learning.  _arXiv preprint arXiv:2204.09715_ , 2022.
  * Charles et al. [2022b] Zachary Charles, Kallista Bonawitz, Stanislav Chiknavaryan, Brendan McMahan, and Blaise Agüera y Arcas.  Federated select: A primitive for communication- and memory-efficient federated learning, 2022b.  URL https://arxiv.org/abs/2208.09432.
  * Wang et al. [2020] Hongyi Wang, Kartik Sreenivasan, Shashank Rajput, Harit Vishwakarma, Saurabh Agarwal, Jy-yong Sohn, Kangwook Lee, and Dimitris Papailiopoulos.  Attack of the tails: Yes, you really can backdoor federated learning.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 16070–16084. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/file/b8ffa41d4e492f0fad2f13e29e1762eb-Paper.pdf.
  * Yang et al. [2022] Tien-Ju Yang, Dhruv Guliani, Françoise Beaufays, and Giovanni Motta.  Partial variable training for efficient on-device federated learning.  In _ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 4348–4352. IEEE, 2022.
  * Hinton et al. [2015] Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al.  Distilling the knowledge in a neural network.  _arXiv preprint arXiv:1503.02531_ , 2(7), 2015.
  * Stanton et al. [2021] Samuel Don Stanton, Pavel Izmailov, Polina Kirichenko, Alexander A Alemi, and Andrew Gordon Wilson.  Does knowledge distillation really work?  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=7J-fKoXiReA.
  * Srivastava et al. [2014] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov.  Dropout: a simple way to prevent neural networks from overfitting.  _The journal of machine learning research_ , 15(1):1929–1958, 2014.
  * Cheng et al. [2022] Gary Cheng, Zachary Charles, Zachary Garrett, and Keith Rush.  Does federated dropout actually work?  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 3387–3395, 2022.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * Krizhevsky [2009] Alex Krizhevsky.  Learning multiple layers of features from tiny images.  Technical report, University of Toronto, 2009.
  * Andreux et al. [2020] Mathieu Andreux, Jean Ogier du Terrail, Constance Beguier, and Eric W Tramel.  Siloed federated learning for multi-centric histopathology datasets.  In _Domain Adaptation and Representation Transfer, and Distributed and Collaborative Learning_ , pages 129–139. Springer, 2020.
  * TFF [2019] TFF.  Tensorflow federated stack overflow dataset.  _Online: https://www. tensorflow. org/federated/api_docs/python/tff/simulation/datasets/stackoverflow_ , 2019.
  * Zoph et al. [2020] Barret Zoph, Ekin D Cubuk, Golnaz Ghiasi, Tsung-Yi Lin, Jonathon Shlens, and Quoc V Le.  Learning data augmentation strategies for object detection.  In _European conference on computer vision_ , pages 566–583. Springer, 2020.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, _Advances in Neural Information Processing Systems 32_ , pages 8024–8035. Curran Associates, Inc., 2019.  URL http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf.
  * Moritz et al. [2018] Philipp Moritz, Robert Nishihara, Stephanie Wang, Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang, William Paul, Michael I Jordan, et al.  Ray: A distributed framework for emerging $\\{$AI$\\}$ applications.  In _13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)_ , pages 561–577, 2018.
  * Bureau [2021] U.S. Census Bureau.  Percentage distribution of household income in the u.s. in 2020.  In Statista, September 2021.  Retrieved May 18, 2022, from https://www.statista.com/statistics/203183/percentage-distribution-of-household-income-in-the-us.
  * Ross [2014] Sheldon M. Ross.  _Introduction to Probability Models_.  Academic Press, San Diego, CA, USA, 11th edition, 2014.
  * Newman [1960] Donald J Newman.  The double dixie cup problem.  _The American Mathematical Monthly_ , 67(1):58–61, 1960.
  * Arivazhagan et al. [2019] Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary.  Federated learning with personalization layers.  _arXiv preprint arXiv:1912.00818_ , 2019.

##  Appendix A Appendix

###  A.1 Statistical Analysis

######  Lemma 1.

Given $I$ indices, and one index is chosen at each round equally randomly. The
expected number of rounds of choosing all indices at least once is

$\displaystyle I\left({1\over I}+{1\over I-1}+\dots+{1\over 1}\right),$

which is the same as

$\displaystyle I\int_{0}^{\infty}\left(1-(1-e^{-t})^{I}\right)dt.$

###### Proof.

We denote the expected number of rounds to choose exactly $i$ indices at least
once as $E(i)$. Then we have $E(1)=1$, because, after the first round, one
index is chosen. After the first round, the expected number of rounds to
choose a new index is ${I\over I-1}$, because one of the remaining $I-1$ out
of the total $I$ indices needs to be chosen. That is, $E(2)=E(1)+{I\over
I-1}$. Similarly, we have

$\displaystyle E(i)=E(i-1)+{I\over I+1-i},\qquad\forall i=2,\dots,I.$

Thus, we have

$\displaystyle E(I)=E(I-1)+I=E(I-2)+{I\over 2}+{I\over 1}=\dots=I\left({1\over
I}+{1\over I-1}+\dots+{1\over 1}\right).$

The lemma is proved. ∎

It shows that the expected number of rounds to choose all indices at least
once is $I\log(I)$ when $I\rightarrow\infty$. This proof can not be
generalized to the case for choosing all indices at least $m$ times for $m\geq
2$. Therefore, we provide alternative proof for it [34, Example 5.17].

###### Alternative proof of Lemma 1.

This proof considers picking the indices as Poisson processes. Assume that the
Poisson process to choose one index has a rate $\lambda=1$. Since the index is
chosen equally randomly, choosing the $j$th index also follows a Poisson
process with a rate $1/I$ for any $j$ [34, Proposition 5.2]. We let $X_{j}$ be
the time to choose the first index $j$, and

$\displaystyle X=\max_{1\leq j\leq I}X_{j}$ (9)

is the time all indices are chosen at least once. Since all $X_{j}$ are
independent with rate $1/I$, we have

$\displaystyle P\\{X<t\\}=$ $\displaystyle P\\{\max\limits_{1\leq j\leq
I}X_{j}<t\\}=P\\{X_{j}<t,\mbox{ for }j=1,\dots,I\\}$ $\displaystyle=$
$\displaystyle(1-e^{-t/I})^{I}.$

Therefore, we have

$\displaystyle\mathbb{E}[X]=\int_{0}^{\infty}P\\{x>t\\}dt=\int_{0}^{\infty}\left(1-(1-e^{-t/I})^{I}\right)dt$

We let $N$ be the number of rounds to choose all indices at least one, and
$T_{i}$ be the $i$th interarrival time of the Poisson process for choosing one
index. Then we have

$\displaystyle X=\sum_{i=1}^{N}T_{i},$

and $T_{i}$ are independent. Thus we have

$\displaystyle\mathbb{E}[X|N]=N\mathbb{E}[T_{i}]=N,$

and which gives

$\displaystyle\mathbb{E}[X]=\mathbb{E}\\{\mathbb{E}[X|N]\\}=\mathbb{E}[N].$

Thus we have

$\displaystyle\mathbb{E}[N]=\int_{0}^{\infty}\left(1-(1-e^{-t/I})^{I}\right)dt=I\int_{0}^{\infty}\left(1-(1-e^{-t})^{I}\right)dt.$

The lemma is proved. ∎

Next, we will present the lemma for choosing each index at least $m$ times.

######  Lemma 2.

Given $I$ indices, and one index is chosen at each round equally randomly. The
expected number of rounds of choosing all indices at least $m$ times is

$\displaystyle I\int_{0}^{\infty}\left(1-(1-S_{m}(t)e^{-t})^{I}\right)dt,$

where

$\displaystyle S_{m}(y):=1+y+{y^{2}\over
2!}+\dots+{y^{m-1}\over(m-1)!}=\sum_{l=0}^{m-1}{y^{l}\over l!}.$ (10)

###### Proof.

We consider picking the indices as Poisson processes again. Assume that the
Poisson process to choose one index has a rate $\lambda=1$. Since the index is
chosen equally randomly, choosing the $j$th index also follows a Poisson
process with a rate of $1/I$ for any $j$. We let $X_{j}$ be the time to choose
index $j$ for the $m$th time, and

$\displaystyle X=\max_{1\leq j\leq I}X_{j}$ (11)

is the time all indices are chosen at least $m$ times. Since all $X_{j}$ are
independent with rate $1/I$, we have

$\displaystyle P\\{X<t\\}=$ $\displaystyle P\\{\max\limits_{1\leq j\leq
I}X_{j}<t\\}=P\\{X_{j}<t,\mbox{ for }j=1,\dots,I\\}$ $\displaystyle=$
$\displaystyle(1-S_{m}(t/I)e^{-t/I})^{I}.$

Therefore, we have

$\displaystyle\mathbb{E}[X]=\int_{0}^{\infty}P\\{x>t\\}dt.$

We let $N$ be the number of rounds to choose all indices at least $m$ times,
and $T_{i}$ be the $i$th interarrival time of the Poisson process for choosing
one index. Then we have

$\displaystyle X=\sum_{i=1}^{N}T_{i},$

and $T_{i}$ are independent. Thus we have

$\displaystyle\mathbb{E}[X|N]=N\mathbb{E}[T_{i}]=N,$

and which gives

$\displaystyle\mathbb{E}[X]=\mathbb{E}\\{\mathbb{E}[X|N]\\}=\mathbb{E}[N].$

Thus we have

$\displaystyle\mathbb{E}[N]=\int_{0}^{\infty}\left(1-(1-S_{m}(t/I)e^{-t/I})^{I}\right)dt=I\int_{0}^{\infty}\left(1-(1-S_{m}(t)e^{-t})^{I}\right)dt.$

The lemma is proved. ∎

It shows that the expected number of rounds to choose all indices at least
once is $I\log(I)+I(m-1)\log\log I$ when $I\rightarrow\infty$ [35].

###  A.2 Formal Definition of Selective Aggregation Scheme

Formally speaking, let $\mathcal{M}\subset\mathcal{N}$ be the set of selected
clients from the client pool from which the server pulls model parameters at
round $j$. Let $\theta_{[i,k]}$ be the $k^{th}$ parameter of layer $i$ of the
global model and $\theta_{m,[i,k]}$ be the $k^{th}$ parameter of layer $i$ of
client $m$. We denote $\mathcal{M}_{k}\subset\mathcal{M}$ as the set of
clients updating the $k^{th}$ parameter. The model parameters are aggregated
as follows:

$\displaystyle\theta_{[i,k]}=\frac{1}{\sum_{\begin{subarray}{c}m\in\mathcal{M}_{k}\end{subarray}}p_{m}}\sum_{\begin{subarray}{c}m\in\mathcal{M}_{k}\end{subarray}}p_{m}\theta_{m,[i,k]},$
(12)

The client weight $p_{m}$ is assigned based on factors like the client model
capacity, the number of data points a client has, etc. Throughout the paper,
unless otherwise stated, the weight of all clients is assumed to be the same,
i.e, $p_{m}=1/N$.

###  A.3 Ablation Study: Impact of Different Weighing Schemes

[12] reported that weighting clients is important to improving model accuracy.
Therefore, we did an ablation study and evaluated three client weighting
schemes: (1) model size-based weighting scheme: client weight is proportional
to the number of kernels in the model; (2) model update-based weighting
scheme: client weight is proportional to the number of updates; and (3) hybrid
weighting scheme: client weight is proportional to both (1) model size and (2)
model update.

Table 5 lists the results. As shown, the performance of the three weighting
schemes is not significantly better than the non-weighting scheme. Therefore,
we used the non-weighting scheme in FedRolex.

Table 5: Impact of weighting schemes on model accuracy under high data
heterogeneity.

| Weighting Scheme | Local Model Accuracy | Global Model Accuracy
---|---|---|---
CIFAR-10 | Non-Weighting | 95.95 (±0.81) | 69.44 (±1.50)
Model Size-based | 95.98 (±0.67) | 69.09 (±1.42)
Model Update-based | 96.01 (±0.71) | 68.83 (±0.89)
Hybrid | 96.05 (±0.96) | 68.78 (±0.89)
CIFAR-100 | Non-Weighting | 81.58 (±0.59) | 56.57 (±0.15)
Model Size-based | 81.23 (±1.56) | 56.99 (±0.27)
Model Update-based | 81.23 (±1.07) | 56.63 (±0.36)
Hybrid | 81.49 (±1.07) | 56.71 (±0.20)

|
---|---
(i) | (ii)

Figure 6:  Impact of inter-round kernel overlap on global model accuracy under
low and high data heterogeneity for (i) CIFAR-10 and (ii) CIFAR-100.

###  A.4 Ablation Study: Impact of Overlapping Kernels

We also studied the impact of overlapping kernels between rounds using
ResNet-18 and CIFAR-10/CIFAR-100 as an example. Specifically, we extracted
sub-models using a rolling window that advances and loops over all the kernels
of each convolution layer in the global model in strides. Let the degree of
overlap between each stride of the rolling window be $r\in[0,1]$. In each
iteration, each convolution layer in the global model is advanced by
$1+\lfloor\beta_{n}\left(1-r\right)K_{i}\rfloor$ where
$\lfloor\;\cdot\;\rfloor$ is the floor function. In FedRolex, $r=1$, i.e., the
kernels are advanced by $1$ from one iteration to the next iteration.

Figure 6 shows the impact of different $r$ on global model accuracy. As shown,
the value of $r$ does have some influence on the global model accuracy, but
the impact is non-linear and inconsistent.

###  A.5 Ablation Study: Impact of Client Participation Rate

In our main paper, we followed prior arts [14, 24, 15, 36, 9] and used a 10%
client participation rate. To examine the effect of client participation rate,
we conducted experiments with both lower (5%) and higher (20%) client
participation rates using CIFAR-10 as an example for FedRolex, HeteroFL and
Federated Dropout.

The results are summarized in Table 6. As shown, FedRolex consistently
outperforms both Federated Dropout and HeteroFL across 5%, 10% and 20% client
participation rates.

Table 6: Performance of FedRolex, HeteroFL, and Federated Dropout under
different client participation rates.

|  | Client Participation Rate
---|---|---
|  | 5% | 10% | 20%
CIFAR-10 | HeteroFL | 48.43 (+/- 1.78) | 63.90 (+/-2.74) | 65.07 (+/- 2.17)
Federated Dropout | 42.06 (+/- 1.29) | 46.64 (+/-3.05) | 55.20 (+/- 4.64)
FedRolex | 57.90 (+/- 2.72) | 69.44 (+/-1.50) | 71.85 ( +/- 1.22)

###  A.6 Communication and Computation Costs of FedRolex

To calculate the communication cost, we use the average size of the models
sent by all the participating clients per round as the metric. To calculate
the computation overhead, we calculate the FLOPs and numbers of parameters in
the models of all the participating clients per round and take the average as
the metric. To put these metrics in context, we also calculate the upper and
lower bounds of the communication cost and computation overhead (i.e., all the
clients were using the same largest model and smallest model, respectively).

Table 7 lists the results. As shown, compared to the upper bound, FedRolex
significantly reduces the communication cost and computation overhead while
being able to achieve comparable model accuracy. Compared to the lower bound,
although FedRolex has higher communication cost and computation overhead, the
model accuracy achieved is much higher than the lower bound. These results
indicate that FedRolex is able to achieve comparable high model accuracy as
the upper bound with much less communication cost and computation overhead.

Table 7: Computation and communication costs of FedRolex compared to upper and
lower bounds represented by homogeneous settings with largest and smallest
models respectively.

| Homogeneous (largest) | FedRolex | Homogeneous (smallest)
---|---|---|---
Average Number of Parameters per Client (Million) | 11.1722 | 2.9781232 | 0.04451
Average FLOPs per Client (Million) | 557.656 | 149.048384 | 2.41318
Average Model Size per Client (MB) | 42.62 | 11.36 | 0.17

###  A.7 Experimental Setup Details

Experimental Setup Details for Table 3. The experimental setup for PT-based
methods is listed in Table 8. The experimental setup for model-homogeneous
baselines was slightly different from the PT-based methods and hence is listed
separately in Table 9.

Table 8: Experimental setup details of PT-based methods in Table 3 on
CIFAR-10, CIFAR-100 and Stack Overflow.

|  | CIFAR-10 | CIFAR-100 | Stack Overflow
---|---|---|---|---
Local Epoch |  | 1 | 1 | 1
Cohort SIze |  | 10 | 10 | 200
Batch Size |  | 10 | 24 | 24
Initial Learning Rate |  | 2.00E-04 | 1.00E-04 | 2.00E-04
Decay Schedule | High Data Heterogeneity | 800, 1500 | 1000, 1500 | 600, 800
Low Data Heterogeneity | 800, 1250 | 1000, 1500
Decay Factor |  | 0.1 | 0.1 | 0.1
Communication Rounds | High Data Heterogeneity | 2500 | 3500 | 1200
Low Data Heterogeneity | 2000 | 3500
Optimizer |  | SGD | SGD | SGD
Momentum |  | 0.9 | 0.9 | 0.9
Weight Decay |  | 5.00E-04 | 5.00E-04 | 5.00E-04

Table 9: Experimental setup details of model-homogeneous baselines in Table 3
on CIFAR-10 and CIFAR-100 and Stack Overflow.

|  | CIFAR-10 | CIFAR-100 | Stack Overflow
---|---|---|---|---
Local Epoch |  | 1 | 1 | 1
Cohort Size |  | 10 | 10 | 200
Batch Size |  | 10 | 24 | 24
Initial Learning Rate |  | 2.00E-04 | 1.00E-04 | 2.00E-04
Decay Schedule | High Data Heterogeneity | 500, 1000 | 1000, 1500 | 300
Low Data Heterogeneity | 500, 1000 | 1000, 1500
Decay Factor |  | 0.1 | 0.1 | 0.1
Communication Rounds$\sim$ | High Data Heterogeneity | 1250 | 3500 | 1000
Low Data Heterogeneity | 1500 | 3500
Optimizer |  | SGD | SGD | SGD
Momentum |  | 0.9 | 0.9 | 0.9
Weight Decay |  | 5.00E-04 | 5.00E-04 | 5.00E-04

Experimental Setup Details for Figure 3. The experimental setup details are
tabulated in Tables 10 and 11.

Table 10: Experimental setup for results shown in Figure 3. $\rho$ between
$0.0$ and $0.5$ in $0.1$ increments.

Dataset | $\rho$ |  | 0.0 | 0.1 | 0.2 | 0.3 | 0.4
---|---|---|---|---|---|---|---
CIFAR-10 |  High Heterogeneity |  | Decay
---
Schedule
| 500, 1000 | 500, 1000 | 500, 1000 | 700, 1200 | 700, 1200
| Communication
---
Rounds
| 1250 | 1250 | 1250 | 1500 | 1500
Low Heterogeneity |  | Decay
---
Schedule
| 500, 1000 | 500, 1000 | 500, 1000 | 700, 1200 | 700, 1200
| Communication
---
Rounds
| 1250 | 1250 | 1250 | 1500 | 1500
CIFAR-100 |  High Heterogeneity |  | Decay
---
Schedule
| 1000, 1500 | 1000, 1500 | 1000, 1500 | 1000, 1500 | 1000, 1500
| Communication
---
Rounds
| 2000 | 2000 | 2000 | 2000 | 2000
Low Heterogeneity |  | Decay
---
Schedule
| 1000, 1500 | 1000, 1500 | 1000, 1500 | 1000, 1500 | 1000, 1500
| Communication
---
Rounds
| 2000 | 2000 | 2000 | 2000 | 2000
Stack Overflow |  High Heterogeneity |  | Decay
---
Schedule
| 800 | 800 | 800 | 800 | 800
| Communication
---
Rounds
| 1500 | 1500 | 1500 | 1500 | 1500
Low Heterogeneity |  | Decay
---
Schedule
| 800 | 800 | 800 | 800 | 800
| Communication
---
Rounds
| 1500 | 1500 | 1500 | 1500 | 1500

Table 11: Experimental setup for results shown in Figure 3. $\rho$ between
$0.5$ and $1.0$ in $0.1$ increments.

| $\rho$ |  | 0.6 | 0.7 | 0.8 | 0.9 | 1.0
---|---|---|---|---|---|---|---
Dataset |
CIFAR-10 |  High Heterogeneity |  | Decay
---
Schedule
| 700, 1200 | 700, 1200 | 500, 1000 | 500, 1000 | 500, 1000
| Communication
---
Rounds
| 1500 | 1500 | 1250 | 1250 | 1250
Low Heterogeneity |  | Decay
---
Schedule
| 700, 1200 | 700, 1200 | 500, 1000 | 500, 1000 | 500, 1000
| Communication
---
Rounds
| 1500 | 1500 | 1250 | 1250 | 1250
CIFAR-100 |  High Heterogeneity |  | Decay
---
Schedule
| 1000, 1500 | 1000, 1500 | 1000, 1500 | 1000, 1500 | 1000, 1500
| Communication
---
Rounds
| 2000 | 2000 | 2000 | 2000 | 2000
Low Heterogeneity |  | Decay
---
Schedule
| 1000, 1500 | 1000, 1500 | 1000, 1500 | 1000, 1500 | 1000, 1500
| Communication
---
Rounds
| 2000 | 2000 | 2000 | 2000 | 2000
Stack Overflow |  High Heterogeneity |  | Decay
---
Schedule
| 800 | 800 | 800 | 800 | 800
| Communication
---
Rounds
| 1500 | 1500 | 1500 | 1500 | 1500
Low Heterogeneity |  | Decay
---
Schedule
| 800 | 800 | 800 | 800 | 800
| Communication
---
Rounds
| 1500 | 1500 | 1500 | 1500 | 1500

Experimental Setup Details for Figure 4. The experimental setup details are
tabulated in Table 12.

Table 12: Experimental setup for results shown in Figure 4

Dataset | $\gamma$ |  | 2 | 4 | 8 | 16
---|---|---|---|---|---|---
CIFAR-10 |  High Heterogeneity |  | Decay
---
Schedule
| 800, 1200 | 800, 1200 | 800, 1200 | 800, 1200
| Communication
---
Rounds
| 1500 | 1500 | 1500 | 1500
Low Heterogeneity |  | Decay
---
Schedule
| 800, 1200 | 800, 1200 | 800, 1200 | 800, 1200
| Communication
---
Rounds
| 1500 | 1500 | 1500 | 1500
CIFAR-100 |  High Heterogeneity |  | Decay
---
Schedule
| 800, 1200 | 800, 1200 | 800, 1200 | 800, 1200
| Communication
---
Rounds
| 1500 | 1500 | 1500 | 1500
Low Heterogeneity |  | Decay
---
Schedule
| 800, 1200 | 800, 1200 | 800, 1200 | 800, 1200
| Communication
---
Rounds
| 1500 | 1500 | 1500 | 1500
Stack Overflow |  High Heterogeneity |  | Decay
---
Schedule
| 800 | 800 | 800 | 800
| Communication
---
Rounds
| 1500 | 1500 | 1500 | 1500
Low Heterogeneity |  | Decay
---
Schedule
| 800 | 800 | 800 | 800
| Communication
---
Rounds
| 1500 | 1500 | 1500 | 1500

$1\times$$\nicefrac{{1}}{{2}}\times$$\nicefrac{{1}}{{4}}\times$$\nicefrac{{1}}{{8}}\times$$\nicefrac{{1}}{{16}}\times$ |  | Model Capacity | Annual Household Income
---|---
$\nicefrac{{1}}{{16}}\times$ | $<\$75,000$
$\nicefrac{{1}}{{8}}\times$ | $\$75,000-\$100,000$
$\nicefrac{{1}}{{4}}\times$ | $\$100,000-\$150,000$
$\nicefrac{{1}}{{2}}\times$ | $\$150,000-\$200,000$
$1\times$ | $>\$200,000$

Figure 7: Mapping between real-world annual household income and model
capacity. Table 13: Experimental setup for Table 4 for CIFAR-10, CIFAR-100 and
Stack Overflow.

|  | CIFAR-10 | CIFAR-100 | Stack Overflow
---|---|---|---|---
Local Epoch |  | 1 | 1 | 1
Cohort SIze |  | 10 | 10 | 200
Batch Size |  | 10 | 24 | 24
Initial Learning Rate |  | 2.00E-04 | 1.00E-04 | 2.00E-04
Decay Schedule | High Heterogeneity | 800, 1500 | 1000, 1500 | 600, 800
Low Heterogeneity | 800, 1250 | 1000, 1500
Decay Factor |  | 0.1 | 0.1 | 0.1
Communication Rounds | High Heterogeneity | 2500 | 3500 | 1200
Low Heterogeneity | 2000 | 3500
Optimizer |  | SGD | SGD | SGD
Momentum |  | 0.9 | 0.9 | 0.9
Weight Decay |  | 5.00E-04 | 5.00E-04 | 5.00E-04

###  A.8 Algorithm Pseudocodes

The pseudocodes for HeteroFL and Federated Dropout are given in Algorithms 2
and 3 respectively. Their differences from FedRolex are marked using blue
color.

1

2Initialization ; $\theta^{(0)}$, $\mathcal{N}$

Input : $D_{n}\;\beta_{n}\;\forall n\in\mathcal{N}$,

Output : $\theta^{J}$

3 Server Executes

4 for  _$j\leftarrow 0$ to $J-1$_ do

5 Sample subset $\mathcal{M}$ from $\mathcal{N}$

6 Broadcast
$\theta^{(j)}_{m,[i\;;\;{0,1,\;...\;\lfloor\beta_{n}K_{i}\rfloor-1}]}\forall
i\;\text{and}\;m\in\mathcal{M}$

7 for  _each client $m\in\mathcal{M}$_ do

8 clientStep(_$\theta^{(j)}_{m}$ , $D_{m}$_)

9 end for

10 Aggregate $\theta^{(j+1)}_{[i,k]}$ according to Equation (12)

11

12 end for

13 Subroutine  _clientStep(_$\theta^{(j)}_{n}$ , $D_{n}$_)_

14 $m_{n}\longleftarrow len(D_{n})$

15 for  _$k\leftarrow 0$ to $m_{n}$_ do

16 $\theta_{n}\longleftarrow\theta_{n}-\eta\nabla l(\theta_{n};d_{n,k})$

17 end for

18 return $\theta_{n}$

19

Algorithm 2 HeteroFL

1

2Initialization ; $\theta^{(0)}$, $\mathcal{N}$

Input : $D_{n}\;\beta_{n}\;\forall n\in\mathcal{N}$,

Output : $\theta^{J}$

3 Server Executes

4 for  _$j\leftarrow 0$ to $J-1$_ do

5 Sample subset $\mathcal{M}$ from $\mathcal{N}$

6 Broadcast
$\theta^{(j)}_{m,[i\;;\;{k_{1},...,k_{\lfloor\beta_{n}K_{i}\rfloor}}]}\forall
i\;\text{and}\;m\in\mathcal{M}$

7 for  _each client $m\in\mathcal{M}$_ do

8 clientStep(_$\theta^{(j)}_{m}$ , $D_{m}$_)

9 end for

10 Aggregate $\theta^{(j+1)}_{[i,k]}$ according to Equation (12)

11

12 end for

13 Subroutine  _clientStep(_$\theta^{(j)}_{n}$ , $D_{n}$_)_

14 $m_{n}\longleftarrow len(D_{n})$

15 for  _$k\leftarrow 0$ to $m_{n}$_ do

16 $\theta_{n}\longleftarrow\theta_{n}-\eta\nabla l(\theta_{n};d_{n,k})$

17 end for

18 return $\theta_{n}$

19

Algorithm 3 Federated Dropout