# Personalized Federated Learning with Multiple Known Clusters

Boxiang Lyu Booth School of Business, The University of Chicago, Chicago, IL.
<EMAIL_ADDRESS>Filip Hanzely Toyota Technological Institute at
Chicago, Chicago, IL<EMAIL_ADDRESS>Mladen Kolar Booth School of
Business, The University of Chicago, Chicago, IL.
<EMAIL_ADDRESS>
(First draft: April 25th, 2022)

###### Abstract

We consider the problem of personalized federated learning when there are
known cluster structures within users. An intuitive approach would be to
regularize the parameters so that users in the same cluster share similar
model weights. The distances between the clusters can then be regularized to
reflect the similarity between different clusters of users. We develop an
algorithm that allows each cluster to communicate independently and derive the
convergence results. We study a hierarchical linear model to theoretically
demonstrate that our approach outperforms agents learning independently and
agents learning a single shared weight. Finally, we demonstrate the advantages
of our approach using both simulated and real-world data.

Keywords: personalized federated learning, multi-task learning, distributed
optimization

##  1 Introduction

Smart phones, voice assistants, and wearable devices are everywhere in our
lives, constantly collecting data on our behavior and habits. Federated
learning (FL) is a recently introduced framework developed to use this rich
source of data, while minimizing the intrusion of clients’ privacy. Although
traditional machine learning methods often require the aggregation of client
data in a central server, federated learning avoids such a requirement,
allowing models to be trained with mostly local computation and occasional
server-wide communication rounds (McMahan et al., 2017; Yang et al., 2019;
Bonawitz et al., 2019; Li et al., 2020; Kairouz et al., 2019).

Data sets used in FL tasks are heterogeneous in nature, as they are collected
from clients who are heterogeneous in nature. Tailoring the learned model to
each client through personalized FL has garnered a tremendous amount of
interest in recent years (Hanzely et al., 2020a; Fallah et al., 2020; Deng et
al., 2020; Dinh et al., 2020; Mansour et al., 2020). A central theme is
learning a single global model and a personalized model for each client
simultaneously in the training process (Hanzely and Richtárik, 2020; Dinh et
al., 2020; Li et al., 2021; Fallah et al., 2020). While it is typical to
assume that clients can be grouped into a single cluster in existing
literature, in disciplines such as education, psychology, or economics, it is
often assumed that clients can be grouped into multiple clusters using known
information. Hierarchical linear models are frequently used to model data with
such a cluster structure Raudenbush and Bryk (1986); Raudenbush (1988); Bryk
and Raudenbush (1987, 1992); Hofmann (1997); Stephen and Anthony (2002).

Motivated by the prevalence of hierarchical models in social sciences, we
propose a hierarchical, multi-cluster approach to personalization in federated
learning. In particular, we leverage the known hierarchical structure and
simultaneously learn (1) a global model for all clients, (2) a cluster-
specific model for each client cluster, and (3) a personalized model for each
client. More specifically, we develop a loopless algorithm for fitting the
hierarchical linear model and derive its convergence rates and optimal
parameters. Our algorithm allows each client cluster to communicate only
within the group and allows each cluster to determine when to aggregate local
updates independently of other clusters. Such an independence could improve
convergence when the client clusters align with the communication graph.
Existing empirical research has shown that extra communication rounds within
different geographical clusters can improve convergence rates (Huang et al.,
2019; Briggs et al., 2020) and the proposed algorithm can leverage this
property. In the finite-sum setting, we further develop an accelerated,
variance reduced, stochastic variant of the algorithm, which has been shown to
be minimax optimal in terms of communication rounds and oracle calls in the
single cluster setting (Hanzely et al., 2020a).

Some existing work studied federated learning under the assumption that
clients can be grouped into different clusters without assuming that these
clusters are known. Our work emphasizes the optimality of the proposed
approach when clusters are given a priori. In particular, the prior literature
mostly focuses on the convergence rates of the optimization procedure and
identifies the conditions under which the underlying cluster structures can be
recovered. A key problem remains unaddressed: even assuming the cluster
structure is known, can we prove that these structures can improve upon simple
baselines such as training using only local data or training a single model?
More importantly, can we show the optimality of the clustered formulation even
in the simple case where the cluster labels are given?

To answer these questions, we establish statistical properties for the
estimator obtained in our framework. We show that under a simple problem
setting, our approach recovers the best linear unbiased estimator of the
clients’ local weights, dominating both training a single model for all
clients and training a unique model for each client independently. This result
complements Li et al. (2021) that studied the single-group setting. We further
show that the unbiased restriction cannot be weakened. Even in the single
cluster setting, we can construct a counterexample where a James-Stein
estimator outperforms the proposed method. Our result also complements Chen et
al. (2021) by identifying a regime in which a personalized approach is at
least as good as the recommended alternatives, noting that structural
information can be used to develop personalization schemes that outperform the
proposed alternatives.

Finally, we demonstrate the empirical effectiveness of our model on the DMEF
Customer Lifetime Value data set Blattberg et al. (2009) and compare it with a
recently proposed method in targeted marketing Bumbaca et al. (2020),
developing a personalized marketing model using data from a leading non-profit
organization in the United States Blattberg et al. (2009).

###  1.1 Related work

Our work is related to the literature on personalized federated learning,
distributed multi-task learning, and Bayesian hierarchical models.

There has been a lot of focus on personalization in federated learning over
the past few years. Karimireddy et al. (2020) is one of the first works to
discuss heterogeneity among clients. Fallah et al. (2020) adapted model-
agnostic meta-learning (MAML) algorithms for personalization. Deng et al.
(2020) developed a method that interpolates between client-specific parameters
and global parameters. Mansour et al. (2020) suggests three different
approaches for personalized federated learning and provides learning theoretic
guarantees. Li et al. (2021); Hanzely and Richtárik (2020); Dinh et al. (2020)
studied a personalization approach where each client has a parameter whose
distance to the average of the parameters is regularized either explicitly or
implicitly and developed different optimization techniques. Li et al. (2018)
studied optimizing a similar regularized loss in a non-personalized setting,
where a single model is trained for all clients. More specifically, Li et al.
(2021) provided some theoretical justification on the robustness and fairness
of the procedure in the single cluster regime, Hanzely et al. (2021) provided
a unified analysis of different optimization techniques, and Dinh et al.
(2020) studied the optimization problem assuming that individual clients can
exactly evaluate a proximal operator. These approaches can be viewed as a
special instance of distributed multi-task learning with graph regularization
Wang et al. (2018).

A different line of work on personalization assumes that client-specific
parameters may be drawn from an unknown mixture distribution and
simultaneously group clients and learn model parameters using a single
algorithm (Mansour et al., 2020; Ghosh et al., 2020; Sattler et al., 2020;
Smith et al., 2017; Briggs et al., 2020; Huang et al., 2019). Similar
approaches have been studied in multi-task learning (Kumar and Daume III,
2012; Jacob et al., 2008; Zhang and Yang, 2017; Zhang and Yeung, 2012; Zhou et
al., 2011a, b; Bakker and Heskes, 2003). While these approaches focus on the
setting where the structure and parameters of the cluster are learned
simultaneously, in many domains, the clusters are known a priori and given,
for example, by known covariates such as age, gender, and geographical
location (Raudenbush and Bryk, 1986; Raudenbush, 1988; Bryk and Raudenbush,
1987, 1992; Hofmann, 1997; Stephen and Anthony, 2002; Lee and Nelder, 1996;
Daniels and Gatsonis, 1999). Such a structure can be used to develop
hierarchical models with improved personalization without having to separately
cluster the clients. The data example that we investigate in detail in Section
6 comes from marketing, where hierarchical models have a long history (Naik
and Peters, 2009; Bumbaca et al., 2017, 2020; Hooley et al., 1999; French and
Russell-Bennett, 2015). Given the ubiquity of mobile devices, combining these
marketing models with federated learning could better help the industry
implement state-of-the-art marketing research.

Our optimization procedure is related to the methods used to optimize the
objectives commonly found in hierarchical federated learning (Abad et al.,
2020; Wang et al., 2020; Liu et al., 2020; Briggs et al., 2020; Wainakh et
al., 2020). However, while the existing literature focuses on finding a single
model for all clients, our optimization procedure learns personalized models
for each client and each cluster and also learns a joint global model. The
design of the optimization algorithm is related to loopless procedures in
distributed optimization (Zhao et al., 2021; Li, 2021; Qian et al., 2021a).
These procedures remove the inner loops, thereby simplifying the algorithm. In
single-machine settings, such simplifications have been shown to outperform
their loopy counterparts (Kovalev et al., 2020).

Concurrent to our work, Marfoq et al. (2021) studies federated multi-task
learning under a mixture of distributions, focusing on nonasymptotic
convergence rates. Our work further shows the optimality of our approach in
terms of generalization error in addition to convergence analysis.
Additionally, instead of analyzing an expectation-maximization-inspired
approach, our work uses a loopless gradient-based algorithm that has been
shown to enjoy optimal communication complexity in the single cluster regime.
Furthermore, Duan and Wang (2022) discusses a general framework for multi-task
learning and shows that the approach can be adapted to various concepts of
task relatedness. Although the loss function discussed here is similar, we
further develop a federated optimization procedure and characterize the bounds
on the communication and computation complexity of the federated learning
algorithm used to minimize the loss.

###  1.2 Notation

For any vector $v\in\mathbb{R}^{d}$, we use $\|v\|$ to denote its $\ell_{2}$
norm. For any finite set $A$, we use $|A|$ to denote its cardinality and
$A[i]$ to denote the $i$-th element in $A$ according to some arbitrary order.
The $d$-dimensional identity matrix is denoted as $I_{d}\in\mathbb{R}^{d\times
d}$. The Kronecker product between two conforming matrices $B,C$ is denoted as
$B\otimes C$.

##  2 Model Formulation

Suppose that there are $n$ clients with their individual data sets for whom we
would like to fit personalized models. Furthermore, suppose that these clients
are divided into $k$ known clusters. Generally speaking, if two clients belong
to the same cluster, then we expect their models to be more similar than if
the clients belong to different clusters. Let $\mathcal{I}_{j}$,
$j=1,\ldots,k$, be the set of clients belonging to the cluster $j$. We use
$f_{i}(\cdot)$ to denote the loss function for client $i=1,\ldots,n$.
Throughout the paper, we assume that the loss function $f_{i}(\cdot)$ is
strongly convex and smooth. In particular, we make the following assumption.

######  Assumption 2.1.

The loss function $f_{i}$ is $\mu$-strongly convex and $L$-smooth.

We use “local”, “cluster”, and “global“ to denote variables, functions, and
values associated with individual clients, different clusters, and the entire
network, respectively.

Let $\theta_{i}\in\mathbb{R}^{d}$ denote the model parameter for client $i$.
We focus on minimizing the following objective function:

$\min_{\\{\theta_{i}\\}_{i=1}^{n}}F(\\{\theta_{i}\\}_{i=1}^{n})\coloneq\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\left(f_{i}(\theta_{i})+\frac{(1-\alpha_{j})\gamma_{i}}{2}\left\|\theta_{i}-\bar{\theta}_{j}\right\|^{2}+\frac{\alpha_{j}\gamma_{i}}{2}\left\|\theta_{i}-\bar{\theta}\right\|^{2}\right),$
(2.1)

where $\alpha_{j}\in\mathbb{R}_{\geq 0}$, $j=1,\ldots,k$, are the tuning
parameters that control the regularization strength in each cluster,
$\gamma_{i}\in\mathbb{R}_{\geq 0}$, $i=1,\ldots,n$, are the tuning parameters
specific to each client, and $\\{\bar{\theta}_{j}\\}_{j=1}^{n},\bar{\theta}$
denote the weight averages of the parameters in the cluster $j$ and the entire
network, respectively. That is,

$\bar{\theta}_{j}=\frac{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\theta_{i}}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}},\quad
j=1,\ldots,k,\quad\text{and}\quad\bar{\theta}=\frac{\sum_{j=1}^{k}\alpha_{j}\bar{\theta}_{j}}{\sum_{j=1}^{k}\alpha_{j}}=\frac{\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}\theta_{i}}{\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}}.$
(2.2)

The objective is comprised of three terms: the client-specific loss and two
regularization terms, summed over all clients. The first regularizer penalizes
the distance between the local parameter $\theta_{i}$ and the cluster averages
$\bar{\theta}_{j}$, while the second regularizer penalizes the distance
between the local parameter and the global average $\bar{\theta}$. By changing
the cluster-specific parameter $\alpha_{j}$ from $0$ to $1$ we can interpolate
between the two regimes: when $\alpha_{j}=0$ for all $j$, we train $k$
personalized models—one for each cluster, independently of other clusters;
when $\alpha_{j}=1$ for all $j$, we train the single-cluster model studied in
Hanzely and Richtárik (2020); Dinh et al. (2020); Li et al. (2021).

The objective in Eq. 2.1 is different from commonly used objective functions
in multi-task learning (Wang et al., 2018; Zhou et al., 2011a, b; Jacob et
al., 2008). When the cluster structure is known, the multi-task learning
objective can be written as

$\begin{split}\min_{\\{\theta_{i}\\}_{i=1}^{n},\\{w_{j}\\}_{j=1}^{k},\bar{w}}&F_{MTL}(\\{\theta_{i}\\}_{i=1}^{n},\\{w_{j}\\}_{j=1}^{k},\bar{w})\\\
&\quad\coloneq\sum_{j=1}^{k}\left(\sum_{i\in\mathcal{I}_{j}}\left(f_{i}(\theta_{i})+\frac{\gamma_{i}}{2}\|\theta_{i}-w_{j}\|^{2}\right)+\frac{\lambda_{j}}{2}\|w_{j}-\bar{w}\|^{2}\right),\end{split}$
(2.3)

where the penalty parameter $\gamma_{i}$ regularizes the distance between the
local parameter and the cluster parameter $w_{j}$, while the penalty parameter
$\lambda_{j}$, $j=1,\ldots,k$, regularizes the distance between the cluster
parameter and the global parameter $\bar{w}$. Furthermore, in (2.3) we
optimize both the local weights $\\{\theta_{i}\\}_{i=1}^{n}$ and the average
local parameters, $\\{w_{j}\\}_{j=1}^{k}$ and $\bar{w}$. The two forms,
however, have same stationary points.

######  Proposition 2.2.

Suppose that 2.1 holds for loss functions $\\{f_{i}\\}_{i=1}^{n}$. Let
$\\{\lambda_{j}\\}_{j=1}^{k},\\{\gamma_{i}\\}_{i=1}^{n}$ be any set of tuning
parameters for Eq. 2.3. Let $\\{\widehat{w}_{j}\\}_{j=1}^{k}$,
$\widehat{\bar{w}}$, and $\\{\widehat{\theta}_{i}^{\prime}\\}_{i=1}^{n}$ be
the minimizers of (2.3). Fix

$\alpha_{j}=\frac{\lambda_{j}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}\gamma_{i}},\quad
j=1,\ldots,k,$

and let $\\{\widehat{\theta}_{i}\\}_{i=1}^{n}$ denote the unique minimizer of
Eq. 2.1. We then have $\widehat{\theta}_{i}=\widehat{\theta}_{i}^{\prime}$ for
all $i\in[n]$.

###### Proof.

See Section A.1. ∎

As we will demonstrate in the sequel, while Eqs. 2.1 and 2.3 have the same
stationary points, optimizing the former does not require us to keep track of
cluster or network averages, allowing these parameters to be calculated on the
fly. This feature of Eq. 2.1 better suits the federated learning setting,
removing the server’s need to keep track of additional variables during the
optimization process.

To further motivate the objective in Eq. 2.1, we show that the maximum
likelihood estimate of all parameters in a hierarchical generalized linear
model can be viewed as the minimizer of Eq. 2.1 (Lee and Nelder, 1996; Stephen
and Anthony, 2002; Bryk and Raudenbush, 1992). Specifically, let

$\displaystyle\bar{\theta}_{j}^{*}\sim\mathcal{N}\left(\bar{\theta}^{*},\sigma_{\text{center}}^{2}I_{d}\right),\quad
j=1,\ldots,k,$

where $\mathcal{N}(\cdot,\cdot)$ denotes a Gaussian distribution and
$\bar{\theta}^{*}\in\mathbb{R}^{d}$. Furthermore, assume that for
$j=1,\ldots,k$ and $i\in\mathcal{I}_{j}$, we have

$\displaystyle\theta_{i}^{*}$
$\displaystyle\sim\mathcal{N}\left(\bar{\theta}^{*}_{j},\sigma_{j}^{2}I_{d}\right),$
$\displaystyle y_{i}\mid\theta_{i}^{*},X_{i}$ $\displaystyle\sim
p_{Y}(y_{i};u^{-1}(X_{i}^{T}\theta_{i}^{*}),\tau),$

where $X_{i}\in\mathbb{R}^{n_{i}\times d}$ is a matrix of observations,
$y_{i}\in\mathbb{R}^{n_{i}}$ is the response vector, $u$ is a known link
function, $p_{Y}(\cdot;\cdot,\cdot)$ is the probability density function of an
exponential family distribution. The negative log-likelihood for this model is

$\begin{split}&\ell(\\{\theta_{i}\\}_{i=1}^{n},\\{\bar{\theta}_{j}\\}_{j=1}^{k},\bar{\theta};\\{X_{i},y_{i}\\}_{i=1}^{n},\tau)=\\\
&\quad\sum_{j=1}^{k}\Biggl{(}\sum_{i\in\mathcal{I}_{j}}\left(\log
p_{Y}(y_{i};u^{-1}(X^{T}_{i}\theta_{i}^{*}),\tau)+\frac{1}{2\sigma_{\text{cluster
}j}^{2}}\left\|\theta_{i}-\bar{\theta}_{j}\right\|^{2}\right)\\\ &\hskip
200.0003pt+\frac{1}{2\sigma_{\text{center}}^{2}}\left\|\bar{\theta}_{j}-\bar{\theta}\right\|^{2}\Biggr{)},\end{split}$

which is a special instance of Eq. 2.3. Then, by Proposition 2.2, minimizing
Eq. 2.1 with suitable tuning parameters corresponds to the maximum likelihood
estimation.

##  3 Algorithm and Convergence Analysis

We develop a loopless SGD-style algorithm to minimize the objective in Eq.
2.1. In each iteration, the algorithm randomly decides between descending on
the loss functions $f_{i}$’s, descending on the regularizer controlling the
distance between the local weights and the centers of the local clusters, and
descending on both regularizers. The first kind of descent step can be done in
parallel on all machines, the same as the local computation step in local SGD
(Stich, 2018). The second kind requires clients within the same cluster to
communicate with one another, but does not require between-cluster
communication. Only when we simultaneously perform gradient descent on both
the within-cluster and between-cluster regularizers do we communicate across
clusters.

###  3.1 Technical Preliminaries

We start by introducing an additional notation. Recall that $i=1,\ldots,n$ is
used to index clients, $j=1,\dots,k$ is used to index clusters,
$\mathcal{I}_{j}$ is the set of clients belonging to the cluster $j$, with
cardinality $|\mathcal{I}_{j}|$. Let
$\Theta_{j}\in\mathbb{R}^{|\mathcal{I}_{j}|d}$ be the weight vector formed by
stacking the weights of the clients in the cluster $j$ and
$\Theta\in\mathbb{R}^{nd}$ be the weight vector formed by stacking the weights
of all clients; that is,

$\bm{\theta}_{j}=(\theta_{\mathcal{I}_{j}[1]},\ldots,\theta_{\mathcal{I}_{j}[|\mathcal{I}_{j}|]})\in\mathbb{R}^{|\mathcal{I}_{j}|d}\quad\text{and}\quad\bm{\theta}=(\theta_{1},\ldots,\theta_{n})\in\mathbb{R}^{nd}.$
(3.1)

Without loss of generality, assume that clients are ordered according to the
cluster to which they belong, that is,
$\mathcal{I}_{1}=\\{1,\ldots,|\mathcal{I}_{1}|\\}$ and
$\mathcal{I}_{k}=\\{|\mathcal{I}_{k-1}|+1,\ldots,|\mathcal{I}_{k-1}|+|\mathcal{I}_{k}|\\}$.
We may then write the cluster-specific and global regularizers as

$\displaystyle\psi_{j}\left(\bm{\theta}_{j};\left\\{\gamma_{i}\right\\}_{i\in\mathcal{I}_{j}}\right)=\frac{1}{2}\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\left\|\theta_{i}-\bar{\theta}_{j}\right\|^{2},\qquad
j=1,\ldots,k,$
$\displaystyle\varphi\left(\bm{\theta};\left\\{\gamma_{i}\right\\}_{i=1}^{n},\left\\{\alpha\right\\}_{j=1}^{k}\right)=\frac{1}{2}\sum_{j=1}^{k}\alpha_{j}\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\left\|\theta_{i}-\bar{\theta}\right\|^{2}.$

The loss function in Eq. 2.1 can now be rewritten as

$\min_{\bm{\theta}}F(\bm{\theta})=\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}f_{i}(\theta_{i})+\sum_{j=1}^{k}(1-\alpha_{j})\psi_{j}\left(\bm{\theta}_{j};\left\\{\gamma_{i}\right\\}_{i\in\mathcal{I}_{j}}\right)+\varphi\left(\bm{\theta};\left\\{\gamma_{i}\right\\}_{i=1}^{n},\left\\{\alpha\right\\}_{j=1}^{k}\right).$

We explicitly calculate the gradients of the cluster-specific regularizers and
the global regularizer in the following proposition.

######  Proposition 3.1.

We have that

$\displaystyle\nabla_{\theta_{i}}\psi_{j}\left(\bm{\theta}_{j};\left\\{\gamma_{i}\right\\}_{i\in\mathcal{I}_{j}}\right)=\gamma_{i}\left(\theta_{i}-\bar{\theta}_{j}\right),\qquad
i\in\mathcal{I}_{j},j=1,\ldots,k;$
$\displaystyle\nabla_{\theta_{i}}\varphi\left(\bm{\theta};\left\\{\gamma_{i}\right\\}_{i=1}^{n},\left\\{\alpha\right\\}_{j=1}^{k}\right)=\alpha_{j}\gamma_{i}\left(\theta_{i}-\bar{\theta}\right),\qquad
i\in\mathcal{I}_{j},j=1,\ldots,k.$

###### Proof.

See Section A.2. ∎

###  3.2 Asynchronous Loopless Local Gradient Descent (Async-L2GD)

We formally introduce our main algorithm, Async-L2GD, in this section.
Algorithm 1 provides the pseudocode. We present an accelerated, variance
reduced, and stochastic variant of the algorithm for the finite-sum setting in
Section 4.

Algorithm 1 Async-L2GD

1: Input: $\theta_{1}^{0}=\dots=\theta_{n}^{0}=\bm{0}_{d}\in\mathbb{R}^{d}$,
step size $\eta>0$, probabilities $p_{1},\dots,p_{k},p_{0}\in[0,1]$, fractions
$\tau_{1},\dots,\tau_{k}\in[0,1]$.

2: for $t=1,2,\dots$ do

3: $\xi_{0}=1$ with probability $p_{0}$ and 0 with probability $1-p_{0}$.

4: if $\xi_{0}=1$ then

5: All Clusters $j=1,\dots,k$ compute cluster average
$\bar{\theta}_{j}^{t}=\frac{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\theta_{i}^{t}}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}.$

6: All Clusters communicate with one another and calculate network average
$\bar{\theta}^{t}=\frac{\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}\theta_{i}^{t}}{\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}}$.

7: All Clusters $j=1,\dots,k$ compute step towards both cluster average and
network average for all clients $i\in\mathcal{I}_{j}$
$\theta_{i}^{t+1}=\left(1-\frac{\eta\gamma_{i}(\alpha_{j}+\tau_{j}(1-\alpha_{j}))}{p_{0}}\right)\theta_{i}^{t}+\frac{\eta\gamma_{i}}{p_{0}}(\alpha_{j}\bar{\theta}^{t}+\tau_{j}(1-\alpha_{j})\bar{\theta}_{j}^{t})$.

8: else

9: for All Clusters $j=1,\dots,k$ in parallel do

10: $\xi_{j}=1$ with probability $p_{j}$ and 0 with probability $1-p_{j}$.

11: if $\xi_{j}=1$ then

12: Compute cluster average
$\bar{\theta}_{j}^{t}=\frac{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\theta_{i}^{t}}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}.$

13: Compute step towards cluster average for all clients $i\in\mathcal{I}_{j}$
$\theta_{i}^{t+1}=\left(1-\frac{\eta\gamma_{i}(1-\tau_{j})(1-\alpha_{j})}{(1-p_{0})p_{j}}\right)\theta_{i}^{t}+\frac{\eta\gamma_{i}(1-\tau_{j})(1-\alpha_{j})}{(1-p_{0})p_{j}}\bar{\theta}_{j}^{t}$.

14: else

15: All Clients $i\in\mathcal{I}_{j}$ perform a local gradient descent step
$\theta_{i}^{t+1}=\theta_{i}^{t}-\frac{\eta}{(1-p_{0})(1-p_{j})}\nabla
f_{i}(\theta_{i}^{t})$.

16: end if

17: end for

18: end if

19: end for

At the beginning of each round in Algorithm 1, the network randomly decides
whether to aggregate the averages between all clusters or not. If a global
aggregation round is performed, each cluster first calculates its cluster
average and then communicates with one another (or a central server) to
compute the between-cluster aggregate.

In a round that does not involve communication between clusters, we allow each
cluster to independently decide whether to communicate within the cluster or
not by randomly sampling $\xi_{j}\sim\text{Bernoulli}(p_{j})$. By allowing
$p_{j}$ to be different across clusters, we effectively allow different
clusters to have different communication schedules that are specific to them.
In real-world applications, this flexibility can be appreciated. For example,
when we partition clients according to their geographical locations, we
effectively allow clients in different regions to communicate according to
different schedules, which could reduce communication latency and energy
consumption (Abad et al., 2020; Liu et al., 2020). We characterize the impact
of $\\{p_{j}\\}_{j=0}^{k}$ on the number of communication rounds in
Proposition 3.2.

######  Proposition 3.2 (Expected Number of Communication Rounds).

Suppose that Algorithm 1 is run for $T$ rounds. The expected number of
communication rounds between clusters is $p_{0}(1-p_{0})T$. The expected
number of communication rounds within the cluster $j$ is
$(1-p_{0})p_{j}(1-p_{j})T$.

###### Proof.

By Lemma 4.3 in Hanzely et al. (2020a), the expected number of between-cluster
communication rounds is given by $p_{0}(1-p_{0})T$. A within-cluster
communication round occurs only when $\xi_{0}=0$, which occurs with
probability $1-p_{0}$. ∎

As can be seen from Algorithm 1, any communication between clusters
necessarily implies that all clusters have computed within-cluster average.
Intuitively speaking, it would then be more efficient to optimize
$\psi_{j}\left(\bm{\theta}_{j};\\{\gamma_{i}\\}_{i\in\mathcal{I}_{j}}\right)$
whenever the cluster averages are computed in both the within-cluster and
between-cluster aggregation rounds. Unfortunately, a naive implementation of
this idea would lead to a biased gradient oracle. Suppose that we make a
gradient descent on $\psi_{j}$ and $\varphi$ with the same learning rate.
Since we descend on $\psi_{j}$ in both types of communication rounds, using
the same learning rate for both $\psi_{j}$ and $\varphi$ would effectively
cause us to descend twice on $\psi_{j}$, as the stochastic gradient oracle is
biased towards updating $\psi_{j}$ more frequently. We introduce variables
$\tau_{1},\ldots,\tau_{j}$ to scale the effective step size for $\psi_{j}$ in
the two different types of communication rounds, thereby ensuring that the
stochastic gradient oracle is unbiased.

Algorithm 1 induces a stochastic gradient oracle that is equivalent to SGD
with the following oracle. Assuming $p_{j},\tau_{j}\in(0,1)$, $j=0,\ldots,k$,
Algorithm 1 defines a stochastic gradient oracle for $F(\bm{\theta})$, denoted

$G(\bm{\theta})=(G_{1}(\theta_{1})^{T},\ldots,G_{n}(\theta_{n})^{T})^{T}\in\mathbb{R}^{nd},$

where for each $j$ and $i\in\mathcal{I}_{j}$,
$G_{i}(\theta_{i})\in\mathbb{R}^{d}$ and

$G_{i}(\theta_{i})=\begin{cases}\frac{\gamma_{i}\alpha_{j}}{p_{0}}\left(\theta_{i}^{t}-\bar{\theta}^{t}\right)+\frac{\gamma_{i}\tau_{j}(1-\alpha_{j})}{p_{0}}\left(\theta_{i}^{t}-\bar{\theta_{j}^{t}}\right),&\text{
if $\xi_{0}=1$}\\\
\frac{\gamma_{i}(1-\tau_{j})(1-\alpha_{j})}{(1-p_{0})p_{j}}\left(\theta_{i}^{t}-\bar{\theta}_{j}^{t}\right),&\text{
if $\xi_{0}=0$ and $\xi_{j}=1$}\\\ \frac{1}{(1-p_{0})(1-p_{j})}\nabla
f_{i}(\theta_{i}^{t}),&\text{ if $\xi_{0}=\xi_{j}=0$}.\end{cases}$ (3.2)

Note that $G_{i}(\theta_{i})$ is an unbiased estimator of
$\nabla_{\theta_{i}}F(\bm{\theta})$ for all $i$, ensuring that
$G(\bm{\theta})$ is an unbiased estimator of
$\nabla_{\bm{\theta}}F(\bm{\theta})$. Intuitively speaking, we can view
Algorithm 1 as a stochastic gradient descent procedure on $F(\bm{\theta})$
with step size $\eta$, using a noisy gradient oracle $G(\cdot)$.

The variance of $G(\bm{\theta})$ with respect to $\\{\xi_{j}\\}_{j=0}^{k}$
depends on a combination of $\\{p_{j}\\}_{j=0}^{k}$ and
$\\{\tau_{j}\\}_{j=1}^{k}$. Fixing $\tau_{j}$ to some arbitrary value is
suboptimal, and we should properly tune the parameter. Intuitively speaking,
when $p_{j}$ is small, within-cluster communications are less frequent. A
larger $\tau_{j}$ allows between-cluster communication rounds to “help out”
more when optimizing
$\psi_{j}\left(\bm{\theta}_{j};\\{\gamma_{i}\\}_{i\in\mathcal{I}_{j}}\right)$,
instead of relying on within-cluster communication. We discuss how they should
be adjusted according to the frequency of local and global aggregation rounds,
that is, how $\\{\tau_{j}\\}_{j=1}^{k}$ should be chosen given
$\\{p_{j}\\}_{j=0}^{k}$.

######  Proposition 3.3.

Suppose $\tau_{j}={p_{0}}(p_{0}+2(1-p_{0})p_{j})^{-1}$, $j=1,\dots,k$. Then

$\mathbb{E}_{\\{\xi_{j}\\}_{j=0}^{k}}[\|G(\bm{\theta})-G(\widehat{\bm{\theta}})\|^{2}]\\\
\leq\frac{2}{p_{0}}\left\|\nabla_{\bm{\theta}}\varphi(\bm{\theta})-\nabla_{\bm{\theta}}\varphi(\widehat{\bm{\theta}})\right\|^{2}+\sum_{j=1}^{k}\frac{2(1-\alpha_{j})^{2}}{p_{0}+2(1-p_{0})p_{j}}\|\nabla_{\bm{\theta}_{j}}\psi_{j}(\bm{\theta}_{j})-\nabla_{\bm{\theta}_{j}}\psi_{j}(\widehat{\bm{\theta}}_{j})\|^{2}\\\
+\frac{1}{1-p_{0}}\sum_{j=1}^{k}\frac{1}{1-p_{j}}\|\nabla_{\bm{\theta}_{j}}F_{j}(\bm{\theta}_{j})-\nabla_{\bm{\theta}_{j}}F_{j}(\widehat{\bm{\theta}}_{j})\|^{2}.$

###### Proof.

See Section A.3. ∎

Intuitively, $\tau_{j}$ balances between the probability of a communication
round between clusters, $p_{0}$, and the probability of a communication round
within a cluster, $(1-p_{0})p_{j}$. Therefore, it reduces the equivalent
gradient oracle variance. Consider the extreme case where $p_{0}\neq 0$ while
$p_{j}=0$, that is, there are no communication rounds within a cluster. Our
choice for $\tau_{j}$ is then exactly 1, which means that we will optimize
$\psi_{j}\left(\bm{\theta}_{j};\\{\gamma_{i}\\}_{i\in\mathcal{I}_{j}}\right)$
only during communication rounds between clusters, which is expected.

###  3.3 Convergence Analysis

We begin our convergence analysis by analyzing the convexity and smoothness of
the loss function, $F(\bm{\theta})$. First, we show that the regularizers
$\psi_{j}$ and $\varphi$ are convex and smooth.

######  Proposition 3.4.

The regularizers are convex and smooth. In particular,

  1. 1.

For all $j$,
$\psi_{j}\left(\bm{\theta}_{j};\left\\{\gamma_{i}\right\\}_{i\in\mathcal{I}_{j}}\right)$
is convex and $\max_{i\in\mathcal{I}_{j}}\gamma_{i}$-smooth in
$\bm{\theta}_{j}$.

  2. 2.

$\varphi\left(\bm{\theta};\left\\{\gamma_{i}\right\\}_{i=1}^{n},\left\\{\alpha\right\\}_{j=1}^{k}\right)$
is convex and
$\max_{j=1,\ldots,k}\max_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}$-smooth in
$\bm{\theta}$.

###### Proof.

See Section A.4. ∎

Under 2.1, from Proposition 3.4 it follows that the loss function
$F(\bm{\theta})$ is $\mu$-strongly convex in $\bm{\theta}$ and has a unique
minimizer for any set of penalty parameters
$(\\{\alpha_{j}\\}_{j=1}^{k},\\{\gamma_{i}\\}_{i=1}^{n})$. Let
$\widehat{\Theta}(\\{\alpha_{j}\\}_{j=1}^{k},\\{\gamma_{i}\\}_{i=1}^{n})$ be
the corresponding minimizer. When the penalty parameters
$\\{\alpha_{j}\\}_{j=1}^{k},\\{\gamma_{i}\\}_{i=1}^{n}$ are clear from the
context, we write

$\displaystyle\widehat{\Theta}=\widehat{\Theta}(\\{\alpha_{j}\\}_{j=1}^{k},\\{\gamma_{i}\\}_{i=1}^{n}),\qquad\psi_{j}(\cdot)=\psi_{j}\left(\cdot;\\{\gamma_{i}\\}_{i\in\mathcal{I}_{j}}\right),\quad
j=1,\ldots,k,$
$\displaystyle\varphi(\cdot)=\varphi\left(\cdot;\\{\gamma_{i}\\}_{i=1}^{n},\\{\alpha_{j}\\}_{j=1}^{k}\right),\quad
j=1,\ldots,k.$

Finally, we use
$F_{j}(\bm{\theta}_{j})=\sum_{i\in\mathcal{I}_{j}}f_{i}(\theta_{i})$ to denote
the average unregularized loss function in each cluster. The following theorem
provides the convergence rate.

######  Theorem 3.5 (Convergence Rate).

Suppose 2.1 holds and $\\{\tau_{j}\\}_{j=1}^{k}$ are set according to
Proposition 3.3. Let

$\mathcal{L}=\max\left\\{\frac{2}{p_{0}}\max_{j=1,\ldots,k}\max_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i},\max_{j=1,\dots,k}\frac{2(1-\alpha_{j})\max_{i\in\mathcal{I}_{j}}\gamma_{i}}{p_{0}+2(1-p_{0})p_{j}},\frac{L}{1-p_{0}}\max_{j=1,\dots,k}\frac{1}{1-p_{j}}\right\\},$
(3.3)

and

$\begin{split}\sigma_{\widehat{\bm{\theta}}}^{2}=&\frac{2}{p_{0}}\|\nabla_{\bm{\theta}}\varphi(\widehat{\bm{\theta}})\|^{2}+\sum_{j=1}^{k}\frac{2(1-\alpha_{j})^{2}}{p_{0}+2(1-p_{0})p_{j}}\|\nabla_{\theta_{j}}\psi_{j}(\widehat{\bm{\theta}}_{j})\|^{2}\\\
&\qquad+\frac{1}{1-p_{0}}\sum_{j=1}^{k}\frac{1}{1-p_{j}}\|\nabla_{\bm{\theta}_{j}}F_{j}(\widehat{\bm{\theta}}_{j})\|^{2}.\end{split}$
(3.4)

If the step size satisfies $\eta\leq\frac{1}{2\mathcal{L}}$, then

$\mathbb{E}[\|\bm{\theta}^{t}-\widehat{\bm{\theta}}\|^{2}]\leq\left(1-\eta\mu\right)^{t}\|\bm{\theta}^{0}-\widehat{\bm{\theta}}\|^{2}+\frac{2\eta\sigma_{\widehat{\bm{\theta}}}^{2}}{\mu}.$

###### Proof.

See Section A.5. ∎

Compared with the single cluster result in Hanzely and Richtárik (2020), our
convergence rate involves a few additional terms. By allowing each cluster to
determine when to communicate individually, the expected smoothness
coefficient and the variance of the gradient oracle at the optimum are more
complex, as they incorporate both local and global communication frequencies,
$\\{p_{j}\\}_{j=1}^{k}$ and $p_{0}$. When we set $p_{j}=0,\alpha_{j}=1$ for
all $j$, we recover the single cluster convergence rates given in Theorem 4.5
of Hanzely and Richtárik (2020) up to constant factors.

The parameters $p_{j}$, $j=0,\dots,k$, that control the frequency of
communication can be tuned by minimizing $\mathcal{L}$ in principle.
Unfortunately, as $\mathcal{L}$ is effectively a maximum taken over $2k+1$
different terms, directly minimizing the expression is infeasible. We instead
consider minimizing the following upper bound on $\mathcal{L}$:

$\tilde{\mathcal{L}}=\max\left\\{\frac{2}{p_{0}}C_{1},\max_{j=1,\ldots,k}\frac{2C_{2}}{p_{0}+2(1-p_{0})p_{j}},\frac{L}{1-p_{0}}\max_{j=1,\ldots,k}\frac{1}{1-p_{j}}\right\\},$
(3.5)

where
$C_{1}=\max_{j=1,\ldots,k}\max_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}$ and
$C_{2}=\max_{j=1,\ldots,k}\max_{i\in\mathcal{I}_{j}}(1-\alpha_{j})\gamma_{i}$.
The choice of parameters depends on the relationship between $C_{1}$ and
$C_{2}$.

######  Corollary 3.6.

Suppose that $\\{\tau_{j}\\}_{j=1}^{k}$ are set according to Proposition 3.3.

When $C_{2}>C_{1}$, setting $\eta=\frac{1}{2\mathcal{L}}$,
$p_{0}=\frac{2C_{1}}{C_{1}+C_{2}+L}$, and
$p_{j}=\frac{C_{2}-C_{1}}{C_{2}-C_{1}+L}$ ensures that the optimal number of
iterations is in
$\mathcal{O}\left(\frac{(C_{1}+C_{2}+L)}{\mu}\log\frac{1}{\epsilon}\right)$,
the number of communication rounds between clusters is in
$\mathcal{O}\left(\frac{C_{1}(C_{2}-C_{1}+L)}{(C_{1}+C_{2}+L)\mu}\log\frac{1}{\epsilon}\right)$,
and the number of communication rounds within a cluster is in
$\mathcal{O}\left(\frac{L(C_{2}-C_{1})}{(C_{2}-C_{1}+L)\mu}\log\frac{1}{\epsilon}\right)$
for all clusters.

When $C_{2}\leq C_{1}$, setting $\eta=\frac{1}{2\mathcal{L}}$,
$p_{0}=\frac{2C_{1}}{2C_{1}+L}$, and $p_{j}=0$ ensures that the optimal number
of iterations is in
$\mathcal{O}\left(\frac{(C_{1}+L)}{\mu}\log\frac{1}{\epsilon}\right)$, the
number of communication rounds between clusters is in
$\mathcal{O}\left(\frac{C_{1}L}{(C_{1}+L)\mu}\log\frac{1}{\epsilon}\right)$,
and the number of communication rounds within a cluster is 0.

###### Proof.

See Section A.6. ∎

We conclude the section by emphasizing that Algorithm 1 does not require a
central server to aggregate information across all clusters. The central
server in Algorithm 1 only serves two purposes: flipping a coin ($\xi_{0}$) to
determine whether a communication round between clusters is necessary and
calculating the global average should there be a communication round between
clusters. The former can be easily decentralized across clusters by asking all
cluster servers to flip a coin and take a majority vote, while the latter can
be implemented by asking all cluster servers to communicate with one another.

##  4 Asynchronous Accelerated Loopless Local SGD with Variance Reduction

The convergence rate of Algorithm 1 when minimizing the loss in Eq. 2.1 is
suboptimal due to the lack of acceleration and variance reduction. In this
section, we propose an accelerated, variance reduced, stochastic variant of
Algorithm 1 tailored to the finite-sum setting. The variant enjoys optimal
communication complexity in a single cluster setting Hanzely et al. (2020a),
regardless of the relationship between regularization strength and smoothness
of loss functions $f_{i}$, outperforming popular alternatives discussed in
Dinh et al. (2020); Hanzely et al. (2021); Li et al. (2021); Mansour et al.
(2020). We further hypothesize that the optimality holds when extended to
multi-cluster setup studied here.

Algorithm 2 Async-AL2SGD+

Input: Step size $\eta$, probabilities $p_{1},\dots,p_{k},p_{0},\rho\in[0,1]$,
fractions $\tau_{1},\dots,\tau_{k}\in[0,1]$.

Initialize: $1<a_{0},a_{1}<1$, $b_{1},b_{2}>0$,
$x_{i}^{0}=y_{i}^{0}=z_{i}^{0}=\theta_{i}^{0}=\bm{0}_{d}\in\mathbb{R}^{d}$

for $t=1,2,\dots$ do

All Clients $i=1,\dots,n$ perform local update
$\theta_{i}^{t}=a_{1}z_{i}^{t}+a_{2}x_{i}^{t}+(1-a_{1}-a_{2})y_{i}^{t}$.

$\xi_{0}=1$ with probability $p_{0}$ and 0 with probability $1-p_{0}$

if $\xi_{0}=1$ then

All Clusters $j=1,\dots,k$ compute cluster average
$\bar{\theta}_{j}^{t}=\frac{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\theta_{i}^{t}}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}$

All Clusters aggregate network average
$\bar{\theta}^{t}=\frac{\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}\theta_{i}^{t}}{\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}}$

All Clients calculate gradient estimate according to Eq. 4.1

Set $y_{i}^{t+1}=\theta_{i}^{t}-\eta g_{i}^{t}$

else

for All Clusters $j=1,\dots,k$ in parallel do

$\xi_{j}=1$ with probability $p_{j}$ and 0 with probability $1-p_{j}$

if $\xi_{j}=1$ then

Compute cluster average
$\bar{\theta}_{j}^{t}=\frac{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\theta_{i}^{t}}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}$
and send it back to each client

All Clients calculate gradient estimate according to Eq. 4.1

Set $y_{i}^{t+1}=\theta_{i}^{t}-\eta g_{i}^{t}$

else

All Clients calculate gradient estimate according to Eq. 4.1

Set $y_{i}^{t+1}=\theta_{i}^{t}-\eta g_{i}^{t}$

end if

end for

end if

All Clients $i=1,\dots,n$ update:
$z_{i}^{t+1}=b_{1}z_{i}^{t}+(1-b_{1})\theta_{i}^{t}+\frac{b_{2}}{\eta}(y_{i}^{t+1}-\theta_{i}^{t})$

$\xi^{\prime}=1$ with probability $\rho$ and 0 with probability $1-\rho$

if $\xi^{\prime}=0$ then

For all Clients $i=1,\dots,n$: $x_{i}^{t+1}=x_{i}^{t}$

else

For all Clients $i=1,\dots,n$ update $x_{i}^{t+1}=y_{i}^{t+1}$, and evaluate
and store $\nabla f_{i}(x_{i}^{t+1})$

All Clusters communicate, compute averages
$\bar{x}^{t}=\frac{\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}x_{i}^{t}}{\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}}$,
$\bar{x}_{j}^{t}=\frac{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}x_{i}^{t}}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}$
for all $j$, and send them back to the clients.

end if

end for

We assume that the local loss has a finite sum structure over smooth and
strongly convex functions, a common assumption in the literature on
accelerated variance reduced algorithms (Hanzely et al., 2020a; Hanzely and
Richtárik, 2020; Kovalev et al., 2020). We formally characterize our
assumption below.

######  Assumption 4.1.

The loss function $f_{i}$, $i=1,\ldots,n$, has the following finite structure:

$f_{i}(\theta)=\frac{1}{n_{i}}\sum_{l=1}^{n_{i}}\tilde{f}_{i,l}(\theta),$

where $\widetilde{f}_{i,l}$ is $\tilde{L}$-smooth and $\mu$-strongly convex,
$l=1,\ldots,n_{i}$.

We define a stochastic gradient estimate for all clients, similar to the
construction for Algorithm 1. For a client $i$ that belongs to the cluster
$j$, the variance reduced stochastic gradient is

$\begin{split}g_{i}^{t}&=\nabla
f_{i}(x_{i}^{t})+\alpha_{j}\gamma_{i}(x_{i}^{t}-\bar{x}^{t})+(1-\alpha_{j})\gamma_{i}(x_{i}^{t}-\bar{x}_{j}^{t})\\\
&\quad\quad+\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{0}=1\\}\frac{\gamma_{i}\alpha_{j}}{p_{0}}(\theta_{i}^{t}-\bar{\theta}^{t}-(x_{i}^{t}-\bar{x}^{t}))\\\
&\quad\quad+\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{0}=0\\}\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{j}=1\\}\frac{\gamma_{i}(1-\tau_{j})(1-\alpha_{j})}{(1-p_{0})p_{j}}\left((\theta_{i}^{t}-\bar{\theta}_{j}^{t})-(x_{i}^{t}-\bar{x}_{j}^{t})\right)\\\
&\quad\quad+\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{0}=0\\}\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{j}=0\\}\frac{1}{(1-p_{0})(1-p_{j})}\left(\nabla\tilde{f}_{i,l}(\theta_{i}^{t})-\nabla\tilde{f}_{i,l}(x_{i}^{t})\right),\end{split}$
(4.1)

where $l$ is selected uniformly at random at each iteration for every client.
At a high level, Eq. 4.1 defines a stochastic gradient oracle for a finite-sum
composite optimization problem. The different realizations of
$\\{\xi_{0},\ldots,\xi_{j}\\}$ determine the type of communication round to
execute, if any, at any given step. For example, when $\xi_{0}=1$, a
communication round between the client clusters is executed, while when
$x_{0}=0$ and $\xi_{1}=1$, the first client cluster executes a communication
round within the cluster. The following lemma provides a bound on the variance
of the stochastic gradient oracle in Eq. 4.1.

######  Lemma 4.2.

Suppose that 4.1 holds and that $\\{\tau_{j}\\}_{j=1}^{k}$ are selected as in
Proposition 3.3. Let

$\mathcal{L}=\max\left\\{\frac{2}{p_{0}}\max_{j=1,\ldots,k}\max_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i},\max_{j=1,\dots,k}\frac{2(1-\alpha_{j})\max_{i\in\mathcal{I}_{j}}\gamma_{i}}{p_{0}+2(1-p_{0})p_{j}},\frac{\tilde{L}}{1-p_{0}}\max_{j=1,\dots,k}\frac{1}{1-p_{j}}\right\\}.$

Then

$\mathbb{E}\left[\|\bm{g}^{t}-\nabla F(\bm{x}^{t})\|^{2}\right]\leq
2\mathcal{L}D_{F}(\bm{\theta}^{t},\bm{x}^{t}),$ (4.2)

where $\bm{g}^{t}=(g_{1}^{t},\ldots,g_{n}^{t})^{T}\in\mathbb{R}^{nd}$ is the
variance reduced stochastic gradient oracle, where $g_{i}^{t}$ is defined in
Eq. 4.1, and $D_{F}(x_{1},x_{2})\coloneq F(x_{1})-F(x_{2})-\langle\nabla
F(x_{2}),x_{1}-x_{2}\rangle$ is the Bregman divergence induced by the loss
function $F$ in Eq. 2.1.

###### Proof.

See Section A.7. ∎

We use the oracle to construct an instance of L-Katyusha, a loopless,
variance-reduced, accelerated algorithm Qian et al. (2021b). Algorithm 2
provides the pseudocode, while Theorem 4.3 provides the convergence rate for
Algorithm 2.

######  Theorem 4.3.

Suppose that the conditions of Lemma 4.2 are satisfied. Let

$\displaystyle
L_{F}=\tilde{L}+\max_{i=1,\ldots,n}\gamma_{i},\quad\eta=\frac{1}{4}\max\\{L_{F},\mathcal{L}\\}^{-1},$
$\displaystyle
a_{1}=\min\left\\{\frac{1}{2},\sqrt{\eta\mu\max\left\\{\frac{1}{2},\frac{a_{2}}{\rho}\right\\}}\right\\},\quad
a_{2}=\frac{\mathcal{L}}{2\max\\{L_{F},\mathcal{L}\\}},$ $\displaystyle
b_{1}=1-b_{2}\mu,\quad b_{2}=\frac{1}{\max\\{2\mu,4a_{1}/\eta\\}}.$

Then the iteration complexity of Algorithm 2 is

$\mathcal{O}\left(\left(\frac{1}{\rho}+\sqrt{\frac{\mathcal{L}}{\rho\mu}}\right)\log\frac{1}{\epsilon}\right).$

###### Proof.

See Section A.8. ∎

In addition to the iteration complexity, it is easy to obtain communication
bounds, gradient complexity, as well as optimal parameters. Similar to
Corollary 3.6, directly minimizing $\mathcal{L}$ over all the parameters is
infeasible and we consider the following upper bound on $\mathcal{L}$ instead:

$\tilde{\mathcal{L}}=\max\left\\{\frac{2C_{1}}{p_{0}},\frac{2C_{2}}{p_{0}+2(1-p_{0})p_{j}},\frac{\tilde{L}}{1-p_{0}}\max_{j=1,\ldots,k}\frac{1}{1-p_{j}}\right\\},$

where
$C_{1}=\max_{j=1,\ldots,k}\max_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}$ and
$C_{2}=\max_{j=1,\ldots,k}\max_{i\in\mathcal{I}_{j}}(1-\alpha_{j})\gamma_{i}$.
The upper bound is virtually the same as the one in Eq. 3.5, and we have the
following.

######  Corollary 4.4.

Consider Algorithm 2 with a fixed $\rho$ and the tuning parameters set as:

$\displaystyle
L_{F}=\tilde{L}+\max_{i=1,\ldots,n}\gamma_{i},\quad\eta=\frac{1}{4}\max\\{L_{F},\mathcal{L}\\}^{-1},$
$\displaystyle
a_{1}=\min\left\\{\frac{1}{2},\sqrt{\eta\mu\max\left\\{\frac{1}{2},\frac{a_{2}}{\rho}\right\\}}\right\\},\quad
a_{2}=\frac{\mathcal{L}}{2\max\\{L_{F},\mathcal{L}\\}},$ $\displaystyle
b_{1}=1-b_{2}\mu,\quad b_{2}=\frac{1}{\max\\{2\mu,4a_{1}/\eta\\}}.$

Furthermore, $\\{\tau_{j}\\}_{j=1}^{k}$ is set according to Proposition 3.3.

When $C_{2}>C_{1}$, setting $p_{0}=\frac{2C_{1}}{C_{1}+C_{2}+\tilde{L}}$ and
$p_{j}=\frac{C_{2}-C_{1}}{C_{2}-C_{1}+\tilde{L}}$ ensures that the optimal
number of iterations is in
$\mathcal{O}\left(\sqrt{\frac{C_{1}+C_{2}+\tilde{L}}{\mu}}\log\frac{1}{\epsilon}\right)$,
the number of communication rounds between clusters and within a cluster is in
$\mathcal{O}\left(\frac{C_{1}(C_{2}-C_{1}+\tilde{L})}{(C_{1}+C_{2}+\tilde{L})\sqrt{(C_{1}+C_{2}+\tilde{L})\mu}}\log\frac{1}{\epsilon}\right)$
and
$\mathcal{O}\left(\frac{(C_{2}-C_{1})\tilde{L}}{(C_{2}-C_{1}+\tilde{L})\sqrt{(C_{1}+C_{2}+\tilde{L})\mu}}\log\frac{1}{\epsilon}\right)$,
respectively.

When $C_{2}\leq C_{1}$, setting $p_{0}=\frac{2C_{1}}{2C_{1}+\tilde{L}}$ and
$p_{j}=0$ ensures the optimal number of iterations is in
$\mathcal{O}\left(\sqrt{\frac{C_{1}+\tilde{L}}{\mu}}\log\frac{1}{\epsilon}\right)$,
the number of communication rounds between clusters is in
$\mathcal{O}\left(\frac{C_{1}\tilde{L}}{(C_{1}+\tilde{L})\sqrt{(C_{1}+\tilde{L})\mu}}\log\frac{1}{\epsilon}\right)$,
and the number of communication rounds within a cluster is 0.

###### Proof.

See Section A.9. ∎

##  5 Case Study: Hierarchical Linear Model

In the previous section, we have answered how to minimize the objective in Eq.
2.1. Next, we provide a statistical model of personalization under which the
minimizer of Eq. 2.1 corresponds to an estimator that outperforms the common
alternatives. More precisely, we show that the minimizer of Eq. 2.1 strictly
outperforms both (a) training a single global model for all clients and (b)
training a separate model for each client independent of the data of other
clients. Unlike the analysis in Li et al. (2021), we consider the
hierarchical, multi-cluster regime. Existing approaches to personalized
federated learning often use loss functions similar to those discussed in Eqs.
2.1 and 2.3. Therefore, it is important to understand the statistical
properties of the corresponding minimizers.

Although the minimizer Eq. 2.1 outperforms commonly used alternatives, we also
provide two alternative estimators that are hard to efficiently compute in a
federated learning setting, yet dominate our proposed estimator. Specifically,
they achieve a lower mean squared error. The efficient implementation of the
two alternatives remains an open question for future research.

The statistical model we consider in this section is based on a hierarchical
linear model with Gaussian priors (Stephen and Anthony, 2002). Nature first
draws the cluster centers from a Gaussian distribution with unknown mean and
then draws each client’s parameter from a Gaussian distribution centered at
the cluster center the client belongs to. More precisely, for an unknown
parameter $\bar{\theta}^{*}\in\mathbb{R}^{d}$, our model is:

$\begin{array}[]{lll}\bar{\theta}_{j}^{*}=\bar{\theta}^{*}+\bar{\xi}_{j},&\bar{\xi}_{j}\sim\mathcal{N}(0,\bar{\sigma}^{2}I_{d}),&j=1,\ldots,k,\\\
\theta_{i}^{*}=\bar{\theta}_{j}^{*}+\xi_{i},&\xi_{i}\sim\mathcal{N}(0,\bar{\sigma}_{j}^{2}I_{d}),&i\in\mathcal{I}_{j},\\\
y_{i}=X_{i}\theta_{i}^{*}+\epsilon_{i},&\epsilon_{i}\sim\mathcal{N}(0,\sigma_{i}^{2}I_{n_{i}}),&i\in\mathcal{I}_{j},\end{array}$
(5.1)

where $\bar{\theta}_{j}^{*}\in\mathbb{R}^{d}$ represents the center of the
cluster $j$, $\theta_{i}^{*}\in\mathbb{R}^{d}$ represents the client-specific
parameter, and $(X_{i},y_{i})\in\mathbb{R}^{n_{i}\times
d}\times\mathbb{R}^{n_{i}}$ is the data set on the $i$-th client.

When estimating all client parameters simultaneously in Eq. 5.1, we obtain the
following maximum likelihood estimation problem:

$\displaystyle\min_{\\{\theta_{i}\\}_{i=1}^{m}}\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\left(\frac{1}{\sigma_{i}^{2}}\|y_{i}-X_{i}\theta_{i}\|^{2}+\frac{\gamma_{i}\alpha_{j}}{2}\|\theta_{i}-\bar{\theta}\|^{2}+\frac{\gamma_{i}(1-\alpha_{j})}{2}\|\theta_{i}-\bar{\theta}_{j}\|\right),$

where

$\displaystyle\bar{\theta}$
$\displaystyle=\left(\sum_{j^{\prime}=1}^{k}\sum_{i^{\prime}\in\mathcal{I}_{j^{\prime}}}\gamma_{i^{\prime}}\alpha_{j^{\prime}}\right)^{-1}\sum_{j=1}^{k}\sum_{i=\in\mathcal{I}_{j}}\gamma_{i}\alpha_{j}\theta_{i};$
$\displaystyle\bar{\theta}_{j}$
$\displaystyle=\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\gamma_{i^{\prime}}\right)^{-1}\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\theta_{i},\
j=1,\ldots,k.$

The objective is an instance of Eq. 2.1, and by Proposition 2.2, is equivalent
to

$\min_{\\{\theta_{i}\\}_{i=1}^{m},\\{w_{j}\\}_{j=1}^{k},\bar{w}}\quad\sum_{j=1}^{k}\left(\frac{\lambda_{j}}{2}\|w_{j}-\bar{w}\|^{2}+\sum_{i\in\mathcal{I}_{j}}\left(\frac{1}{2\sigma_{i}^{2}}\|y_{i}-X_{i}\theta_{i}\|^{2}+\frac{\gamma_{i}}{2}\|\theta_{i}-w_{j}\|^{2}\right)\right).$
(5.2)

We focus on Eq. 5.2 for convenience and show that, when
$\\{\lambda_{j}\\}_{j=1}^{k}$, $\\{\gamma_{i}\\}_{i=1}^{n}$ are properly tuned
and $\bar{\sigma}^{2}$, $\bar{\sigma}_{j}^{2}$ are known, the resulting
minimizers $\\{\widehat{\theta}_{i}\\}_{i=1}^{n}$ attain the smallest mean
squared error among a class of linear unbiased estimators (Kariya and Kurata,
2004).

######  Theorem 5.1.

Suppose that $\lambda_{j}=\left(\bar{\sigma}^{2}\right)^{-1}$,
$\gamma_{i}=\left(\bar{\sigma}_{j}^{2}\right)^{-1}$, and
$X_{i}^{T}X_{i}=\beta_{i}I_{d}$ for some $\beta_{i}\in\mathbb{R}$,
$i\in\mathcal{I}_{j}$, $j=1,\ldots,k$. Then $\widehat{\theta}_{i}$, obtained
as the minimizer of Eq. 5.2, is the best linear unbiased estimator of
$\theta_{i}^{*}$ given
$\\{(X_{i},y_{i})\\}\cup\\{\widehat{\theta}_{i^{\prime}}^{d}\\}_{i^{\prime}\neq
i}$, where $\widehat{\theta}_{i}^{d}=(X_{i}^{T}X_{i})^{-1}(X_{i}y)$.

###### Proof.

See Section A.10. ∎

To provide an intuition behind Theorem 5.1, we first consider the single
cluster setting in Li et al. (2021). Focusing on an arbitrary cluster $j$, we
know that for all $i\in\mathcal{I}_{j}$, $\widehat{\theta}_{i}$ consists of
two parts: the first part estimates $\theta_{i}^{*}$ using only data on the
$i$-th client and the second part estimates $\bar{\theta}_{j}^{*}$ using other
clients’ parameters. Since Eq. 5.2 aggregates information across clients only
by regularizing the distance between weight estimates, the estimate for
$\bar{\theta}_{j}^{*}$ cannot depend directly on
$\\{(X_{i^{\prime}},y_{i^{\prime}})\\}_{i^{\prime}\neq i}$ and is constructed
instead through $\\{\widehat{\theta}_{i^{\prime}}^{d}\\}_{i^{\prime}\neq i}$.
In our multi-cluster setting, the estimator $\widehat{\theta}_{i}$ operates in
a similar fashion, as can be seen from the proof of Theorem 5.1. Since we do
not have direct access to the data of other clients, the estimates for
$\bar{\theta}_{j}^{*},\bar{\theta}^{*}$ are constructed indirectly using
$\\{\widehat{\theta}_{i^{\prime}}^{d}\\}_{i^{\prime}\neq i}$. A direct
consequence of Theorem 5.1 is the following.

######  Corollary 5.2.

Suppose that the conditions of Theorem 5.1 hold. Let
$\widehat{\theta}^{\text{all}}=\left(\sum_{i=1}^{n}X_{i}^{T}X_{i}\right)^{-1}\left(\sum_{i=1}^{n}X_{i}y_{i}\right)$.
Then

$\mathbb{E}[\|\widehat{\theta}_{i}-\theta_{i}^{*}\|^{2}]\leq\min\left\\{\mathbb{E}[\|\widehat{\theta}_{i}^{d}-\theta_{i}^{*}\|^{2}],\mathbb{E}[\|\widehat{\theta}^{all}-\theta_{i}^{*}\|^{2}]\right\\},$

where $\widehat{\theta}_{i}$ and $\widehat{\theta}_{i}^{d}$ are defined in
Theorem 5.1, and the expectation is taken over the randomness in Eq. 5.1.

Corollary 5.2 illustrates existence of a regime under which a personalized
estimator consistently outperforms the alternatives, learning a single model
for all clients without any personalization $\widehat{\theta}^{all}$ and
learning a model independently for each client $\widehat{\theta}_{i}^{d}$,
proposed in Chen et al. (2021), highlighting the effectiveness and necessity
of personalization within highly structured problems. Note that
$\widehat{\theta}_{i}^{d}$ and $\widehat{\theta}^{all}$ can both be written as
unbiased linear estimators of $\theta_{i}^{*}$ given
$\\{(X_{i},y_{i})\\}\cup\\{\widehat{\theta}_{i^{\prime}}^{d}\\}_{i^{\prime}\neq
i}$. The optimality of $\widehat{\theta}_{i}$ among this class of estimators
ensures that its mean squared error is no greater than these alternatives.

###  5.1 Limitations of Unbiased Estimators

We have shown that solving Eq. 5.2 consistently outperforms common
alternatives $\widehat{\theta}^{all}$ and $\widehat{\theta}_{i}^{d}$. In
particular, regularizing the distance between client model parameters and
average model parameters provides a viable approach for personalization. See
also Hanzely and Richtárik (2020); Li et al. (2021); Hanzely et al. (2021);
Dinh et al. (2020). Furthermore, our result complements Chen et al. (2021),
identifying a regime in which personalization consistently outperforms
learning a single model and learning models independently. Unfortunately, we
cannot guarantee that our approach is optimal among all possible estimators.
We discuss two alternative estimators that result in a lower mean squared
error, but are hard to implement in federated learning setting.

First, we note that Eq. 5.2 is the maximum likelihood estimator when we
simultaneously estimate all clients’ parameters. However, for any particular
client, we can derive an unbiased linear estimator with smaller mean squared
error, by marginalizing other clients’ and clusters’ parameters.

######  Proposition 5.3.

Suppose $\\{(X_{i},y_{i})\\}_{i=1}^{n}$ are generated according to the model
in Eq. 5.1. For any $i$, there exists a linear unbiased estimator
$\tilde{\theta}_{i}$ of $\theta_{i}^{*}$ that satisfies
$\mathbb{E}[\|\tilde{\theta}_{i}-\theta_{i}^{*}\|^{2}]\leq\mathbb{E}[\|\widehat{\theta}_{i}-\theta_{i}^{*}\|^{2}]$.

###### Proof.

See Section A.11. ∎

By marginalizing other parameters, we derive $\tilde{\theta}_{i}$, the best
linear unbiased estimator of $\theta_{i}^{*}$ given
$\\{(X_{i^{\prime}},y_{i^{\prime}})\\}_{i^{\prime}=1}^{n}$, which includes
data from other clients. An explicit form for the equation that
$\tilde{\theta}_{i}$ solves can be found in Section A.11. By contrast,
$\widehat{\theta}_{i}$ only has access to other clients’ weight estimates,
$\\{\widehat{\theta}_{i^{\prime}}^{d}\\}_{i^{\prime}\neq i}$, but does not
have direct access to their data. We emphasize, however, that doing so for all
the clients is costly, as we need to solve a separate generalized least
squares problem for all clients. On the other hand, learning
$\widehat{\theta}_{i}$ for all clients can be done simultaneously by
optimizing Eq. 5.2.

An alternative estimator of $\theta_{i}^{*}$ can be constructed based on
James-Stein estimator James and Stein (1992). Suppose that $X\sim
N(\xi,I_{p})$ is a $p$-dimensional Gaussian random vector with mean $\xi$ and
covariance $I_{p}$. The James-Stein Estimator of $\xi$ is defined as

$\widehat{\xi}^{\rm JS}=\left(1-\frac{p-2}{\|X\|_{2}^{2}}\right)X.$

Let $\widehat{\xi}^{\rm MLE}$ be the maximum likelihood estimator of $\xi$,
that is, $\widehat{\xi}^{\rm MLE}=X$. Then, for all $p\geq 3$,

$\mathbb{E}\left[\left\|\widehat{\xi}^{\rm
JS}-\xi\right\|^{2}\right]\leq\mathbb{E}\left[\left\|\widehat{\xi}^{\rm
MLE}-\xi\right\|^{2}\right].$

See James and Stein (1992) for a proof.

We construct a biased estimator that dominates $\widehat{\theta}_{i}$ even in
a simplified, single cluster regime, under which we are effectively solving a
$d$-dimensional point estimation problem.

######  Proposition 5.4.

Consider a single-cluster model under which $y_{i}\sim
N(\theta_{i}^{*},I_{d})$, where $\theta_{i}^{*}\sim
N(\bar{\theta}^{*},I_{d})$, $i=1,\dots,n$, and $\bar{\theta}^{*}\in\mathbb{R}$
is an unknown parameter. If $d>3$, then there exists a biased estimator
$\tilde{\theta}^{JS}_{i}$ such that
$\mathbb{E}[\|\tilde{\theta}_{i}^{JS}-\theta_{i}^{*}\|^{2}]\leq\mathbb{E}[\|\widehat{\theta}_{i}-\theta_{i}^{*}\|^{2}]$,
where $\widehat{\theta}_{i}$ is the best linear unbiased estimator for
$\theta_{i}^{*}$ given $\\{y_{i}\\}_{i=1}^{n}$.

###### Proof.

See Section A.12. ∎

Efficiently implementing the biased estimator is non-trivial. The estimator
requires careful adjustment of a shrinkage coefficient to achieve a smaller
mean squared error. How to tune this shrinkage coefficient efficiently in a
federated learning setting is unclear. As a result, the minimizer of Eq. 5.2
is a great practical alternative.

The two estimators provided in this section also point out that alternative
approaches to personalized federated learning (Deng et al., 2020; Li et al.,
2021; Hanzely and Richtárik, 2020) are not optimal from a statistical point of
view under a hierarchical linear model. While these two estimators are
impractical in federated learning setting, it remains an open question how to
approximate them with efficient computation and communication, while also
respecting the privacy concerns in federated learning.

##  6 Numerical Results

We illustrate the performance of our algorithm on both simulated data and a
real-world marketing data set. We focus on the generalization error of the
minimizer of Eq. 2.1, rather than on the optimization performance of
Algorithms 1 and 2, since related approaches have been studied in the single-
cluster regime Li et al. (2021); Hanzely and Richtárik (2020); Hanzely et al.
(2020a, 2021); Dinh et al. (2020). Our aim is to complement those studies.

We consider a modified version of Algorithm 1, where one coin toss is used to
determine whether the clients should perform a local step or a communication
round. During a communication round, each machine minimizes its distance to
the cluster and global average simultaneously. Our experiments focus on the
generalization behavior of the estimator rather than on the optimization
error. Therefore, such a simplification does not affect the validity of our
results and simplifies the implementation. Scripts for replicating the
experiments can be found this GitHub repository.

###  6.1 Simulation Studies

We compare our algorithm with the tuning parameters set as in Theorem 5.1
against three baselines: i) learning a single model for all clients, ii)
learning each client’s model independently, and iii) learning personalized
models centered around a single point Hanzely and Richtárik (2020); Li et al.
(2021); Dinh et al. (2020).

Simulation data are generated from the following hierarchical linear model.
The center of all clusters is
$\bar{\theta}^{*}=\mathbf{0}_{20}\in\mathbb{R}^{20}$, where $\mathbf{0}_{20}$
is an all-zero vector. There are 20 clusters, each with 20 clients. The
cluster centers and client parameters are generated as:

$\displaystyle\bar{\theta}_{j}^{*}\sim\mathcal{N}(\bar{\theta}^{*},I_{20}),\quad
j=1,\dots,k,$
$\displaystyle\theta_{i}^{*}\sim\mathcal{N}(\bar{\theta}_{j}^{*},I_{20}),\quad
i\in\mathcal{I}_{j}.$

On each client, we generate a data matrix $X_{i}\in\mathbb{R}^{m\times 20}$,
where $m$ is the number of samples on the $i$-th client and is selected from
$\\{1,5,25,50,100,200\\}$. Each entry in $X_{i}$ is drawn i.i.d. from the
standard Gaussian distribution, $\mathcal{N}(0,1)$. Subsequently, the response
$y_{i}$ is drawn from the linear model:

$y_{i}\mid X_{i}\sim\mathcal{N}(X_{i}\theta_{i}^{*},I_{n_{i}}).$

Under the data generating procedure described above, the three baselines take
the following form:

  1. 1.

Training a single-model for all clients:
$\widehat{\theta}^{\mathbf{sm}}=\left(\sum_{i=1}^{n}X_{i}^{T}X_{i}\right)^{\dagger}\left(\sum_{i=1}^{n}X_{i}^{T}y_{i}\right)$.

  2. 2.

Entirely locally-trained estimator: for each client $i$,
$\widehat{\theta}_{i}^{\mathbf{lt}}=(X_{i}^{T}X_{i})^{\dagger}X_{i}^{T}y_{i}$,
where $(X_{i}^{T}X_{i})^{\dagger}$ is the Moore-Penrose pseudo-inverse of the
empirical covariance matrix.

  3. 3.

Training a single-cluster personalized model (Hanzely and Richtárik, 2020; Li
et al., 2021; Dinh et al., 2020): the objective function is given by

$\min_{\\{\theta_{i}\\}}\frac{1}{n}\sum_{i=1}^{n}\left(\frac{1}{2}\|X_{i}\theta_{i}-y_{i}\|^{2}+\frac{\lambda^{\mathbf{sc}}}{2}\|\theta_{i}-\bar{\theta}\|^{2}\right),$

where $\bar{\theta}=n^{-1}\sum_{i=1}^{n}\theta_{i}$ and
$\lambda^{\mathbf{sc}}>0$ is user-chosen parameter that controls the strength
of personalization. The minimizer of the objective can be obtained as
$\widehat{\theta}_{i}^{\mathbf{sc}}=(X_{i}^{T}X_{i}+\lambda^{\mathbf{sc}}I_{20})^{-1}\left((X_{i}^{T}X_{i})\widehat{\theta}_{i}^{\mathbf{lt}}+\lambda^{\mathbf{sc}}\widehat{\bar{\theta}}^{\mathbf{sc}}\right)$,
where

$\widehat{\bar{\theta}}^{\mathbf{sc}}=\left(I_{20}-\frac{\lambda^{\textbf{sc}}}{n}\sum_{i=1}^{n}\left(X_{i}^{T}X_{i}+\lambda^{\mathbf{sc}}I_{20}\right)^{-1}\right)^{-1}\left(\frac{1}{n}\sum_{i=1}^{n}\left(X_{i}^{T}X_{i}+\lambda^{\mathbf{sc}}I_{20}\right)^{-1}(X_{i}^{T}X_{i})\widehat{\theta}_{i}^{\mathbf{lt}}\right).$

We tune $\lambda^{\mathbf{sc}}$ over 20 evenly spaced points in $[0.01,2]$
using cross validation.

Following Theorem 5.1, we set $\lambda_{j}=1$ for all $j$ and $\gamma_{i}=1$
for all $i$. We use Proposition 2.2 to convert the maximum likelihood
estimation problem into the form in Eq. 2.1, which we minimize using the
simplified algorithm with communication probability $p=0.1$, stepsize
$\eta=10^{-4}$, and maximum number of iterations $50000$. We use
$\\{\widehat{\theta}_{i}^{\mathbf{our}}\\}$ to denote the estimators produced
by our model.

We measure the performance of different estimators using the $\ell_{2}$
distance between the estimates of the parameters of the clients and their
actual parameters, that is, $\|\widehat{\theta}_{i}-\theta^{*}_{i}\|_{2}^{2}$
for
$\widehat{\theta}_{i}\in\\{\widehat{\theta}^{\mathbf{sm}},\widehat{\theta}^{\mathbf{lt}}_{i},\widehat{\theta}_{i}^{\mathbf{sc}},\widehat{\theta}_{i}^{\mathbf{our}}\\}$.
The results are averaged over five independent runs.

To demonstrate that our model consistently outperforms baselines, we pick two
specific choices of $m$, $m=10$, and $m=100$. Intuitively, as $m$ increases,
local training becomes more viable, whereas a smaller $m$ means that training
a single model could be more beneficial. Here, we show that our suggested
approach outperforms both these alternatives regardless of $m$. Figure 1 and
Table 1 confirm that our proposed method consistently outperforms alternatives
for $m\in\\{1,5,10,25,100,200\\}$.

| $m=10$ | $m=100$
---|---|---
| Avg. ($\pm$ SD.) | Max | Avg. ($\pm$ SD.) | Max
$\widehat{\theta}^{\mathbf{lt}}$ | 4.50 $(\pm 0.981)$ | 8.224 | 0.494 $(\pm 0.093)$ | 0.768
$\widehat{\theta}^{\mathbf{sm}}$ | 6.11 $(\pm 0.968)$ | 9.025 | 6.243 $(\pm 1.003)$ | 9.479
$\widehat{\theta}^{\mathbf{sc}}$ | 4.46 $(\pm 0.958)$ | 8.112 | 0.494 $(\pm 0.093)$ | 0.763
$\widehat{\theta}^{\mathbf{our}}$ | 3.46 $(\pm 0.689)$ | 8.676 | 0.489 $(\pm 0.093)$ | 0.748
Table 1:  The average
($\frac{1}{n}\sum_{i=1}^{n}\|\widehat{\theta}_{i}-\theta_{i}^{*}\|^{2}$) and
maximum ($\max_{i}\|\widehat{\theta}_{i}-\theta_{i}^{*}\|^{2}$) $\ell_{2}$
distance between parameter estimates and true parameters for
$m\in\\{10,100\\}$. The proposed method (last row) consistently outperforms
existing methods in two settings. Figure 1: The average $\ell_{2}$ distance
versus the number of local samples.

###  6.2 Application: Donor Response

We illustrate our algorithm on a real data set that contains donations and
solicitation histories from a leading nonprofit organization in the US
Blattberg et al. (2009). We follow the experimental setup described in Bumbaca
et al. (2017). For each solicitation record, we use its recency and frequency
as covariates, where recency is defined as the number of days since the
donor’s last donation and frequency is the number of past donations. We define
a solicitation as successful when the donor has made a donation after the
current solicitation attempt and before the next solicitation attempt and
model the probability of a successful solicitation using logistic regression
fitted on log-transformed covariates.

We view each ZIP code as an individual client and group the clients using the
median household income of the clients in the ZIP code, based on the data
obtained from www.unitedstateszipcodes.org. Using the income brackets defined
in Snider (2019), we group the ZIP code into 4 different clusters: poor-or-
near-poor, lower-middle-class, middle-class, and upper-middle-class. ZIP codes
with no recorded median household income are grouped into a fifth category,
and ZIP codes with less than 5 solicitations are removed. We retain 29490
clients and 5 clusters after processing.

We tune the parameters
$\\{\lambda_{j}\\}_{j=1}^{k},\\{\gamma_{i}\\}_{i=1}^{n}$, defined in Eq. 2.3,
by performing a grid search over
$\\{10^{-2},10^{-1.875},10^{-1.750},\ldots,10^{1.875},10^{2}\\}$. The chosen
values are then used to calculate the corresponding values for
$\\{\alpha_{j}\\}_{j=1}^{k}$. We set the test-train ratio to 0.2,
communication probability to 0.1, stepsize to $10^{-3}$, and the maximum
number of iterations to 5000.

We record the accuracy and cross-entropy for each individual client. We set
the cross-entropy to 100 for clients whose cross-entropy overflows. We then
tune $\lambda,\gamma$ with cross-validation, based on the average cross-
entropy taken across all the clients. Using cross-validation, we decide on
$\lambda_{j}=0.01$ for all $j$ and $\gamma_{i}=1.360$ for all $j$

Figure 2: Distribution of cross-entropy losses at each ZIP code. From left to
right: personalized model using selected parameters, models trained using
entirely local data, and a single model trained on all data.

Figure 2 visualizes the performance of our model compared to naive baselines.
Specifically, it characterizes the distribution of the clients’ cross-entropy.
To complement Figure 2, we include more quantitative results in Table 2, which
contains summary statistics on the clients’ cross-entropy. From the table, we
observe that our model is comparable to training a single model on all data in
terms of averages and quartiles, while it has a lower maximum client-specific
cross-entropy. In particular, we can see from Figure 2 that our model has a
smaller percentage of clients with average cross-entropy above 1.0: 183 ZIP
codes have cross-entropy loss greater than 1 in the personalized model,
whereas 202 have cross-entropy greater than 1 in the single-model alternative.
Our model also significantly outperforms training a local model for each ZIP
code separately.

| Avg. ($\pm$ SD.) | 25% | 50% | 75% | Max
---|---|---|---|---|---
Locally-trained | $0.413(\pm 0.932)$ | 0.173 | 0.247 | 0.313 | 6
Single-model | $0.211(\pm 0.176)$ | 0.058 | 0.201 | 0.273 | 2.76
Our Model | $0.215(\pm 0.179)$ | 0.069 | 0.209 | 0.283 | 2.668
Table 2: Summary statistics of cross-entropy losses at each ZIP code. First column is the client-specific cross entropy, second column the first quartile, third column second quartile (median), and fourth column third quartile. In the fifth column we report the maximum client-specific cross-entropy. | Avg. ($\pm$ SD.) | 25% | 50% | 75% | Perf. Ratio
---|---|---|---|---|---
Locally-trained | $0.940(\pm 0.070)$ | 0.919 | 0.947 | 1 | 0.972
Single-model | $0.942(\pm 0.067)$ | 0.921 | 0.948 | 1 | 0.986
Our Model | $0.942(\pm 0.066)$ | 0.921 | 0.948 | 1 | 1
Table 3: Summary statistics of accuracy on each ZIP code. First column is the
average client-specific accuracy, second column the first quartile, third
column second quartile (median), and fourth column third quartile. We define
performance ratio as the proportion of clients with accuracy no less than the
model being compared to, and record it in the fifth column.

Although our model attains slightly higher cross-entropy than training a
single model for all clients, the accuracy of both models are almost the same,
as we can see from Table 3. We conjecture that the higher cross-entropy is due
to a lack of emphasis on tuning the learning rates and optimizing until
convergence for our proposed personalized model.

##  7 Conclusions and Future Directions

We propose a new approach to personalization in federated learning when there
are multiple known clusters among the clients. Our algorithm is based on a
variant of loopless local gradient descent that allows each cluster to have
its own communication schedule. The estimator is shown to be optimal among a
class of unbiased linear estimators and performs better than commonly used
alternatives. We empirically demonstrated our estimator on both simulated and
real-world data.

In the future, we will investigate how to obtain an efficient implementation
of alternative approaches described in Section 5.1. We have identified two
estimators that are costly to obtain, yet outperform our proposed method.
Studying their efficient implementation could yield even better methods for
personalized federated learning.

## Acknowledgments

This work is partially supported by the William S. Fishman Faculty Research
Fund at the University of Chicago Booth School of Business. This work was
completed in part with resources provided by the University of Chicago
Research Computing Center.

## References

  * Abad et al. (2020) M. S. H. Abad, E. Ozfatura, D. Gunduz, and O. Ercetin.  Hierarchical federated learning across heterogeneous cellular networks.  In _ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 8866–8870. IEEE, 2020.
  * Bakker and Heskes (2003) B. Bakker and T. Heskes.  Task clustering and gating for bayesian multitask learning.  _Journal of Machine Learning Research_ , 4:83–99, 2003\.
  * Blattberg et al. (2009) R. C. Blattberg, E. C. Malthouse, and S. A. Neslin.  Customer lifetime value: Empirical generalizations and some conceptual questions.  _Journal of Interactive Marketing_ , 23(2):157–168, 2009.
  * Bock (1975) M. E. Bock.  Minimax estimators of the mean of a multivariate normal distribution.  _The Annals of Statistics_ , pages 209–218, 1975.
  * Bonawitz et al. (2019) K. Bonawitz, H. Eichner, W. Grieskamp, D. Huba, A. Ingerman, V. Ivanov, C. Kiddon, J. Konečnỳ, S. Mazzocchi, H. B. McMahan, et al.  Towards federated learning at scale: System design.  _arXiv preprint arXiv:1902.01046_ , 2019.
  * Briggs et al. (2020) C. Briggs, Z. Fan, and P. Andras.  Federated learning with hierarchical clustering of local updates to improve training on non-iid data.  In _2020 International Joint Conference on Neural Networks (IJCNN)_ , pages 1–9. IEEE, 2020.
  * Bryk and Raudenbush (1987) A. S. Bryk and S. W. Raudenbush.  Application of hierarchical linear models to assessing change.  _Psychological bulletin_ , 101(1):147, 1987.
  * Bryk and Raudenbush (1992) A. S. Bryk and S. W. Raudenbush.  _Hierarchical linear models: Applications and data analysis methods._ Sage Publications, Inc, 1992.
  * Bumbaca et al. (2017) F. Bumbaca, S. Misra, and P. E. Rossi.  Distributed markov chain monte carlo for bayesian hierarchical models.  _Available at SSRN 2964646_ , 2017.
  * Bumbaca et al. (2020) F. Bumbaca, S. Misra, and P. E. Rossi.  Scalable target marketing: Distributed markov chain monte carlo for bayesian hierarchical models.  _Journal of Marketing Research_ , 57(6):999–1018, 2020.
  * Chen et al. (2021) S. Chen, Q. Zheng, Q. Long, and W. J. Su.  A theorem of the alternative for personalized federated learning.  _arXiv preprint arXiv:2103.01901_ , 2021.
  * Daniels and Gatsonis (1999) M. J. Daniels and C. Gatsonis.  Hierarchical generalized linear models in the analysis of variations in health care utilization.  _Journal of the American Statistical Association_ , 94(445):29–42, 1999.
  * Deng et al. (2020) Y. Deng, M. M. Kamani, and M. Mahdavi.  Adaptive personalized federated learning.  _arXiv preprint arXiv:2003.13461_ , 2020.
  * Dinh et al. (2020) C. T. Dinh, N. H. Tran, and T. D. Nguyen.  Personalized federated learning with moreau envelopes.  _arXiv preprint arXiv:2006.08848_ , 2020.
  * Duan and Wang (2022) Y. Duan and K. Wang.  Adaptive and robust multi-task learning.  _arXiv preprint arXiv:2202.05250_ , 2022.
  * Fallah et al. (2020) A. Fallah, A. Mokhtari, and A. Ozdaglar.  Personalized federated learning: A meta-learning approach.  _arXiv preprint arXiv:2002.07948_ , 2020.
  * French and Russell-Bennett (2015) J. French and R. Russell-Bennett.  A hierarchical model of social marketing.  _Journal of Social Marketing_ , 2015.
  * Ghosh et al. (2020) A. Ghosh, J. Chung, D. Yin, and K. Ramchandran.  An efficient framework for clustered federated learning.  _arXiv preprint arXiv:2006.04088_ , 2020.
  * Gower et al. (2019) R. M. Gower, N. Loizou, X. Qian, A. Sailanbayev, E. Shulgin, and P. Richtárik.  Sgd: General analysis and improved rates.  In _International Conference on Machine Learning_ , pages 5200–5209. PMLR, 2019.
  * Hanzely and Richtárik (2020) F. Hanzely and P. Richtárik.  Federated learning of a mixture of global and local models.  _arXiv preprint arXiv:2002.05516_ , 2020.
  * Hanzely et al. (2020a) F. Hanzely, S. Hanzely, S. Horváth, and P. Richtarik.  Lower bounds and optimal algorithms for personalized federated learning.  _Advances in Neural Information Processing Systems_ , 33, 2020a.
  * Hanzely et al. (2020b) F. Hanzely, D. Kovalev, and P. Richtarik.  Variance reduced coordinate descent with acceleration: New method with a surprising application to finite-sum problems.  In _International Conference on Machine Learning_ , pages 4039–4048. PMLR, 2020b.
  * Hanzely et al. (2021) F. Hanzely, B. Zhao, and M. Kolar.  Personalized federated learning: A unified framework and universal optimization techniques.  _arXiv preprint arXiv:2102.09743_ , 2021.
  * Hofmann (1997) D. A. Hofmann.  An overview of the logic and rationale of hierarchical linear models.  _Journal of management_ , 23(6):723–744, 1997\.
  * Hooley et al. (1999) G. Hooley, J. Fahy, T. Cox, J. Beracs, K. Fonfara, and B. Snoj.  Marketing capabilities and firm performance: a hierarchical model.  _Journal of market-focused management_ , 4(3):259–278, 1999.
  * Huang et al. (2019) L. Huang, A. L. Shea, H. Qian, A. Masurkar, H. Deng, and D. Liu.  Patient clustering improves efficiency of federated machine learning to predict mortality and hospital stay time using distributed electronic medical records.  _Journal of biomedical informatics_ , 99:103291, 2019.
  * Jacob et al. (2008) L. Jacob, F. Bach, and J.-P. Vert.  Clustered multi-task learning: A convex formulation.  _arXiv preprint arXiv:0809.2085_ , 2008.
  * James and Stein (1992) W. James and C. Stein.  Estimation with quadratic loss.  In _Breakthroughs in statistics_ , pages 443–460. Springer, 1992\.
  * Kairouz et al. (2019) P. Kairouz, H. B. McMahan, B. Avent, A. Bellet, M. Bennis, A. N. Bhagoji, K. Bonawitz, Z. Charles, G. Cormode, R. Cummings, et al.  Advances and open problems in federated learning.  _arXiv preprint arXiv:1912.04977_ , 2019.
  * Karimireddy et al. (2020) S. P. Karimireddy, S. Kale, M. Mohri, S. Reddi, S. Stich, and A. T. Suresh.  Scaffold: Stochastic controlled averaging for federated learning.  In _International Conference on Machine Learning_ , pages 5132–5143. PMLR, 2020.
  * Kariya and Kurata (2004) T. Kariya and H. Kurata.  _Generalized least squares_.  John Wiley & Sons, 2004.
  * Kovalev et al. (2020) D. Kovalev, S. Horváth, and P. Richtárik.  Don’t jump through hoops and remove those loops: Svrg and katyusha are better without the outer loop.  In _Algorithmic Learning Theory_ , pages 451–467. PMLR, 2020.
  * Kubokawa (1991) T. Kubokawa.  An approach to improving the james-stein estimator.  _Journal of Multivariate Analysis_ , 36(1):121–126, 1991.
  * Kumar and Daume III (2012) A. Kumar and H. Daume III.  Learning task grouping and overlap in multi-task learning.  _arXiv preprint arXiv:1206.6417_ , 2012.
  * Lee and Nelder (1996) Y. Lee and J. A. Nelder.  Hierarchical generalized linear models.  _Journal of the Royal Statistical Society: Series B (Methodological)_ , 58(4):619–656, 1996.
  * Li et al. (2018) T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith.  Federated optimization in heterogeneous networks.  _arXiv preprint arXiv:1812.06127_ , 2018.
  * Li et al. (2020) T. Li, A. K. Sahu, A. Talwalkar, and V. Smith.  Federated learning: Challenges, methods, and future directions.  _IEEE Signal Processing Magazine_ , 37(3):50–60, 2020.
  * Li et al. (2021) T. Li, S. Hu, A. Beirami, and V. Smith.  Ditto: Fair and robust federated learning through personalization.  In _International Conference on Machine Learning_ , pages 6357–6368. PMLR, 2021.
  * Li (2021) Z. Li.  Anita: An optimal loopless accelerated variance-reduced gradient method.  _arXiv preprint arXiv:2103.11333_ , 2021.
  * Liu et al. (2020) L. Liu, J. Zhang, S. Song, and K. B. Letaief.  Client-edge-cloud hierarchical federated learning.  In _ICC 2020-2020 IEEE International Conference on Communications (ICC)_ , pages 1–6. IEEE, 2020.
  * Mansour et al. (2020) Y. Mansour, M. Mohri, J. Ro, and A. T. Suresh.  Three approaches for personalization with applications to federated learning.  _arXiv preprint arXiv:2002.10619_ , 2020.
  * Marfoq et al. (2021) O. Marfoq, G. Neglia, A. Bellet, L. Kameni, and R. Vidal.  Federated multi-task learning under a mixture of distributions.  _arXiv preprint arXiv:2108.10252_ , 2021.
  * McMahan et al. (2017) B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In _Artificial intelligence and statistics_ , pages 1273–1282. PMLR, 2017.
  * Naik and Peters (2009) P. A. Naik and K. Peters.  A hierarchical marketing communications model of online and offline media synergies.  _Journal of Interactive Marketing_ , 23(4):288–299, 2009.
  * Nesterov et al. (2018) Y. Nesterov et al.  _Lectures on convex optimization_ , volume 137.  Springer, 2018.
  * Petersen et al. (2008) K. B. Petersen, M. S. Pedersen, et al.  The matrix cookbook.  _Technical University of Denmark_ , 7(15):510, 2008.
  * Qian et al. (2021a) X. Qian, H. Dong, P. Richtárik, and T. Zhang.  Error compensated loopless svrg, quartz, and sdca for distributed optimization.  _arXiv preprint arXiv:2109.10049_ , 2021a.
  * Qian et al. (2021b) X. Qian, Z. Qu, and P. Richtárik.  L-svrg and l-katyusha with arbitrary sampling.  _Journal of Machine Learning Research_ , 22:1–49, 2021b.
  * Raudenbush and Bryk (1986) S. Raudenbush and A. S. Bryk.  A hierarchical model for studying school effects.  _Sociology of education_ , pages 1–17, 1986.
  * Raudenbush (1988) S. W. Raudenbush.  Educational applications of hierarchical linear models: A review.  _journal of Educational Statistics_ , 13(2):85–116, 1988.
  * Sattler et al. (2020) F. Sattler, K.-R. Müller, and W. Samek.  Clustered federated learning: Model-agnostic distributed multitask optimization under privacy constraints.  _IEEE transactions on neural networks and learning systems_ , 2020\.
  * Smith et al. (2017) V. Smith, C.-K. Chiang, M. Sanjabi, and A. Talwalkar.  Federated multi-task learning.  _arXiv preprint arXiv:1705.10467_ , 2017.
  * Snider (2019) S. Snider.  Where do i fall in the american economic class system.  _US News & World Report_, 2019.
  * Stephen and Anthony (2002) R. Stephen and B. Anthony.  Hierarchical linear models, 2002.
  * Stich (2018) S. U. Stich.  Local sgd converges fast and communicates little.  _arXiv preprint arXiv:1805.09767_ , 2018.
  * Wainakh et al. (2020) A. Wainakh, A. S. Guinea, T. Grube, and M. Mühlhäuser.  Enhancing privacy via hierarchical federated learning.  In _2020 IEEE European Symposium on Security and Privacy Workshops (EuroS &PW)_, pages 344–347. IEEE, 2020.
  * Wang et al. (2020) J. Wang, S. Wang, R.-R. Chen, and M. Ji.  Local averaging helps: Hierarchical federated learning and convergence analysis.  _arXiv preprint arXiv:2010.12998_ , 2020.
  * Wang et al. (2018) W. Wang, J. Wang, M. Kolar, and N. Srebro.  Distributed stochastic multi-task learning with graph regularization.  _arXiv preprint arXiv:1802.03830_ , 2018.
  * Yang et al. (2019) Q. Yang, Y. Liu, Y. Cheng, Y. Kang, T. Chen, and H. Yu.  Federated learning.  _Synthesis Lectures on Artificial Intelligence and Machine Learning_ , 13(3):1–207, 2019.
  * Zhang and Yang (2017) Y. Zhang and Q. Yang.  A survey on multi-task learning.  _arXiv preprint arXiv:1707.08114_ , 2017.
  * Zhang and Yeung (2012) Y. Zhang and D.-Y. Yeung.  A convex formulation for learning task relationships in multi-task learning.  _arXiv preprint arXiv:1203.3536_ , 2012.
  * Zhao et al. (2021) H. Zhao, Z. Li, and P. Richtárik.  Fedpage: A fast local stochastic gradient method for communication-efficient federated learning.  _arXiv preprint arXiv:2108.04755_ , 2021.
  * Zhou et al. (2011a) J. Zhou, J. Chen, and J. Ye.  Clustered multi-task learning via alternating structure optimization.  _Advances in neural information processing systems_ , 2011:702, 2011a.
  * Zhou et al. (2011b) J. Zhou, J. Chen, and J. Ye.  Malsar: Multi-task learning via structural regularization.  _Arizona State University_ , 21, 2011b.

##  Appendix A Missing Lemmas and Proofs

We provide detailed proofs for the results in the main text.

###  A.1 Proof of Proposition 2.2

The gradient of $F(\\{\theta_{i}\\}_{i=1}^{n})$ w.r.t. $\theta_{i_{0}}$ for
any $i_{0}\in\\{1,\ldots,n\\}$ is given by

$\nabla_{\theta_{i_{0}}}F(\\{\theta_{i}\\}_{i=1}^{n})=\nabla
f_{i_{0}}(\theta_{i_{0}})+\alpha_{j_{0}}\gamma_{i_{0}}(\theta_{i_{0}}-\bar{\theta})+(1-\alpha_{j_{0}})\gamma_{i_{0}}(\theta_{i_{0}}-\bar{\theta}_{j_{0}}),$

where we let $j_{0}$ denote the cluster client $i_{0}$ belongs to. The
gradients of Eq. 2.3 w.r.t. $\\{w_{j}\\}_{j=1}^{k},\bar{w}$ are given by

$\displaystyle\nabla_{w_{j}}F_{MTL}(\\{\theta_{i}\\}_{i=1}^{n},\\{w_{j}\\}_{j=1}^{k},\bar{w})=\sum_{i\in\mathcal{I}_{j}}(\theta_{i}-w_{j})+\lambda_{j}(w_{j}-\bar{w}),\,j=1,\ldots,k,$
$\displaystyle\nabla_{\bar{w}}F_{MTL}(\\{\theta_{i}\\}_{i=1}^{n},\\{w_{j}\\}_{j=1}^{k},\bar{w})=\sum_{j=1}^{k}\lambda_{j}(w_{j}-\bar{w}).$

Setting the gradients to zero and solving the resulting linear system, give us

$\displaystyle
w_{j}=\frac{\lambda_{j}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}\bar{w}+\frac{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\theta_{i}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}\gamma_{i}},\qquad
j=1,\ldots,k,$
$\displaystyle\bar{w}=\left(\sum_{j=1}^{k}\frac{\lambda_{j}\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}\right)^{-1}\left(\sum_{j=1}^{k}\frac{\lambda_{j}\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\theta_{i}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}\right).$

Plugging
$\alpha_{j}=\frac{\lambda_{j}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}$
and recalling Eq. 2.2, we note that

$\bar{w}=\frac{\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}\theta_{i}}{\sum_{j=1}^{k}\sum_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}}=\bar{\theta},$

and

$w_{j}=\alpha_{j}\bar{\theta}+(1-\alpha_{j})\bar{\theta}_{j},\qquad
j=1,\ldots,k.$

Therefore, the first order condition for any $\theta_{i_{0}}$ is

$\nabla
f_{i_{0}}(\theta_{i_{0}})+\alpha_{j_{0}}\gamma_{i_{0}}(\theta_{i_{0}}-\bar{\theta})+(1-\alpha_{j_{0}})\gamma_{i_{0}}(\theta_{i_{0}}-\bar{\theta}_{j_{0}})=0,$

which is exactly the first order condition
$\nabla_{\theta_{i}}F(\\{\theta_{i}\\}_{i=1}^{n})=0$. The proof is now
complete.

###  A.2 Proof of Proposition 3.1

We prove the first statement for some $j$ and $i\in\mathcal{I}_{j}$. Since

$\nabla_{\theta_{i}}\bar{\theta}_{j}=\frac{\gamma_{i}}{\sum_{i^{\prime}\in\mathcal{I}_{j}}\gamma_{i^{\prime}}},$

we have

$\displaystyle\nabla_{\theta_{i}}\psi_{j}\left(\bm{\theta}_{j};\left\\{\gamma_{i}\right\\}_{i\in\mathcal{I}_{j}}\right)$
$\displaystyle=\gamma_{i}\left(\theta_{i}-\bar{\theta}_{j}\right)\left(1-\frac{\gamma_{i}}{\sum_{i^{\prime}\in\mathcal{I}_{j}}\gamma_{i^{\prime}}}\right)$
$\displaystyle\quad+\sum_{i^{\prime}\in\mathcal{I}_{j}\backslash\\{i\\}}\gamma_{i^{\prime}}\left(\theta_{i^{\prime}}-\bar{\theta}_{j}\right)\left(-\frac{\gamma_{i}}{\sum_{i^{\prime}\neq\mathcal{I}_{j}}\gamma_{i^{\prime}}}\right)$
$\displaystyle=\gamma_{i}\left(\theta_{i}-\bar{\theta}_{j}\right)-\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{\gamma_{i^{\prime}}}{\sum_{i^{\prime\prime}\in\mathcal{I}_{j}}\gamma_{i^{\prime\prime}}}\left(\theta_{i^{\prime}}-\bar{\theta}_{j}\right)$
$\displaystyle=\gamma_{i}(\theta_{i}-\bar{\theta}_{j}).$

The second statement is proven similarly.

###  A.3 Proof of Proposition 3.3

Note that the gradient oracle, $G(\bm{\theta})$, can be written as

$\displaystyle G(\bm{\theta})$
$\displaystyle=\frac{\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{0}=1\\}}{p_{0}}\left(\nabla_{\bm{\theta}}\varphi(\bm{\theta})+\begin{bmatrix}(1-\alpha_{1})\tau_{1}\nabla_{\bm{\theta}_{1}}\psi_{i}(\bm{\theta}_{1})\\\
(1-\alpha_{2})\tau_{2}\nabla_{\bm{\theta}_{2}}\psi_{2}(\bm{\theta}_{2})\\\
\vdots\\\
(1-\alpha_{k})\tau_{k}\nabla_{\bm{\theta}_{k}}\psi_{i}(\bm{\theta}_{k})\end{bmatrix}\right)+$
$\displaystyle\quad\quad\frac{\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{0}=0\\}}{1-p_{0}}\begin{bmatrix}\frac{\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{1}=1\\}}{p_{1}}(1-\alpha_{1})(1-\tau_{1})\nabla_{\bm{\theta}_{1}}\psi_{1}(\bm{\theta}_{1})+\frac{\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{1}=0\\}}{1-p_{1}}\nabla_{\bm{\theta}_{1}}F_{1}(\bm{\theta}_{1})\\\
\frac{\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{2}=1\\}}{p_{2}}(1-\alpha_{2})(1-\tau_{2})\nabla_{\bm{\theta}_{2}}\psi_{2}(\bm{\theta}_{2})+\frac{\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{2}=0\\}}{1-p_{2}}\nabla_{\bm{\theta}_{2}}F_{2}(\bm{\theta}_{2})\\\
\vdots\\\
\frac{\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{k}=1\\}}{p_{k}}(1-\alpha_{k})(1-\tau_{k})\nabla_{\bm{\theta}_{k}}\psi_{k}(\bm{\theta}_{k})+\frac{\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{k}=0\\}}{1-p_{k}}\nabla_{\bm{\theta}_{k}}F_{k}(\bm{\theta}_{k})\end{bmatrix}.$

We bound the two terms separately using the law of total expectation. Note
that

$\displaystyle\mathbb{E}$
$\displaystyle\left[\|G(\bm{\theta})-G(\bm{\theta}^{*})\|^{2}|\xi_{0}=1\right]\mathbb{P}(\xi_{0}=1)$
$\displaystyle=\frac{1}{p_{0}}\left\|\nabla_{\bm{\theta}}\varphi(\bm{\theta})-\nabla_{\bm{\theta}}\varphi(\bm{\theta}^{*})+\begin{bmatrix}(1-\alpha_{1})\tau_{1}(\nabla_{\bm{\theta}_{1}}\psi_{i}(\bm{\theta}_{1})-\nabla_{\bm{\theta}_{1}}\psi_{1}(\bm{\theta}_{1}^{*}))\\\
(1-\alpha_{2})\tau_{2}(\nabla_{\bm{\theta}_{2}}\psi_{2}(\bm{\theta}_{2})-\nabla_{\bm{\theta}_{2}}\psi_{2}(\bm{\theta}_{2}^{*}))\\\
\vdots\\\
(1-\alpha_{k})\tau_{k}(\nabla_{\bm{\theta}_{k}}\psi_{i}(\bm{\theta}_{k})-\nabla_{\bm{\theta}_{k}}\psi_{k}(\bm{\theta}_{k}^{*}))\end{bmatrix}\right\|^{2}$
$\displaystyle\leq\frac{2}{p_{0}}\left\|\nabla_{\bm{\theta}}\varphi(\bm{\theta})-\nabla_{\bm{\theta}}\varphi(\bm{\theta}^{*})\right\|^{2}+\frac{2}{p_{0}}\left\|\begin{bmatrix}(1-\alpha_{1})\tau_{1}(\nabla_{\bm{\theta}_{1}}\psi_{i}(\bm{\theta}_{1})-\nabla_{\bm{\theta}_{1}}\psi_{1}(\bm{\theta}_{1}^{*}))\\\
(1-\alpha_{2})\tau_{2}(\nabla_{\bm{\theta}_{2}}\psi_{2}(\bm{\theta}_{2})-\nabla_{\bm{\theta}_{2}}\psi_{2}(\bm{\theta}_{2}^{*}))\\\
\vdots\\\
(1-\alpha_{k})\tau_{k}(\nabla_{\bm{\theta}_{k}}\psi_{i}(\bm{\theta}_{k})-\nabla_{\bm{\theta}_{k}}\psi_{k}(\bm{\theta}_{k}^{*}))\end{bmatrix}\right\|^{2}$
$\displaystyle=\frac{2}{p_{0}}\left\|\nabla_{\bm{\theta}}\varphi(\bm{\theta})-\nabla_{\bm{\theta}}\varphi(\bm{\theta}^{*})\right\|^{2}+\frac{2}{p_{0}}\sum_{j=1}^{k}(1-\alpha_{j})^{2}\tau_{j}^{2}\|\nabla_{\bm{\theta}_{1}}\psi_{i}(\bm{\theta}_{1})-\nabla_{\bm{\theta}_{1}}\psi_{1}(\bm{\theta}_{1}^{*})\|^{2},$

where the second line uses the fact that $\mathbb{E}$ is taken over
$\\{\xi_{j}\\}_{j=0}^{k}$, the third line uses the triangle inequality, and
the fourth line follows directly using the fact that $\|\cdot\|^{2}$ is the
squared $\ell_{2}$-norm.

Similarly, we have

$\displaystyle\mathbb{E}\left[\|G(\bm{\theta})-G(\bm{\theta}^{*})\|^{2}|\xi_{0}=0\right]\mathbb{P}(\xi_{0}=0)=\frac{1}{1-p_{0}}\sum_{j=1}^{k}\mathbb{E}_{\xi_{j}}{M_{j}},$

where

$\displaystyle M_{j}=$
$\displaystyle\Biggl{\|}\frac{\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{j}=1\\}}{p_{j}}(1-\alpha_{j})(1-\tau_{j})(\nabla_{\bm{\theta}_{j}}\psi_{j}(\bm{\theta}_{j})-\nabla_{\bm{\theta}_{j}}\psi_{j}(\bm{\theta}_{j}^{*}))$
$\displaystyle\hskip
80.00012pt+\frac{\operatorname*{1{\hskip-2.5pt}\hbox{$\rm{I}$}}\\{\xi_{j}=0\\}}{1-p_{j}}(\nabla_{\bm{\theta}_{j}}F_{j}(\bm{\theta}_{j})-\nabla_{\bm{\theta}_{j}}F_{j}(\bm{\theta}_{j}^{*}))\Biggr{\|}^{2}.$

For a fixed $j$, we have

$\displaystyle\mathbb{E}_{\xi_{j}}[M_{j}]$
$\displaystyle=\mathbb{E}_{\xi_{j}}[M_{j}|\xi_{j}=1]\mathbb{P}(\xi_{j}=1)+\mathbb{E}_{\xi_{j}}[M_{j}|\xi_{j}=0]\mathbb{P}(\xi_{j}=0)$
$\displaystyle=\frac{1}{p_{j}}(1-\alpha_{j})^{2}(1-\tau_{j})^{2}\|\nabla_{\bm{\theta}_{j}}\psi_{j}(\bm{\theta}_{j})-\nabla_{\bm{\theta}_{j}}\psi_{j}(\bm{\theta}_{j}^{*})\|^{2}$
$\displaystyle\qquad+\frac{1}{1-p_{j}}\|\nabla_{\bm{\theta}_{j}}F_{j}(\bm{\theta}_{j})-\nabla_{\bm{\theta}_{j}}F_{j}(\bm{\theta}_{j}^{*})\|^{2}.$

Combining the two bounds, we further have

$\displaystyle\mathbb{E}$
$\displaystyle[\|G(\bm{\theta})-G(\bm{\theta}^{*})\|^{2}]$
$\displaystyle=\mathbb{E}\left[\|G(\bm{\theta})-G(\bm{\theta}^{*})\|^{2}|\xi_{0}=1\right]\mathbb{P}(\xi_{0}=1)+\mathbb{E}\left[\|G(\bm{\theta})-G(\bm{\theta}^{*})\|^{2}|\xi_{0}=0\right]\mathbb{P}(\xi_{0}=0)$
$\displaystyle\leq\frac{2}{p_{0}}\left\|\nabla_{\bm{\theta}}\varphi(\bm{\theta})-\nabla_{\bm{\theta}}\varphi(\bm{\theta}^{*})\right\|^{2}$
$\displaystyle\qquad+\sum_{j=1}^{k}(1-\alpha_{j})^{2}\left(\frac{2}{p_{0}}\tau_{j}^{2}+\frac{(1-\tau_{j})^{2}}{p_{j}(1-p_{0})}\right)\|\nabla_{\bm{\theta}_{j}}\psi_{j}(\bm{\theta}_{j})-\nabla_{\bm{\theta}_{j}}\psi_{j}(\bm{\theta}_{j}^{*})\|^{2}$
$\displaystyle\qquad+\sum_{j=1}^{k}\frac{1}{1-p_{j}}\|\nabla_{\bm{\theta}_{j}}F_{j}(\bm{\theta}_{j})-\nabla_{\bm{\theta}_{j}}F_{j}(\bm{\theta}_{j}^{*})\|^{2}.$

In the above display, only the second term depends on $\tau_{j}$. This term is
minimized when

$\tau_{j}=\frac{p_{0}}{p_{0}+2(1-p_{0})p_{j}}.$

The result immediately follows now.

###  A.4 Proof of Proposition 3.4

We prove the first statement for a fixed $j$. From the proof of Proposition
3.1, we can write the gradient of
$\psi_{j}\left(\bm{\theta}_{j};\left\\{\gamma_{i}\right\\}_{i\in\mathcal{I}_{j}}\right)$
w.r.t. $\bm{\theta}_{j}$ as

$\nabla_{\bm{\theta}_{j}}\psi_{j}\left(\bm{\theta}_{j};\left\\{\gamma_{i}\right\\}_{i\in\mathcal{I}_{j}}\right)=\bm{\theta}_{j}-\begin{bmatrix}\bar{\theta}_{j}\\\
\bar{\theta}_{j}\\\ \vdots\\\ \bar{\theta}_{j}\end{bmatrix}$

and the Hessian as
$\nabla_{\bm{\theta}_{j}\bm{\theta}_{j}}\psi_{j}\left(\bm{\theta}_{j};\left\\{\gamma_{i}\right\\}_{i\in\mathcal{I}_{j}}\right)=I_{d}\otimes
H_{j}$, where $H_{j}\in\mathbb{R}^{|\mathcal{I}_{j}|\times|\mathcal{I}_{j}|}$
is a symmetric matrix with element in row $i_{1}$ and column $i_{2}$
($i_{1},i_{2}\in\mathcal{I}_{j}$) denoted as

$H_{j}[i_{1},i_{2}]=\begin{cases}\gamma_{i_{1}}\left(1-\frac{\gamma_{i_{1}}}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}}\right),\quad&\text{if
$i_{1}=i_{2}$},\\\
-\frac{\gamma_{i_{1}}\gamma_{i_{2}}}{\sum_{i=1\in\mathcal{I}_{j}}\gamma_{i}},\quad&\text{otherwise}.\end{cases}$

For an arbitrary fixed vector $v\in\mathbb{R}^{|\mathcal{I}_{j}|}$, we have

$v^{T}H_{j}v=\sum_{l=1}^{|\mathcal{I}_{j}|}\gamma_{\mathcal{I}_{j}[l]}(v[l])^{2}-\frac{1}{\sum_{l=1}^{|\mathcal{I}_{j}|}\gamma_{\mathcal{I}_{j}[l]}}\left(\sum_{l=1}^{|\mathcal{I}_{j}|}\gamma_{\mathcal{I}_{j}[l]}(v[l])\right)^{2}\\\
=\left(\frac{1}{\sum_{l=1}^{|\mathcal{I}_{j}|}\gamma_{\mathcal{I}_{j}[l]}}\right)\left(\sum_{l=1}^{|\mathcal{I}_{j}|}\frac{\gamma_{\mathcal{I}_{j}[l]}}{\sum_{l^{\prime}=1}^{|\mathcal{I}_{j}|}\gamma_{\mathcal{I}_{j}[l^{\prime}]}}(v[l])^{2}-\left(\sum_{l=1}^{|\mathcal{I}_{j}|}\frac{\gamma_{\mathcal{I}_{j}[l]}}{\sum_{l^{\prime}=1}^{|\mathcal{I}_{j}|}\gamma_{\mathcal{I}_{j}[l^{\prime}]}}v[l]\right)^{2}\right)\geq
0,$

which shows that $H_{j}\succeq 0$. Let
$\Gamma_{j}\in\mathbb{R}^{|\mathcal{I}_{j}|\times|\mathcal{I}_{J}|}$ be a
diagonal matrix with diagonal elements
$\\{\gamma_{\mathcal{I}_{j}[l]}\\}_{l=1}^{|\mathcal{I}_{j}|}$. Then

$\displaystyle
v^{T}\Gamma_{j}v-v^{T}H_{j}v=\frac{1}{\sum_{l=1}^{|\mathcal{I}_{j}|}\gamma_{\mathcal{I}_{j}[l]}}\left(\sum_{l=1}^{|\mathcal{I}_{j}|}\gamma_{\mathcal{I}_{j}[l]}(v[l])\right)^{2}\geq
0,$

which shows that
$H_{j}\preceq\Gamma_{j}\preceq\max_{i\in\mathcal{I}_{j}}\gamma_{i}I_{|\mathcal{I}_{J}|}$.
The first statement now follows, since the eigenvalues of the Kronecker
product of two matrices are the products of the pairs of eigenvalues of two
matrices.

The second statement is established in the same way.

###  A.5 Proof of Theorem 3.5

First, we show that the gradient estimator defined in Eq. 3.2 satisfies the
expected smoothness condition. We then bound the gradient estimator’s second
moment, and apply Theorem 3.1 from Gower et al. (2019) to complete the proof.

######  Lemma A.1 (Expected Smoothness).

Suppose conditions of Theorem 3.5 hold and $\mathcal{L}$ is defined in Theorem
3.5. Then

$\mathbb{E}[\|G(\bm{\theta})-G(\bm{\theta}^{*})\|^{2}]\leq
2\mathcal{L}\left(F(\bm{\theta})-F(\bm{\theta}^{*})\right),$

where the expectation is taken over the randomness in
$\\{\xi_{j}\\}_{j=0}^{k}$.

###### Proof.

For a convex function $g:\mathbb{R}^{d}\to\mathbb{R}$, let
$D_{g}(x,y)=g(x)-g(y)-(\nabla g(y))^{T}(x-y)$, $x,y\in\mathbb{R}^{d}$. By
Proposition 3.4 and Theorem B.1, we have

$\displaystyle\left\|\nabla_{\bm{\theta}}\varphi(\bm{\theta})-\nabla_{\bm{\theta}}\varphi(\bm{\theta}^{*})\right\|^{2}$
$\displaystyle\leq
2\max_{j=1,\ldots,k}\max_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}D_{\varphi}(\bm{\theta},\bm{\theta}^{*}),$
$\displaystyle\left\|\nabla_{\bm{\theta}_{j}}\psi_{j}(\bm{\theta}_{j})-\nabla_{\bm{\theta}_{j}}\psi_{j}(\bm{\theta}_{j}^{*})\right\|^{2}$
$\displaystyle\leq
2\max_{i\in\mathcal{I}_{j}}\gamma_{i}D_{\psi_{j}}(\bm{\theta}_{j},\bm{\theta}_{j}^{*}),$

while 2.1 states that

$\|\nabla_{\bm{\theta}_{j}}F_{j}(\bm{\theta}_{j})-\nabla_{\bm{\theta}_{j}}F_{j}(\bm{\theta}_{j}^{*})\|^{2}\leq
2LD_{F_{j}}(\bm{\theta}_{j},\bm{\theta}^{*}_{j}).$

Plugging into the result of Proposition 3.3, we have

$\displaystyle\mathbb{E}[\|G(\bm{\theta})-G(\bm{\theta}^{*})\|^{2}]$
$\displaystyle\leq\frac{4}{p_{0}}\max_{j=1,\ldots,k}\max_{i\in\mathcal{I}_{j}}\alpha_{j}\gamma_{i}D_{\varphi}(\bm{\theta},\bm{\theta}^{*})$
$\displaystyle\qquad+\sum_{j=1}^{k}\frac{4(1-\alpha_{j})^{2}\max_{i\in\mathcal{I}_{j}}\gamma_{i}}{p_{0}+2(1-p_{0})p_{j}}D_{\psi_{j}}(\bm{\theta}_{j},\bm{\theta}_{j}^{*})$
$\displaystyle\qquad+\frac{2L}{1-p_{0}}\sum_{j=1}^{k}\frac{1}{1-p_{j}}D_{F_{j}}(\bm{\theta}_{j},\bm{\theta}^{*}_{j}).$

Since

$F(\bm{\theta})-F(\bm{\theta}^{*})=D_{F}(\bm{\theta},\bm{\theta}^{*})\\\
=D_{\varphi}(\bm{\theta},\bm{\theta}^{*})+\sum_{j=1}^{k}(1-\alpha_{j})D_{\psi_{j}}(\bm{\theta}_{j},\bm{\theta}_{j}^{*})+\sum_{j=1}^{k}D_{F_{j}}(\bm{\theta}_{j},\bm{\theta}_{j}^{*}),$

we have

$\displaystyle\mathbb{E}$
$\displaystyle[\|G(\bm{\theta})-G(\bm{\theta}^{*})\|^{2}]$ $\displaystyle\leq
2\mathcal{L}\left(D_{\varphi}(\bm{\theta},\bm{\theta}^{*})+\sum_{j=1}^{k}(1-\alpha_{j})D_{\psi_{j}}(\bm{\theta}_{j},\bm{\theta}_{j}^{*})+\sum_{j=1}^{k}D_{F_{j}}(\bm{\theta}_{j},\bm{\theta}_{j}^{*})\right)$
$\displaystyle=2\mathcal{L}(F(\bm{\theta})-F(\bm{\theta}^{*})),$

which completes the proof. ∎

######  Corollary A.2 (Bounded Second Moment).

Suppose conditions of Theorem 3.5 hold and $\mathcal{L}$ and
$\sigma_{\bm{\theta}^{*}}^{2}$ are defined in Theorem 3.5. Then

$\mathbb{E}[\|G(\bm{\theta})\|^{2}]\leq
4\mathcal{L}\mathbb{E}[(F(\bm{\theta})-F(\bm{\theta}^{*})]+2\sigma_{\bm{\theta}^{*}}^{2},$

where the expectation is taken over the randomness in
$\\{\xi_{j}\\}_{j=0}^{k}$.

###### Proof.

From the proof of Proposition 3.3, we have
$\mathbb{E}[\|G(\bm{\theta}^{*})\|^{2}]\leq\sigma_{\bm{\theta}^{*}}^{2}$. The
result follows from Lemma 2.4 of Gower et al. (2019). ∎

###  A.6 Proof of Corollary 3.6

When $C_{2}>C_{1}$, $\tilde{\mathcal{L}}$ is minimized when
$p_{0}=\frac{2C_{1}}{C_{1}+C_{2}+L}$ and
$p_{j}=\frac{C_{2}-C_{1}}{C_{2}-C_{1}+L}$. When $C_{2}\leq C_{1}$,
$\tilde{\mathcal{L}}$ is minimized when $p_{0}=\frac{2C_{1}}{2C_{1}+L}$ and
$p_{j}=0$ for all $j$. Then

$\mathcal{L}\leq\tilde{\mathcal{L}}=\begin{cases}C_{1}+C_{2}+L&\text{ if
$C_{2}>C_{1}$},\\\ 2C_{1}+L&\text{ if $C_{2}\leq C_{1}$}.\end{cases}$

With the upperbounds on $\mathcal{L}$ determined for when $C_{2}>C_{1}$ and
when $C_{2}\leq C_{1}$, we are left with deriving the optimal $\tau_{j}$ under
the two settings. For the choices of $\\{p_{j}\\}_{j=0}^{k}$ we have derived
for the two settings, by Proposition 3.3, we know that when $C_{2}>C_{1}$ we
should set $\tau_{j}=\frac{C_{1}}{C_{2}}$ for all $j$ and when $C_{2}\leq
C_{1}$ we set $\tau_{j}=1$ for all $j$.

For any $\epsilon>0$, set $\eta=(2\mathcal{L})^{-1}$ and
$t={2\mathcal{L}}/{\mu}\log({1}/{\epsilon})$. Then

$(1-\eta\mu)^{t}\leq\exp\left\\{-t\eta\mu\right\\}=\epsilon.$

By Proposition 3.2, we have the following.

  1. 1.

When $C_{2}>C_{1}$, $\mathcal{L}\leq C_{1}+C_{2}+L$ and
$t\leq\frac{2(C_{1}+C_{2}+L)}{\mu}\log\frac{1}{\epsilon}$. Since
$p_{0}=\frac{2C_{1}}{C_{1}+C_{2}+L}$,
$p_{j}=\frac{C_{2}-C_{1}}{C_{2}-C_{1}+L}$, the expected number of between-
cluster communication is at most
$\frac{4C_{1}(C_{2}-C_{1}+L)}{(C_{1}+C_{2}+L)\mu}\log\frac{1}{\epsilon}$, and
the expected number of within-cluster communication is at most
$\frac{2L(C_{2}-C_{1})}{(C_{2}-C_{1}+L)\mu}\log\frac{1}{\epsilon}$.

  2. 2.

When $C_{2}\leq C_{1}$, we have $\mathcal{L}\leq 2C_{1}+L$, therefore
$t\leq\frac{2(2C_{1}+L)}{\mu}$. Recalling the choices for $p_{0}$, the number
of between-cluster communication is at most
$\frac{4C_{1}L}{(2C_{1}+L)\mu}\log\frac{1}{\epsilon}$.

###  A.7 Proof of Lemma 4.2

Similarly to the proof of Lemma A.1, we start by conditioning
$\mathbb{E}[\|\bm{g}^{t}-\nabla F(x^{t})\|^{2}]$ on $\xi_{0}=1$ and
$\xi_{0}=0$, respectively. For $\mathbb{E}[\|\bm{g}^{t}-\nabla
F(x^{t})\|^{2}\mid\xi_{0}=0]$, we further expand
$\mathbb{E}[\|g^{t}_{i}-\nabla_{i}F(x^{t})\|^{2}]$ and condition on
$\xi_{j}=1$ and $\xi_{j}=0$, $j=1,\ldots,k$, $i\in\mathcal{I}_{j}$. These
conditional expectations can then be bounded by Bregman divergences, which
completes the proof.

###  A.8 Proof of Theorem 4.3

Similarly to the proof of Proposition 3.4, we can show that
$\sum_{j=1}^{k}\psi_{j}\left(\bm{\theta}_{j};\left\\{\gamma_{i}\right\\}_{i\in\mathcal{I}_{j}}\right)+\varphi\left(\bm{\theta};\left\\{\gamma_{i}\right\\}_{i=1}^{n},\left\\{\alpha\right\\}_{j=1}^{k}\right)$
is $\max_{i=1,\ldots,n}\gamma_{i}$-smooth and convex in $\bm{\theta}$. The
function $F(\bm{\theta})$ is $\tilde{L}+\max_{i=1,\ldots,n}\gamma_{i}$-smooth
and $\mu$-strongly convex in $\bm{\theta}$ under 4.1.

Algorithm 2 is a special instance of L-Katyusha Hanzely et al. (2020b). For
each $t=1,2,\ldots,$ we obtain an unbiased stochastic gradient estimate
$g_{i}^{t}$ for each client. The local updates of the clients follow the form
of the updates in L-Katyusha. The random variable $\xi^{\prime}$ then controls
how often the algorithm updates the full gradient. Plugging the expected
smoothness of the stochastic gradient oracle, given in Lemma 4.2, into Theorem
4.1 from Hanzely et al. (2020b) completes the proof.

###  A.9 Proof of Corollary 4.4

The proof is similar to the proof of Corollary 3.6 and is omitted.

###  A.10 Proof of Theorem 5.1

Under our model, we have

$\widehat{\theta}_{i}^{d}=\theta_{i}^{*}+\frac{1}{\beta_{i}}X_{i}^{T}\epsilon_{i}=\bar{\theta}_{j}^{*}+\xi_{i}+\frac{1}{\beta_{i}}X_{i}^{T}\epsilon_{i}=\bar{\theta}^{*}+\bar{\xi}_{j}+\xi_{i}+\frac{1}{\beta_{i}}X_{i}^{T}\epsilon_{i}.$

By a direct calculation, solution to Eq. 5.2 can be written as

$\displaystyle\widehat{\theta}_{i}$
$\displaystyle=\frac{\beta_{i}/\sigma_{i}^{2}}{\beta_{i}/\sigma_{i}^{2}+\gamma_{i}}\widehat{\theta}_{i}^{d}+\frac{\gamma_{i}}{\beta_{i}/\sigma_{i}^{2}+\gamma_{i}}\widehat{w}_{j},$
$\displaystyle\widehat{w}_{j}$
$\displaystyle=\frac{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}\widehat{\theta}_{i}}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}+\lambda_{j}}+\frac{\lambda_{j}}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}+\lambda_{j}}\widehat{\bar{w}},$
$\displaystyle\widehat{\bar{w}}$
$\displaystyle=\frac{1}{\sum_{j=1}^{k}\lambda_{j}}\sum_{j=1}^{k}\lambda_{j}\widehat{w}_{j}.$

Expanding the right hand side of $\widehat{w}_{j}$, we have

$\displaystyle\widehat{w}_{j}$
$\displaystyle=\frac{1}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}+\lambda_{j}}\left(\sum_{i\in\mathcal{I}_{j}}\frac{\gamma_{i}\beta_{i}/\sigma_{i}^{2}}{\beta_{i}/\sigma_{i}^{2}+\gamma_{i}}\widehat{\theta}_{i}^{d}+\sum_{i\in\mathcal{I}_{j}}\frac{\gamma_{i}^{2}}{\beta_{i}/\sigma_{i}^{2}+\gamma_{i}}\widehat{w}_{j}\right)+\frac{\lambda_{j}}{\sum_{i\in\mathcal{I}_{j}}\gamma_{i}+\lambda_{j}}\widehat{\bar{w}}.$

Let

$C_{i}=\frac{\gamma_{i}\beta_{i}/\sigma_{i}^{2}}{\beta_{i}/\sigma_{i}^{2}+\gamma_{i}}\qquad\text{and}\qquad\widehat{\bar{\theta}}_{j}=\left(\sum_{i\in\mathcal{I}_{j}}C_{i}\right)^{-1}\left(\sum_{i\in\mathcal{I}_{j}}C_{i}\widehat{\theta}_{i}^{d}\right).$

With this notation, we have

$\displaystyle\widehat{w}_{j}$
$\displaystyle=\frac{\sum_{i\in\mathcal{I}_{j}}C_{i}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}C_{i}}\widehat{\bar{\theta}}_{j}+\frac{\lambda_{j}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}C_{i}}\widehat{\bar{w}}$
and $\displaystyle\widehat{\bar{w}}$
$\displaystyle=\frac{1}{\sum_{j=1}^{k}\lambda_{j}}\sum_{j=1}^{k}\lambda_{j}\left(\frac{\sum_{i\in\mathcal{I}_{j}}C_{i}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}C_{i}}\widehat{\bar{\theta}}_{j}+\frac{\lambda_{j}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}C_{i}}\widehat{\bar{w}}\right).$

Let

$D_{j}=\frac{\lambda_{j}\sum_{i\in\mathcal{I}_{j}}C_{i}}{\lambda_{j}+\sum_{i\in\mathcal{I}_{j}}C_{i}}.$

Then

$\widehat{\bar{w}}=\left(\sum_{j=1}^{k}D_{j}\right)^{-1}\sum_{j=1}^{k}D_{j}\widehat{\bar{\theta}}_{j}.$

Without loss of generality, $\mathcal{I}_{1}=\\{1,\dots,|\mathcal{I}_{1}|\\}$,
$\mathcal{I}_{2}=\\{|\mathcal{I}_{1}|+1,|\mathcal{I}_{1}|+2,\dots,|\mathcal{I}_{1}|+|\mathcal{I}_{2}|\\}$,
and so on. For all $i^{\prime}\in\mathcal{I}_{1}$, $i^{\prime}\neq 1$, we have

$\widehat{\theta}_{i^{\prime}}^{d}=\theta_{1}^{*}-\xi_{1}+\xi_{i^{\prime}}+\frac{1}{\beta_{i^{\prime}}}X_{i^{\prime}}^{T}\epsilon_{i^{\prime}}.$

For all $i^{\prime}\in\mathcal{I}_{j^{\prime}},j^{\prime}\neq 1$, we have

$\widehat{\theta}_{i^{\prime}}^{d}=\theta_{1}^{*}-\xi_{1}-\bar{\xi}_{1}+\bar{\xi}_{j^{\prime}}+\xi_{i^{\prime}}+\frac{1}{\beta_{i^{\prime}}}X_{i^{\prime}}^{T}\epsilon_{i^{\prime}}.$

Then

$\begin{bmatrix}y_{1}\\\ \widehat{\theta}_{2}^{d}\\\ \vdots\\\
\widehat{\theta}_{m}^{d}\end{bmatrix}=\begin{bmatrix}X_{1}\\\ I\\\ \vdots\\\
I\end{bmatrix}\theta_{1}^{*}+\underbrace{\begin{bmatrix}\epsilon_{1}\\\
-\xi_{1}+\xi_{2}+\frac{1}{\beta_{2}}X_{2}^{T}\epsilon_{2}\\\ \vdots\\\
-\xi_{1}-\bar{\xi}_{1}+\bar{\xi}_{k}+\xi_{n}+\frac{1}{\beta_{n}}X_{n}^{T}\epsilon_{n}.\end{bmatrix}}_{\bm{\zeta}_{1}},$
(A.1)

where

$\bm{\zeta}_{1}\sim N\left(0,\begin{bmatrix}\sigma_{1}^{2}&\bm{0}_{n-1}^{T}\\\
\bm{0}_{n-1}&\Omega_{1}\end{bmatrix}\otimes I_{d}\right)$

and for any $i,i^{\prime}\in\\{1,\ldots,n\\}$,

$\Omega_{1}[i,i^{\prime}]=\begin{cases}2\bar{\sigma}_{1}^{2}+\sigma_{i}^{2}/\beta_{i}\quad&\text{if
$i=i^{\prime}\in\mathcal{I}_{1}$}\\\
\bar{\sigma}_{1}^{2}+\bar{\sigma}_{j}^{2}+2\bar{\sigma}^{2}+\sigma_{i}^{2}/\beta_{i}&\text{
if $i=i^{\prime}\in\mathcal{I}_{j}$, $j\neq 1$}\\\ \bar{\sigma}_{1}^{2}&\text{
if $i\in\mathcal{I}_{1}$ or $i^{\prime}\in\mathcal{I}_{1}$, $i\neq
i^{\prime}$}\\\ \bar{\sigma}_{1}^{2}+\bar{\sigma}^{2}&\text{ if
$i\in\mathcal{I}_{j},i^{\prime}\not\in\mathcal{I}_{j}$,
$i,i^{\prime}\not\in\mathcal{I}_{1}$, $i\neq i^{\prime}$}\\\
\bar{\sigma}_{1}^{2}+2\bar{\sigma}^{2}&\text{ if
$i,i^{\prime}\in\mathcal{I}_{j},j\neq 1$, $i\neq i^{\prime}$}\end{cases}.$

The matrix $\Omega_{1}$ can be expressed as

$\Omega_{1}=\begin{bmatrix}{\rm
diag}(\\{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2}\\}_{i=2}^{|\mathcal{I}_{1}|})&\bm{0}_{|\mathcal{I}_{1}|-1}\bm{0}_{n-|\mathcal{I}_{1}|}^{T}\\\
\bm{0}_{n-|\mathcal{I}_{1}|}\bm{0}_{|\mathcal{I}_{1}|-1}^{T}&\Omega_{-1}\end{bmatrix}+\bar{\sigma}_{1}^{2}\mathbf{1}_{n-1}\mathbf{1}_{n-1}^{T},$

where

$\Omega_{-1}=\begin{bmatrix}\Omega_{-1}^{(2)}&\bm{0}_{|\mathcal{I}_{2}|}\bm{0}_{|\mathcal{I}_{3}|}^{T}&\dots&\bm{0}_{|\mathcal{I}_{2}|}\bm{0}_{|\mathcal{I}_{k}|}^{T}\\\
\bm{0}_{|\mathcal{I}_{3}|}\bm{0}_{|\mathcal{I}_{2}|}^{T}&\Omega_{-1}^{(3)}&\dots&\bm{0}_{|\mathcal{I}_{3}|}\bm{0}_{|\mathcal{I}_{k}|}^{T}\\\
\vdots&\vdots&\vdots&\vdots\\\
\bm{0}_{|\mathcal{I}_{k}|}\bm{0}_{|\mathcal{I}_{2}|}^{T}&\bm{0}_{|\mathcal{I}_{k}|}\bm{0}_{|\mathcal{I}_{3}|}^{T}&\dots&\Omega_{-1}^{(k)}\end{bmatrix}+\bar{\sigma}^{2}\mathbf{1}_{n-|\mathcal{I}_{1}|}\mathbf{1}_{n-|\mathcal{I}_{1}|}^{T},$

and

$\Omega_{-1}^{(j)}={\rm
diag}(\\{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}\\}_{i\in\mathcal{I}_{j}})+\bar{\sigma}^{2}\mathbf{1}_{|\mathcal{I}_{j}|}\mathbf{1}_{|\mathcal{I}_{j}|}^{T},\qquad
j\neq 1.$

By Woodbury matrix identity (Petersen et al., 2008),

$\left(\Omega_{-1}^{(j)}\right)^{-1}={\rm
diag}(\\{(\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2})^{-1}\\}_{i\in\mathcal{I}_{j}})-\left((\bar{\sigma}^{2})^{-1}+\sum_{i\in\mathcal{I}_{j}}(\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2})^{-1}\right)^{-1}\\\
\times\begin{bmatrix}(\sigma_{\mathcal{I}_{j}[1]}^{2}/\beta_{\mathcal{I}_{j}[1]}+\bar{\sigma}_{j}^{2})^{-1}\\\
(\sigma_{\mathcal{I}_{j}[2]}^{2}/\beta_{\mathcal{I}_{j}[2]}+\bar{\sigma}_{j}^{2})^{-1}\\\
\vdots\\\
(\sigma_{\mathcal{I}_{j}[|\mathcal{I}_{j}|]}^{2}/\beta_{\mathcal{I}_{j}[|\mathcal{I}_{j}|]}+\bar{\sigma}_{j}^{2})^{-1}\end{bmatrix}\begin{bmatrix}(\sigma_{\mathcal{I}_{j}[1]}^{2}/\beta_{\mathcal{I}_{j}[1]}+\bar{\sigma}_{j}^{2})^{-1}\\\
(\sigma_{\mathcal{I}_{j}[2]}^{2}/\beta_{\mathcal{I}_{j}[2]}+\bar{\sigma}_{j}^{2})^{-1}\\\
\vdots\\\
(\sigma_{\mathcal{I}_{j}[|\mathcal{I}_{j}|]}^{2}/\beta_{\mathcal{I}_{j}[|\mathcal{I}_{j}|]}+\bar{\sigma}_{j}^{2})^{-1}\end{bmatrix}^{T}.$

With this, we have

$\mathbf{1}_{|\mathcal{I}_{j}|}^{T}\left(\Omega_{-1}^{(j)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{j}|}=\left(\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\right)^{-1}+\bar{\sigma}^{2}\right)^{-1},$
$\left(\Omega_{-1}^{(j)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{j}|}=\frac{\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\right)^{-1}}{\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\right)^{-1}+\bar{\sigma}^{2}}\begin{bmatrix}(\sigma_{\mathcal{I}_{j}[1]}^{2}/\beta_{\mathcal{I}_{j}[1]}+\bar{\sigma}_{j}^{2})^{-1}\\\
(\sigma_{\mathcal{I}_{j}[2]}^{2}/\beta_{\mathcal{I}_{j}[2]}+\bar{\sigma}_{j}^{2})^{-1}\\\
\vdots\\\
(\sigma_{\mathcal{I}_{j}[|\mathcal{I}_{j}|]}^{2}/\beta_{\mathcal{I}_{j}[|\mathcal{I}_{j}|]}+\bar{\sigma}_{j}^{2})^{-1}\end{bmatrix},$

and

$\Omega_{-1}^{-1}=\begin{bmatrix}(\Omega_{-1}^{(2)})^{-1}&\bm{0}_{|\mathcal{I}_{2}|}\bm{0}_{|\mathcal{I}_{3}|}^{T}&\dots&\bm{0}_{|\mathcal{I}_{2}|}\bm{0}_{|\mathcal{I}_{k}|}^{T}\\\
\bm{0}_{|\mathcal{I}_{3}|}\bm{0}_{|\mathcal{I}_{2}|}^{T}&(\Omega_{-1}^{(3)})^{-1}&\dots&\bm{0}_{|\mathcal{I}_{3}|}\bm{0}_{|\mathcal{I}_{k}|}^{T}\\\
\vdots&\vdots&\vdots&\vdots\\\
\bm{0}_{|\mathcal{I}_{k}|}\bm{0}_{|\mathcal{I}_{2}|}^{T}&\bm{0}_{|\mathcal{I}_{k}|}\bm{0}_{|\mathcal{I}_{3}|}^{T}&\dots&(\Omega_{-1}^{(k)})^{-1}\end{bmatrix}-\\\
\left((\bar{\sigma}^{2})^{-1}+\sum_{j=2}^{k}\left(\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\right)^{-1}+\bar{\sigma}^{2}\right)^{-1}\right)\\\
\times\begin{bmatrix}\left(\Omega_{-1}^{(2)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{2}|}\\\
\left(\Omega_{-1}^{(3)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{3}|}\\\ \vdots\\\
\left(\Omega_{-1}^{(4)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{4}|}\end{bmatrix}\begin{bmatrix}\left(\Omega_{-1}^{(2)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{2}|}\\\
\left(\Omega_{-1}^{(3)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{3}|}\\\ \vdots\\\
\left(\Omega_{-1}^{(4)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{4}|}\end{bmatrix}^{T}.$

Similarly, we have

$\mathbf{1}_{m-|\mathcal{I}_{1}|}^{T}\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}=\left(\bar{\sigma}^{2}+\left[\sum_{j=2}^{k}\left(\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\right)^{-1}+\bar{\sigma}^{2}\right)^{-1}\right]^{-1}\right)^{-1},$
$\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}\\\
=\frac{\left[\sum_{j=2}^{k}\left(\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\right)^{-1}+\bar{\sigma}^{2}\right)^{-1}\right]^{-1}}{\bar{\sigma}^{2}+\left[\sum_{j=2}^{k}\left(\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\right)^{-1}+\bar{\sigma}^{2}\right)^{-1}\right]^{-1}}\begin{bmatrix}\left(\Omega_{-1}^{(2)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{2}|}\\\
\left(\Omega_{-1}^{(3)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{3}|}\\\ \vdots\\\
\left(\Omega_{-1}^{(4)}\right)^{-1}\mathbf{1}_{|\mathcal{I}_{4}|}\end{bmatrix},$

and

$\Omega_{1}^{-1}=\begin{bmatrix}{\rm
diag}(\\{(\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2})^{-1}\\}_{i=2}^{|\mathcal{I}_{1}|})&\bm{0}_{|\mathcal{I}_{1}|-1}\bm{0}_{m-|\mathcal{I}_{1}|}^{T}\\\
\bm{0}_{m-|\mathcal{I}_{1}|}\bm{0}_{|\mathcal{I}_{1}|-1}^{T}&\Omega_{-1}^{-1}\end{bmatrix}-\\\
\left((\bar{\sigma}_{1}^{2})^{-1}+\sum_{i=2}^{|\mathcal{I}_{j}|}(\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2})^{-1}+\mathbf{1}_{m-|\mathcal{I}_{1}|}^{T}\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}\right)^{-1}\\\
\times\begin{bmatrix}(\sigma_{2}^{2}/\beta_{2}+\bar{\sigma}_{1}^{2})^{-1}\\\
\vdots\\\ (\sigma_{m}^{2}/\beta_{2}+\bar{\sigma}_{1}^{2})^{-1}\\\
\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}\end{bmatrix}\begin{bmatrix}(\sigma_{2}^{2}/\beta_{2}+\bar{\sigma}_{1}^{2})^{-1}\\\
\vdots\\\ (\sigma_{m}^{2}/\beta_{2}+\bar{\sigma}_{1}^{2})^{-1}\\\
\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}\end{bmatrix}^{T}.$

Combining the above expressions, we obtain the generalized least squares
estimate for $\theta^{*}_{1}$, which satisfies

$\displaystyle\left(\beta_{1}/\sigma_{1}^{2}+\left\\{\left[\mathbf{1}_{m-|\mathcal{I}_{1}|}^{T}\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}+\sum_{i=2}^{m}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2}}\right]^{-1}+\bar{\sigma}_{i}^{2}\right\\}^{-1}\right)\widehat{\theta}_{1}^{GLS}=$
$\displaystyle\qquad\frac{1}{\sigma_{1}^{2}}X_{1}^{T}y_{1}+\frac{\left\\{\left[\mathbf{1}_{m-|\mathcal{I}_{1}|}^{T}\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}+\sum_{i=2}^{m}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2}}\right]^{-1}+\bar{\sigma}_{i}^{2}\right\\}^{-1}}{\beta_{1}/\sigma_{1}^{2}+\left\\{\left[\mathbf{1}_{m-|\mathcal{I}_{1}|}^{T}\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}+\sum_{i=2}^{m}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2}}\right]^{-1}+\bar{\sigma}_{i}^{2}\right\\}^{-1}}$
$\displaystyle\hskip
60.00009pt\times\frac{\sum_{i=2}^{m}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2}}\widehat{\theta}_{i}^{d}}{\mathbf{1}_{m-|\mathcal{I}_{1}|}^{T}\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}+\sum_{i=2}^{m}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2}}}$
$\displaystyle\qquad+\frac{\left\\{\left[\mathbf{1}_{m-|\mathcal{I}_{1}|}^{T}\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}+\sum_{i=2}^{m}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2}}\right]^{-1}+\bar{\sigma}_{i}^{2}\right\\}^{-1}}{\beta_{1}/\sigma_{1}^{2}+\left\\{\left[\mathbf{1}_{m-|\mathcal{I}_{1}|}^{T}\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}+\sum_{i=2}^{m}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2}}\right]^{-1}+\bar{\sigma}_{i}^{2}\right\\}^{-1}}$
$\displaystyle\hskip
50.00008pt\times\frac{\left[\sum_{j=2}^{k}\left(\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\right)^{-1}+\bar{\sigma}^{2}\right)^{-1}\right]^{-1}}{\mathbf{1}_{m-|\mathcal{I}_{1}|}^{T}\Omega_{-1}^{-1}\mathbf{1}_{m-|\mathcal{I}_{1}|}+\sum_{i=2}^{m}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{1}^{2}}}$
$\displaystyle\hskip
50.00008pt\times\sum_{j=2}^{k}\frac{\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\right)^{-1}\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\widehat{\theta}_{i^{\prime}}^{d}\right)}{\left(\sum_{i^{\prime}\in\mathcal{I}_{j}}\frac{1}{\sigma_{i}^{2}/\beta_{i}+\bar{\sigma}_{j}^{2}}\right)^{-1}+\bar{\sigma}^{2}}.$

Observe that $\widehat{\theta}_{1}^{GLS}$ is exactly $\widehat{\theta}_{1}$
when $\gamma_{i}=\frac{1}{\bar{\sigma}_{j}^{2}}$ for all
$i\in\mathcal{I}_{j},j=1,\dots,k$ and
$\lambda_{j}=\frac{1}{\bar{\sigma}^{2}}$. Our claim holds by the Gauss-Markov
theorem (Kariya and Kurata, 2004).

###  A.11 Proof of Proposition 5.3

Under our model, for $i\neq 1$, $i\in\mathcal{I}_{j}$,

$y_{i}=X_{i}\theta_{1}^{*}-X_{i}(\xi_{1}+\bar{\xi}_{1}-\bar{\xi}_{j}-\xi_{i})+\epsilon_{i}.$

Therefore, similar to Eq. A.1, we have

$\begin{bmatrix}y_{1}\\\ \widehat{\theta}_{2}^{d}\\\ \vdots\\\
\widehat{\theta}_{m}^{d}\end{bmatrix}=\begin{bmatrix}X_{1}\\\ X_{2}\\\
\vdots\\\ X_{n}\end{bmatrix}\theta_{1}^{*}+\begin{bmatrix}\epsilon_{1}\\\
-X_{2}(\xi_{1}-\xi_{2})+\epsilon_{2}\\\ \vdots\\\
-X_{n}(\xi_{1}+\bar{\xi}_{1}-\bar{\xi}_{k}-\xi_{n})+\epsilon_{n}.\end{bmatrix}.$

Let $\tilde{\theta}_{1}$ be the solution to the generalized least squares
problem defined by the equation above. In general,
$\tilde{\theta}_{1}\neq\widehat{\theta}_{1}$. For example, when $X_{i}\neq I$
for all $i$. Since $\widehat{\theta}_{i}^{d}$ are linear in $(X_{i},y_{i})$,
$\widehat{\theta}_{1}$ is linear in $\\{(X_{i},y_{i})\\}_{i=1}^{n}$. Our
proposition then holds by the Gauss-Markov theorem (Kariya and Kurata, 2004).

###  A.12 Proof of Proposition 5.4

Similarly to Eq. A.1, we have

$\widehat{\theta}_{1}=\frac{n-1}{2n}y_{1}+\frac{n+1}{2n}\bar{y}_{-1},$

where

$\bar{y}_{-1}=\frac{1}{n-1}\sum_{i=2}^{n}y1.$

Then

$\mathbb{E}[\|\widehat{\theta}_{1}-\theta_{1}^{*}\|^{2}]=\frac{(n-1)^{2}}{4n^{2}}d+d+\frac{(n+1)^{2}}{4n^{2}}\mathbb{E}[\|\bar{y}_{-1}-\bar{\theta}^{*}\|^{2}].$

The estimator $\tilde{\theta}^{JS}_{1}$ is obtained as

$\tilde{\theta}_{1}^{JS}=\frac{n-1}{2n}y_{1}+\frac{n+1}{2n}C\bar{y}_{-1},$

where $C\in[0,1]$ is a shrinkage parameter. Then

$\mathbb{E}[\|\tilde{\theta}_{1}^{JS}-\theta_{1}^{*}\|^{2}]=\frac{(n-1)^{2}}{4n^{2}}d+d+\frac{(n+1)^{2}}{4n^{2}}\mathbb{E}[\|C\bar{y}_{-1}-\bar{\theta}^{*}\|^{2}].$

When estimating $\bar{\theta}^{*}$ and $d>3$, there is a $C<1$ such that

$\mathbb{E}[\|C\bar{y}_{-1}-\bar{\theta}^{*}\|^{2}]\leq\mathbb{E}[\|\bar{y}_{-1}-\bar{\theta}^{*}\|^{2}],$

where the exact form of C is discussed in James and Stein (1992); Bock (1975);
Kubokawa (1991). Consequently,

$\mathbb{E}[\|\tilde{\theta}_{1}^{JS}-\theta_{1}^{*}\|^{2}]\leq\mathbb{E}[\|\widehat{\theta}_{1}-\theta_{1}^{*}\|^{2}].$

Since

$\mathbb{E}[\tilde{\theta}_{1}^{JS}\mid\theta_{1}^{*}]=\left(\frac{n-1}{2n}+\frac{C(n+1)}{2n}\right)\theta_{1}^{*}=\left(1-\frac{n+1}{2n}(1-C)\right)\theta_{1}^{*}$

and $C<1$, $\tilde{\theta}_{1}^{JS}$ is biased.

##  Appendix B Useful Results

######  Theorem B.1.

Let $g(x):\mathbb{R}^{d}\to\mathbb{R}$ be an $L_{g}$-smooth and convex
function. Let

$D_{g}(x,y)=g(x)-g(y)-(\nabla g(y))^{T}(x-y),\quad x,y\in\mathbb{R}^{d}.$

Then, for all $x,y\in\mathbb{R}^{d}$, we have

$\|\nabla g(x)-\nabla g(y)\|_{2}^{2}\leq 2L_{g}D_{g}(x,y).$

###### Proof.

Directly follows from (2.1.10) in Theorem 2.1.5 of Nesterov et al. (2018). ∎