# Personalized Federated Learning for Heterogeneous Clients with Clustered
Knowledge Transfer

Yae Jee Cho
ECE Department
Carnegie Mellon University
Pittsburgh, PA 15213
<EMAIL_ADDRESS>&Jianyu Wang
ECE Department
Carnegie Mellon University
Pittsburgh, PA 15213
<EMAIL_ADDRESS>&Tarun Chiruvolu
School of Computer Science
Carnegie Mellon University
Pittsburgh, PA 15213
<EMAIL_ADDRESS>&Gauri Joshi
ECE Department
Carnegie Mellon University
Pittsburgh, PA 15213
<EMAIL_ADDRESS>
###### Abstract

Personalized federated learning (FL) aims to train model(s) that can perform
well for individual clients that are highly data and system heterogeneous.
Most work in personalized FL, however, assumes using the same model
architecture at all clients and increases the communication cost by
sending/receiving models. This may not be feasible for realistic scenarios of
FL. In practice, clients have highly heterogeneous system-capabilities and
limited communication resources. In our work, we propose a personalized FL
framework, PerFed-CKT, where clients can use heterogeneous model architectures
and do not directly communicate their model parameters. PerFed-CKT uses
_clustered co-distillation_ , where clients use logits to transfer their
knowledge to other clients that have similar data-distributions. We
theoretically show the convergence and generalization properties of PerFed-CKT
and empirically show that PerFed-CKT achieves high test accuracy with several
orders of magnitude lower communication cost compared to the state-of-the-art
personalized FL schemes.

##  1 Introduction

The emerging paradigm of federated learning (FL) [1, 2, 3, 4] enabled the use
of data collected by thousands of resource-constrained clients to train
machine learning models without having to transfer the data to the cloud. Most
recent work [5, 6, 7, 8] is focused on algorithms for training a single global
model with edge-clients via FL. However, due to the inherently high data-
heterogeneity across clients [9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19], a
single model that is trained to perform best in expectation for the sum of all
participating clients’ loss functions may not work well for each client [20,
21, 22]. For example, if a single global model was trained for next-word
prediction with all available clients and the input was “I was born in”, the
model will most likely give bad results for the individual clients.

The limited generalization properties of conventionally trained FL models to
clients with scarce local data calls for the design of FL algorithms to train
personalized models that can perform well for individual clients. Several
works have investigated personalized FL [23, 24, 25, 26], including applying
meta-learning [23], training separate models on each client with weighted
aggregation of other clients’ models [24], using the global objective as a
regularizer for training individual models at each client [25], or using
model/data-interpolation with clustering for personalization [26]. However,
the aforementioned work does not consider two critical factors of FL: i) the
computation and memory capabilities can be heterogeneous across clients and
ii) the cost of communicating high-dimensional models with the server can be
prohibitively high [27, 28]. Most work in personalized FL assumes a
homogeneous model architecture across clients, and frequent communication of
the model parameters.

In this work, we propose training personalized models with clustered co-
distillation. Co-distillation [29, 30, 31] is an approach to perform
distributed training across different clients with reduced communication cost
by only exchanging models’ predictions on a common unlabeled dataset instead
of the model parameters. This method adds a regularization term to the local
loss of each client to penalize the client’s prediction from being
significantly different from the predictions of other clients. In the
conventional co-distillation, the regularizing term for each client is the
average of all other participating clients’ predictions. However in FL,
clients’ data can be highly heterogeneous. Thus, forcing each client to follow
the average prediction of all clients can exacerbate its generalization by
learning irrelevant knowledge from clients that have significantly different
data distributions [32]. Hence, we propose a novel _clustered co-distillation_
framework PerFed-CKT, where each client uses the average prediction of only
the clients that have similar data distributions. This way, we prevent each
client from assimilating irrelevant knowledge from unrelated clients.

In short, PerFed-CKT largely improves on current personalized FL strategies in
the following ways:

  * •

Allows model heterogeneity across clients where the architecture and size of
the model for local training can vary across clients.

  * •

Dramatically reduces the communication cost by transferring logits instead of
model parameters between the clients and the server.

  * •

Improves generalization performance for data-scarce clients, while preventing
learning from unrelated clients by clustered knowledge transfer.

We also present theoretical analysis of PerFed-CKT with its convergence
guarantees and generalization performance. The generalization results show
that clustering indeed helps in terms of improvement of the generalization
properties of individual clients. Our experiments demonstrate that for both
model-homogeneous and model-heterogeneous environments, PerFed-CKT can achieve
high test accuracy with several orders of magnitude less communication.

##  2 Background and Related Work

#### Personalized FL.

In personalized FL, the goal is to train a single or several model(s) that can
generalize well to each client’s test dataset. In [23], using meta-learning
for training a global model that better represents each client’s data was
proposed. A similar line of work using the moreau envelope as a regularizer
was proposed in [33]. Work in [24] proposed to find the optimal weighted
combination of models from clients so that each client gets a model that
better represents its target data distribution. The authors in [26] propose
general approaches that can be applied to vanilla FL for personalization,
including client clustering and data/model interpolation.

The aforementioned work however, all requires model homogeneity and direct
communication of model parameters across clients/server. Although [34] does
consider communication cost in personalization by using distillation, it does
not provide any theoretical guarantees and does not consider the high data
heterogeneity across clients that can tamper with the personalization
performance. Moreover they require the presence of a large labeled public
dataset, which is realistically an expensive resource to have access to. Our
work investigates a novel personalized FL framework that allows model
heterogeneity, improves communication-efficiency, and utilizes data
heterogeneity with clustering while using a small unlabeled public dataset.

(a)

(b)

(c)

(d)

Figure 1: Toy example with linear regression for given input data
$\mathbf{X}\in\mathbb{R}^{2}$ and local (true) models
$\mathbf{w}\in\mathbb{R}^{2\times 1}$ for three clients indexed by 0-2; (a):
the true model for each client and the global model from FedAvg. The resulting
global model does not match well with clients’ true model; (b)-(d): the model
for each client resulting from PerFed-CKT with and without clustering (simple
average of logits). PerFed-CKT with clustering yields the model closest to the
true model for all clients.

#### Knowledge Transfer.

Knowledge distillation (KD) [35] is prominently used as a method of knowledge
transfer from a pre-trained larger model to a smaller model [36, 37, 38, 39,
40, 41, 42, 43, 44]. Extending from this conventional KD, co-distillation [32,
45, 46] transfers knowledge across multiple models that are being trained
concurrently. Specifically, each model is trained with the supervised loss
with an additional regularizer term that encourages the model to yield similar
outputs to the outputs of the other models that are also being trained.

Using co-distillation for improved generalization in distributed training has
been recently proposed in [29, 30, 31]. Authors of [29] have shown empirically
that co-distillation indeed improves generalization for distributed learning
but often results in over-regularization, where the trained model’s
performance drops due to overfitting to the regularization term. In [30], co-
distillation was suggested for communication-efficient distributed training,
but not in the personalized FL context where data can be highly heterogeneous
across nodes and presented limited experiments on a handful of nodes with
homogeneous data distributions across the nodes.

Applying co-distillation for personalized FL presents a unique challenge in
that each client’s data distribution can be significantly different from other
clients’ distributions. Using standard co-distillation with the entire
clients’ models can actually worsen clients’ test performance due to learning
irrelevant information from other clients with different data distributions.
We show in our work that this is indeed the case and show that using
clustering to find the clients that have similar data-distributions with each
other and then performing co-distillation within the clusters improves the
personalized model’s performance significantly for each client.

##  3 Proposed Personalized FL Framework: PerFed-CKT

### Problem Formulation

Consider a cross-device FL setup where a large number of $K$ clients (edge-
devices) are connected to a central server. We consider a $N$-class
classification task where each client $k\in[K]$ has its local training dataset
$\mathcal{B}_{k}$ with $|\mathcal{B}_{k}|=m_{k}$ data samples. We denote
$p_{k}=m_{k}/{\sum_{k=1}^{K}m_{k}}$ as the fraction of data for client $k$.
Each data sample $\xi$ is a pair $(\mathbf{x},y)$ where
$\mathbf{x}\in\mathbb{R}^{d}$ is the input and $y\in[1,N]$ is the label. The
dataset $\mathcal{B}_{k}$ is drawn from the local data distribution
$\mathcal{D}_{k}$, where we denote the empirical data distribution of
$\mathcal{B}_{k}$ as $\widehat{\mathcal{D}}_{k}$. In the standard FL [1],
clients aim to collaboratively find the model parameter vector
$\mathbf{w}\in\mathbb{R}^{n}$ that maps the input $\mathbf{x}$ to label $y$,
such that $\mathbf{w}$ minimizes the empirical risk
$F(\mathbf{w})=\sum_{k=1}^{K}p_{k}F_{k}(\mathbf{w})$. The function
$F_{k}(\mathbf{w})$ is the local objective of client $k$, defined as
$F_{k}(\mathbf{w})=\frac{1}{|\mathcal{B}_{k}|}\sum_{\xi\in\mathcal{B}_{k}}f(\mathbf{w},\xi)$
with $f(\mathbf{w},\xi)$ being the composite loss function.

Due to high data heterogeneity across clients, the optimal model parameters
$\mathbf{w}^{*}$ that minimize $F(\mathbf{w})$ can generalize badly to clients
whose local objective $F_{k}(\mathbf{w})$ significantly differs from
$F(\mathbf{w})$. Such clients may opt out of FL, and instead train their own
models $\mathbf{w}_{k}\in\mathbb{R}^{n_{k}}$ by minimizing their local
objectives, where $\mathbf{w}_{k}$ can be heterogeneous in dimension. This can
work well for clients with a large number of training samples (i.e., large
$m_{k}$), since their empirical data distribution $\widehat{\mathcal{D}}_{k}$
becomes similar to $\mathcal{D}_{k}$, ensuring good generalization. However,
if clients have a small number of training samples, which is often the case
[47], the distributions $\mathcal{D}_{k}$ and $\widehat{\mathcal{D}}_{k}$ can
differ significantly, and therefore a model $\mathbf{w}_{k}$ trained only
using the local dataset $\mathcal{B}_{k}$ can generalize badly. Hence,
although clients with small number of training samples are motivated to
participate in FL, the clients may actually not benefit from FL due to the bad
generalization properties coming from other clients with significantly
different data distributions. We show that with our proposed PerFed-CKT,
clients with small training samples still benefit from participating in FL,
improving generalization by being clustered with clients with similar data
distributions.

#### Objective with Clustered Knowledge Transfer.

We use co-distillation across different clients in FL for personalization,
where each client co-distills with only the other clients that have similar
model outputs with its own model, namely, clustered knowledge transfer. With
clustered knowledge transfer, each client learns from clients that have
similar data distributions to improve its generalization performance.Formally,
we consider each client having access to its private dataset $\mathcal{B}_{k}$
and a public dataset $\mathcal{P}$, consisting of unlabeled data. The public
dataset $\mathcal{P}$ is used as a reference dataset for co-distillation
across clients111Unlike the majority of existing distillation methods which
require expensive labeled data, we show the feasibility of leveraging
unlabeled datasets. Note that unlabeled data can be either achieved by
existing datasets or a pre-trained generator (e.g., GAN).. The classification
models $\mathbf{w}_{k},~{}k\in[K]$ output soft-decisions (logits) over the
pre-defined number of classes $N$, which is a probability vector over the $N$
classes. We refer to the soft-decision of model $\mathbf{w}_{k}$ over any
input data $\mathbf{x}$ in either the private or public dataset as
$s(\mathbf{w}_{k},\mathbf{x}):\mathbb{R}^{n_{k}}\times(\mathcal{B}_{k}\cup\mathcal{P})\rightarrow\Delta_{N}$,
where $\Delta_{N}$ stands for the probability simplex over $N$. For notational
simplicity, we define $\mathbf{s}_{k}\in\mathbb{R}^{|\mathcal{P}|\times N}$ as
$s(\mathbf{w}_{k},\mathbf{x})\in\mathbb{R}^{1\times N}$,
$\mathbf{x}\in\mathcal{P}$ stacked into rows for each $\mathbf{x}$. We
similarly define
$\overline{\mathbf{s}}_{k}=\sum_{i=1}^{K}\alpha_{k,i}\mathbf{s}_{i}$.

The clients are connected via a central aggregating server. Each client seeks
to find the model parameter $\mathbf{w}_{k}$ that minimizes the empirical risk
$\Phi_{k}(\mathbf{w}_{k};\overline{\mathbf{s}}_{k})$, where
$\Phi_{k}(\mathbf{w}_{k};\overline{\mathbf{s}}_{k})$ is a sum of the empirical
risk of its own local training data $F_{k}(\mathbf{w}_{k})$ and the
regularization term as follows:

$\displaystyle\Phi_{k}(\mathbf{w}_{k};\overline{\mathbf{s}}_{k})=F_{k}(\mathbf{w}_{k})+\underbrace{\frac{\lambda}{|\mathcal{P}|}\sum_{\mathbf{x}\in\mathcal{P}}\|\overline{\mathbf{s}}_{k}(\mathbf{x})-s(\mathbf{w}_{k},\mathbf{x})\|_{2}^{2}}_{\text{regularization
term}}$ (1)

The term
$\overline{\mathbf{s}}_{k}(\mathbf{x})=\sum_{i=1}^{K}\alpha_{k,i}s(\mathbf{w}_{i},\mathbf{x})$
denotes the weighted average of the logits from all clients for an arbitrary
set of weights for client $k$, i.e., $\\{\alpha_{k,i}\\}_{i\in[K]}$ such that
$\sum_{i=1}^{K}\alpha_{k,i}=1,~{}\forall k\in[K]$. The term $\lambda$
modulates the weight of the regularization term. The weight $\alpha_{k,i}$ for
each client $i,~{}i\in[K]$ with respect to client $k$ results from clustering
the logits by the $\ell_{2}$-norm distance so that clients with similar logits
will have higher weights for each others’ logits. The aggregated logit
information with weights, $\overline{\mathbf{s}}_{k}$, is calculated and sent
by the server to the clients. Details of how the weights
$\alpha_{k,i},{i\in[K]}$ for each client $k$ are calculated and how the logits
are communicated are elaborated in more detail in the subsequent Algorithm
subsection. Before going into details of the algorithm we first give more
intuition on the formulation of the regularization term in the next paragraph.

#### Regularization Term.

Without the regularization term in 1, minimizing
$\Phi_{k}(\mathbf{w}_{k};\overline{\mathbf{s}}_{k})$ with regards to
$\mathbf{w}_{k}$ is analogous to locally training in isolation for minimizing
the local objective function $F_{k}(\mathbf{w}_{k})$ for client $k$. If we
have the regularization term with
$\alpha_{k,i}=1/K,~{}\forall~{}k,~{}i\in[K]$, co-distillation is implemented
without clustering, using all of the clients’ knowledge. We show in the next
toy example and in a generalization bound derived for ensemble models in
personalization in Appendix D of why setting the values
$\alpha_{k,i},k,i\in[K]$ via clustering is critical to improve
personalization.

We consider a toy example with linear regression where we have three clients
with true models as in Figure 1(a) where true model 0 and 1 are similar to
each other but true model 2 is different. The global model trained from
vanilla FedAvg does not match well with any true models as shown in Figure
1(a). If we minimize 1 with respect to $\mathbf{w}_{k}$ without clustering,
i.e., $\alpha_{k,i}=1/K,~{}\forall~{}k,~{}i\in[K]$, the output local model
also diverges from the true model for each client, especially for client 0 and
1, due to the heterogeneity across the true models (see Figure 1(b)-(d)).
Finally, if we minimize 1 with clustering so that for client $k$, higher
weight $\alpha_{k,i}$ is given to the client $i$ that has a similar true model
to client $k$, and smaller weight is given to the other client that has a
different true model, the output local model of client $k$ gets close to its
true model. This is further explored theoretically in Theorem 4.2. PerFed-CKT
is based on this motivation where we set the weights for co-distillation in
$\overline{\mathbf{s}}_{k}(\mathbf{x})=\sum_{i=1}^{K}\alpha_{k,i}s(\mathbf{w}_{i},\mathbf{x})$
so that each client $k$ sets higher $\alpha_{k,i}$, $i\in[K]$ for client $i$
that has smaller difference between $s(\mathbf{w}_{i},\mathbf{x})$ and
$s(\mathbf{w}_{k},\mathbf{x})$. Details of the setup for Figure 1 are in
Appendix E.

### Algorithm.

We minimize 1 with respect to $\mathbf{w}_{k}$ for each client $k$ on its own
device with only communicating the logits instead of the actual model
$\mathbf{w}_{k}$, with the server. With $(t,r)$ denoting the communication
round $t$ and local iteration $r$, we define
$\overline{\mathbf{s}}_{k}^{(t,0)}=\sum_{i=1}^{K}\alpha_{k,i}^{(t,0)}\mathbf{s}_{i}^{(t,0)}$
for $t\in[0,T-1]$ and $r\in[0,\tau-1]$ where
$\overline{\mathbf{s}}_{k}^{(t,0)}$ is fixed for all $r$ and updated only for
every $t$. The term $T$ and $\tau$ is the total number of communication rounds
and local iterations respectively. Note that the logit information
$\overline{\mathbf{s}}_{k}^{(t,0)}$ is computed and sent by the server to the
clients for every communication round $t$. Details are in the following
paragraphs and Algorithm 1.

#### Client Side Update.

From 1, given $\overline{\mathbf{s}}_{k}^{(t,0)}$ from the server, each
client’s local update rule is:

$\displaystyle\begin{aligned}
\mathbf{w}_{k}^{(t,r+1)}=\mathbf{w}_{k}^{(t,r)}-\eta_{t}&\left[\frac{2\lambda}{|\mathcal{P}_{k}^{(t,r)}|}\sum_{\mathbf{x}\in\mathcal{P}_{k}^{(t,r)}}\nabla
s(\mathbf{w}_{k}^{(t,r)},\mathbf{x})^{T}\left(s(\mathbf{w}_{k}^{(t,r)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t,0)}(\mathbf{x})\right)\right.\\\
&\left.+\frac{1}{|\xi_{k}^{(t,r)}|}\sum_{\xi\in|\xi_{k}^{(t,r)}|}\nabla
f(\mathbf{w}_{k}^{(t,r)},\xi)\right]\end{aligned}$ (2)
$\displaystyle~{}~{}\quad\quad\quad\triangleq\mathbf{w}_{k}^{(t,r)}-\eta_{t}\mathbf{g}_{k}(\mathbf{w}_{k}^{(t,r)};\overline{\mathbf{s}}_{k}^{(t,0)})$
(3)

The term $\mathbf{w}_{k}^{(t,r)}$ denotes the local model parameters of client
$k$, $\eta_{t}$ is the learning rate, $\xi_{k}^{(t,r)}$ is the mini-batch
randomly sampled from client $k$’s local dataset $\mathcal{B}_{k}$, and
$\mathcal{P}_{k}^{(t,r)}$ is the mini-batch randomly sampled from the public
dataset from client $k$. We also denote the updated local model of client $k$
after all $\tau$ local iterations for round $t$ as
$\mathbf{w}_{k}^{(t+1,0)}=\mathbf{w}_{k}^{(t,\tau)}$.

For PerFed-CKT, we consider partial client participation where for every
communication round $t$, $m$ clients are selected with probability $p_{k}$
without replacement from $k\in[K]$. We denote the set of selected clients as
$\mathcal{S}^{(t,0)}$ that is fixed for all local iterations $r\in[0,\tau-1]$.
If a client $k\in[K]$ was most recently selected in the previous communication
round $t^{\prime}<t$, and selected again for the current communication round
$t$, we assume that
$\mathbf{w}_{k}^{(t,0)}=\mathbf{w}_{k}^{(t^{\prime},\tau)}$. In other words,
we retrieve the most recently updated local model for the client that is
selected for the next communication round and use that model for local
training. Each client $k\in\mathcal{S}^{(t,0)}$ takes $\tau\geq 1$ local
updates before sending its logits back to the server where each local update
step follows the update in 2.

#### Server Side Clustered Knowledge Aggregation.

After the $\tau$ local iterations, each client $k\in\mathcal{S}^{(t,0)}$ sends
the logits from its updated local model to the server. The logits are denoted
as $\mathbf{s}_{k}^{(t+1,0)}\in\mathbb{R}^{|\mathcal{P}|\times N}$ which is
$s(\mathbf{w}_{k}^{(t+1,0)},\mathbf{x})=s(\mathbf{w}_{k}^{(t,\tau)},\mathbf{x})\in\mathbb{R}^{1\times
N}$ stacked in to rows for each $\mathbf{x}\in\mathcal{P}$. The server uses
$c$-means clustering (also known conventionally as the $k$-means clustering
algorithm [48]) to cluster the received $m$ different set of logits,
$\mathbf{s}_{k}^{(t+1,0)},k\in\mathcal{S}^{(t)}$ to $c$ clusters, where $c$ is
an integer such that $1\leq c\leq m$. The server also gets the next set of
selected clients $\mathcal{S}^{(t+1,0)}$ and sends the centroids
$\\{\mathbf{c}_{i}^{(t+1,0)}\\}_{i\in[c]}$ for each cluster to the clients in
$\mathcal{S}^{(t+1,0)}$. Each client $k^{\prime}\in\mathcal{S}^{(t+1,0)}$ then
determines the centroid that is closest to its current model’s logit as

$\displaystyle\overline{\mathbf{s}}_{k^{\prime}}^{(t+1,0)}=\operatorname*{arg\,min}_{\\{\mathbf{c}_{i}^{(t+1,0)}\\}_{i\in[c]}}\|\mathbf{c}_{i}^{(t+1,0)}-\mathbf{s}_{k^{\prime}}^{(t+1,0)}\|_{2}^{2}$
(4)

and uses it for the local update in 2. Since we defined
$\overline{\mathbf{s}}_{k^{\prime}}^{(t+1,0)}=\sum_{i=1}^{K}\alpha_{k^{\prime},i}^{(t+1,0)}\mathbf{s}_{i}^{(t+1,0)}$,
4 gives a natural selection of $\alpha_{k^{\prime},i}^{(t+1,0)}$ which gives
higher weight to the $\mathbf{s}_{i}^{(t+1,0)},~{}i\in\mathcal{S}^{(t,0)}$
that is closer to client $k^{\prime}$’s logits, i.e.,
$\mathbf{s}_{k^{\prime}}^{(t+1,0)}$. PerFed-CKT can set $\alpha_{k,i}=0$ for
client $i$ if its logit is significantly different from the logit of client
$k$ or if it was not included in the previous set of selected clients.

Algorithm 1 Personalized Federated Learning with Clustered Knowledge Transfer
(PerFed-CKT)

1: Input: $m,~{}\\{p_{k}\\}_{k\in[K]}$, mini-batch size $b$,$b^{\prime}$ for
private, public data each, number of clusters $c$

2: Output: $\\{\mathbf{w}_{k}\\}_{k\in[K]}$

3: Initialize: $\\{\mathbf{s}_{k}^{(0,0)}\\}_{k\in\mathcal{S}^{(-1,0)}}$,
selected set of $m$ clients $\mathcal{S}^{(-1,0)}$

4: For ${t=0,...,T-1}$ communication rounds do:

5: Global server do:

6: Cluster $\\{\mathbf{s}_{k}^{(t,0)}\\}_{k\in\mathcal{S}^{(t-1,0)}}$ by
$c$-means clustering

7: Get centroids $\\{\mathbf{c}_{i}^{(t,0)}\\}_{i\in[c]}$ for each cluster

8: Select $m$ clients for $\mathcal{S}^{(t,0)}$ without replacement from $[K]$
by the dataset ratio $\\{p_{k}\\}_{k\in[K]}$

9: Send centroids $\\{\mathbf{c}_{i}^{(t,0)}\\}_{i\in[c]}$ to clients
$k\in\mathcal{S}^{(t,0)}$

10: Clients $k\in\mathcal{S}^{(t,0)}$ in parallel do:

11: Get $\mathbf{s}_{k}^{(t,0)}$ for current local model
$\mathbf{w}_{k}^{(t,0)}$, and find
$\overline{\mathbf{s}}_{k}^{(t,0)}=\operatorname*{arg\,min}_{\\{\mathbf{c}_{i}^{(t,0)}\\}_{i\in[c]}}\|\mathbf{c}_{i}^{(t,0)}-\mathbf{s}_{k}^{(t,0)}\|_{2}^{2}$.

12: For $r=0,...,\tau-1$ local iterations do:

13: Create mini-batch $\xi_{k}^{(t,r)}$ from sampling $b$ sam- ples uniformly
at random from $\mathcal{B}_{k}$, and mini- batch $\mathcal{P}_{k}^{(t,r)}$
from sampling $b^{\prime}$ samples uniformly at random from $\mathcal{P}$

14: Update
$\mathbf{w}_{k}^{(t,r+1)}\leftarrow\mathbf{w}_{k}^{(t,r)}-\eta\mathbf{g}_{k}(\mathbf{w}_{k}^{(t,r)};\overline{\mathbf{s}}_{k}^{(t,0)})$

15: Send $\mathbf{s}_{k}^{(t+1,0)}=\mathbf{s}_{k}^{(t,\tau)}$ for the updated
local model $\mathbf{w}_{k}^{(t,\tau)}$ back to the server

We show in the subsequent sections that our proposed PerFed-CKT indeed
converges and improves the generalization performance of the individual
clients’ personalized models by clustering. We also empirically show that
PerFed-CKT achieves high test accuracy as state-of-the-art (SOTA) personalized
FL algorithms with drastically smaller communication cost.

##  4 Theoretical Analysis of PerFed-CKT

In this section, we analyze the convergence and generalization properties of
PerFed-CKT, specifically highlighting the effect of clustered knowledge
transfer to generalization.

### Convergence Analysis

Here, we present the convergence guarantees of PerFed-CKT with regards to the
objective function
$\Phi_{k}(\mathbf{w}_{k}^{(t,0)};\overline{\mathbf{s}}_{k}^{(t,0)})$ as
$t\rightarrow\infty$ with $\tau=1$. We use the following assumptions for our
analysis:

######  Assumption 4.1.

The composite loss function $f(\mathbf{w},\xi)$ is Lipschitz-continuous and
Lipschitz-smooth for all $\mathbf{w},~{}\xi$, and therefore
$F_{1}(\mathbf{w}),~{}...,~{}F_{k}(\mathbf{w})$ are all $L_{f}$-continuous and
$L_{p}$-smooth for all $\mathbf{w}$.

######  Assumption 4.2.

Each $F_{1},~{}...,~{}F_{k}$ is bounded below by a scalar $F_{k,\inf}$ over
its domain for $k\in[K]$.

######  Assumption 4.3.

For the mini-batch $\xi_{k}$ uniformly sampled at random from
$\mathcal{B}_{k}$, the resulting stochastic gradient is unbiased, that is,
$\mathbb{E}\left[\frac{1}{|\xi_{k}|}\sum_{\xi\in\xi_{k}}\nabla
f(\mathbf{w}_{k},\xi)\right]=\nabla F_{k}(\mathbf{w}_{k})$.

######  Assumption 4.4.

The stochastic gradient’s expected squared norm is uniformly bounded, i.e.,
$\mathbb{E}\left\|\frac{1}{|\xi_{k}|}\sum_{\xi\in\xi_{k}}\nabla
f(\mathbf{w}_{k},\xi)\right\|^{2}\leq G^{2}$ for $k=1,...,K$.

######  Assumption 4.5.

$s(\mathbf{w},\mathbf{x})$ is $L_{s}-$continuous and $L_{g}-$smooth for all
$\mathbf{w}$ and $x$.

Now we present the convergence guarantees for PerFed-CKT in Theorem 4.1 below:

######  Theorem 4.1.

With 4.1-4.5, after running PerFed-CKT (Algorithm 1) for $t=T$ iterations on
client $k\in[K]$ with $K$ total clients participating, with the learning rate
satisfying
$\sum_{t=0}^{\infty}\eta_{t}=\infty,~{}\sum_{t=0}^{\infty}\eta_{t}^{2}<\infty$,
we have that the norm of the gradient of
$\Phi_{k}(\mathbf{w}_{k}^{(t,0)};\overline{\mathbf{s}}_{k}^{(t,0)})$ with
respect to $\mathbf{w}_{k}^{(t,0)}$ given $\overline{\mathbf{s}}_{k}^{(t,0)}$
goes to zero with probability 1 as $T\rightarrow\infty$, i.e., for every
client $k$,

$\displaystyle\lim_{t\rightarrow\infty}\|\nabla_{\mathbf{w}_{k}^{(t,0)}}\Phi_{k}(\mathbf{w}_{k}^{(t,0)};\overline{\mathbf{s}}_{k}^{(t,0)})\|=0$
(5)

The proof for Theorem 4.1 is presented in the Appendix A. Theorem 4.1 shows
that our proposed algorithm PerFed-CKT indeed converges to a first-order
stationary point with respect to $\mathbf{w}_{k}$ given
$\overline{\mathbf{s}}_{k}$ where the norm of the gradient of our main
objective function $\Phi_{k}(\mathbf{w}_{k};\overline{\mathbf{s}}_{k})$ with
respect to $\mathbf{w}_{k}$ is 0.

### Generalization Performance

Now we show the theoretical grounds for clustered knowledge distillation in
regards to the generalization performance for personalized FL in the problem
of linear regression. We also present a generalization bound for ensemble
models in the context of personalization in Appendix D. For $K$ clients in
total, we consider a Bayesian framework as in [25] where we have $\theta$
uniformly distributed on $\mathbb{R}^{d}$, and each device has its data
distributed with parameters $\mathbf{w}_{k}=\theta+\zeta_{k}$ where
$\zeta_{k}\sim\mathcal{N}(0,\upsilon_{k}^{2}\mathbf{I}_{d})$ and
$\mathbf{I}_{d}$ is the $d\times d$ identity matrix and $\upsilon_{k}$ is
unique to the client’s task. Suppose we have
$\mathbf{y}_{k}=\mathbf{X}_{k}\mathbf{w}_{k}+\mathbf{z},~{}k\in[K]$ where
$\mathbf{y}_{k}\in\mathbb{R}^{n},~{}\mathbf{X}_{k}\in\mathbb{R}^{n\times d}$,
and $\mathbf{z}\in\mathbb{R}^{n}$ such that
$\mathbf{z}\sim\mathcal{N}(0,\sigma^{2}\mathbf{I}_{d})$.

Let us consider a linear regression problem for each device $k$ such that we
have the empirical loss function as
$F_{k}(\mathbf{w}_{k})=\|\mathbf{X}_{k}\mathbf{w}_{k}-\mathbf{y}_{k}\|_{2}^{2}$.
We have that
$\widehat{\mathbf{w}}_{k}=(\mathbf{X}_{k}^{T}\mathbf{X}_{k})^{-1}\mathbf{X}_{k}^{T}\mathbf{y}_{k}$
is a noisy observation of $\mathbf{w}_{k}$ with additive covariance
$\sigma^{2}(\mathbf{X}_{k}^{T}\mathbf{X}_{k})^{-1}$ since
$\widehat{\mathbf{w}}_{k}\sim\mathcal{N}((\mathbf{X}_{k}^{T}\mathbf{X}_{k})^{-1}\mathbf{X}_{k}^{T}\mathbf{y}_{k},\sigma^{2}(\mathbf{X}_{k}^{T}\mathbf{X}_{k})^{-1})$.
Then using Lemma 2 from [25], with the following definitions:

$\displaystyle\Sigma_{k}\coloneqq\sigma^{2}(\mathbf{X}_{k}^{T}\mathbf{X}_{k})^{-1}+\upsilon_{k}^{2}\mathbf{I}_{d}$
(6) $\displaystyle\overline{\Sigma}_{\setminus k}=\left(\sum_{i\in[K],i\neq
k}\Sigma_{i}^{-1}\right)^{-1}$ (7) $\displaystyle\overline{\theta}_{\setminus
k}\coloneqq\overline{\Sigma}_{\setminus k}\sum_{i\in[K],i\neq
k}\Sigma_{i}^{-1}\widehat{\mathbf{w}}_{i}$ (8)

given $\\{\mathbf{X}_{i},\mathbf{y}_{i}\\}_{i\in[K],i\neq k}$ we have that

$\displaystyle\theta=\overline{\theta}_{\setminus k}+\gamma$ (9)

where $\gamma\sim\mathcal{N}(0,\overline{\Sigma}_{\setminus k})$. Further, if
we let

$\displaystyle\widetilde{\Sigma}_{k}\coloneqq\overline{\Sigma}_{\setminus
k}+\upsilon_{k}^{2}\mathbf{I}_{d}$ (10)
$\displaystyle\overline{\Sigma}_{k}\coloneqq\left((\widetilde{\Sigma}_{k})^{-1}+(\sigma^{2}(\mathbf{X}_{k}^{T}\mathbf{X}_{k})^{-1})^{-1}\right)^{-1}$
(11)

given $\\{\mathbf{X}_{i},\mathbf{y}_{i}\\}_{i\in[K]}$, again with Lemma 2 from
[25] we have

$\displaystyle\mathbf{w}_{k}=\overline{\Sigma}_{k}(\sigma^{2}(\mathbf{X}_{k}^{T}\mathbf{X}_{k})^{-1})^{-1}\widehat{\mathbf{w}}_{k}+\overline{\Sigma}_{k}(\widetilde{\Sigma}_{k})^{-1}\overline{\theta}_{\setminus
k}+\vartheta_{k}$ (12)

where $\vartheta_{k}\sim\mathcal{N}(0,\overline{\Sigma}_{k})$. The term for
$\mathbf{w}_{k}$ in 12 uses the fact that $\widehat{\mathbf{w}}_{k}$ is a
noisy observation of $\mathbf{w}_{k}$ with additive noise of zero mean and
covariance $\sigma^{2}(\mathbf{X}_{k}^{T}\mathbf{X}_{k})^{-1}$, and
$\overline{\theta}_{\setminus k}$ is a noisy observation of $\theta$ with
covariance $\overline{\Sigma}_{\setminus k}$. Given all the training samples
from $K$ devices, $\mathbf{w}_{k}$ in 12 is Bayes optimal.

With PerFed-CKT, following 1, we solve the following objective:

$\displaystyle\min_{\mathbf{w}_{k}}\|\mathbf{X}_{k}\mathbf{w}_{k}-\mathbf{y}_{k}\|_{2}^{2}+\lambda_{k}\|\overline{\mathbf{s}}_{k}-s(\mathbf{w}_{k})\|_{2}^{2}$
(13)

where $\lambda_{k}$ is the regularization term as in 1 and
$\overline{\mathbf{s}}_{k}$ and $s(\mathbf{w}_{k})$ each is comparative to the
$\overline{\mathbf{s}}_{k}(\mathbf{x})$ and $s(\mathbf{w}_{k},\mathbf{x})$ in
1 for a single public data point $\mathbf{x}$. Note that in the setting of
linear regression we can set $s(\mathbf{w}_{k})=\mathbf{P}\mathbf{w}_{k}$
where $\mathbf{P}\in\mathbb{R}^{1\times d}$ is the public data (without loss
of generality, we assume single data point for the public data for
simplicity). Accordingly, we set
$\overline{\mathbf{s}}_{k}=\sum_{i=1}^{K}\alpha_{k,i}s(\widehat{\mathbf{w}}_{i})$
for an arbitrary set of weights $\alpha_{k,i},i\in[K]$ for client $k$. Then we
have that the local empirical risk minimizer for 13 is

$\displaystyle\begin{aligned}
\widetilde{\mathbf{w}}_{k}=(\mathbf{X}_{k}^{T}\mathbf{X}_{k}+\lambda_{k}\mathbf{P}^{T}\mathbf{P})^{-1}(\mathbf{X}_{k}^{T}\mathbf{X}_{k}\widehat{\mathbf{w}}_{k}+\lambda_{k}\mathbf{P}^{T}\mathbf{P}\sum_{i=1}^{K}\alpha_{k,i}\widehat{\mathbf{w}}_{i})\end{aligned}$
(14)

Finally, we present the optimal $\lambda_{k}^{*}$ and $\alpha_{k,i}^{*}$ for
any device $k\in[K]$ given the above linear regression problem with PerFed-CKT
in Theorem 4.2.

######  Theorem 4.2.

Assuming $\mathbf{X}_{k}^{T}\mathbf{X}_{k}=\beta\mathbf{I}_{d}$ and
$\mathbf{P}^{T}\mathbf{P}=\nu\mathbf{I}_{d}$ for some constant $\beta,\nu$,
the $\lambda_{k}^{*}$ and $\alpha_{k,i}^{*},i\in[K]$ that minimizes the test
performance on device $k,~{}k\in[K]$ i.e.,

$\displaystyle\lambda_{k}^{*},\alpha_{k,i}^{*},i\in[K]=\operatorname*{arg\,min}_{\lambda_{k},\alpha_{k,i},i\in[K]}\mathbb{E}[F_{k}(\widetilde{\mathbf{w}}_{k})|\widehat{\mathbf{w}}_{k},\overline{\theta}_{\setminus
k}]$ (15)

we have that

$\displaystyle\lambda_{k}^{*}=\sigma^{2}/\upsilon_{k}^{2}\nu,~{}\alpha_{k,i}^{*}=\frac{B_{k}}{\sigma^{2}+\beta\upsilon_{i}^{2}}$
(16)

with $A_{k}=\left(\sum_{i\in[K],i\neq
k}\frac{1}{\sigma^{2}+\beta\upsilon_{i}^{2}}\right)^{-1},~{}B_{k}=\frac{A_{k}(\sigma^{2}+\beta\upsilon_{k}^{2})}{\sigma^{2}+A_{k}\beta\upsilon_{k}^{2}}$.

Theorem 4.2 shows that given the objective function in 13 and the
corresponding minimizer 14, in a data-heterogeneous scenario where
$\upsilon_{k},~{}k\in[K]$ is unique to each client $k$, we have that the
optimal weights $\alpha_{k,i}^{*},~{}i\in[K]$ for client $k$ is in fact
inversely proportional to $\upsilon_{i}$. Intuitively, this means that since
larger $\upsilon_{i}$ leads to a larger divergence from the original $\theta$
in $\mathbf{w}_{i}=\theta+\zeta_{i}$, giving lower weight $\alpha_{k,i}$ to
client $i$ improves generalization of the personalzied model. This gives new
insight into co-distillation for personalization in FL since previous work
[30, 34] only consider scenarios where the weight
$\alpha_{i,k}=1/K,~{}\forall~{}i,k\in[K]$ in a non-personalized FL setting.
The results also present strong motivation for clustered knowledge transfer
for personalized FL. The proof for Theorem 4.2 is presented in Appendix B.
Further discussions on the implications of Theorem 4.2 is presented in
Appendix C.

##  5 Experiments

For all experiments we randomly sample a fraction $(C)$ of clients from
$K=100$ clients per communication round for local training. For the sake of
simplicity and fair comparison across different benchmarks, we use do not
apply any momentum acceleration or weight decay to local training. Further
details of the experimental setup are in Appendix E.

### Experimental Setup

#### Datasets and models.

We evaluate PerFed-CKT with CIFAR10 [49] as the training/test dataset and
CIFAR100 [50] as the public dataset for image classification in mainly two
different scenarios: model homogeneity and heterogeneity. For model
homogeneity, VGG11 [51] is deployed for all clients. For model heterogeneity,
we sample one of VGG13/VGG11/CNN model architecture for each client with the
probability of a larger model getting assigned to a client is proportional to
the client’s dataset size (see Figure 2(a)). We partition data heterogeneously
amongst clients using the Dirichlet distribution $\text{Dir}_{K}(\alpha)$
[52], smaller $\alpha$ leads to higher data size imbalance and degree of label
skew across clients. We set $\alpha=0.01$ to emulate realistic FL scenarios
with large data-heterogeneity (see Figure 2(b)).

(a)

(b)

Figure 2: (a) Proportion of the models architectures deployed across clients
for the model heterogeneity scenario; (b) data-distribution with $\alpha=0.01$
for all clients where larger circle indicates larger dataset size for each
label 0-9 of CIFAR10.

#### Baselines.

We compare PerFed-CKT with SOTA FL algorithms designed to efficiently train
either (i) a single global model at the server (e.g. FedAvg, FedProx,
Scaffold, FedDF) or (ii) personalized model(s) either at the server side as a
global model (GM) or client side (e.g., Per-FedAvg, FedFomo, Ditto,
HypCluster) as a local model (LM). Note that we do not assume a labeled public
dataset, and instead relax the condition to a small222We only use 2000
unlabeled data-samples per experiment that are sampled uniformly at random
from CIFAR100. and unlabeled public dataset and therefore exclude comparison
to methods which require training directly on a labeled public dataset (e.g.,
FedMD).

Table 1: Average test accuracy across the entire clients and total
communication cost (number of parameters communicated per round) for the
model-homogeneous scenario with total number of clients $K=100$. The standard
deviation for the test accuracy across random seeds is shown in the
parenthesis.

|  | $C=0.10$ | $C=0.15$
---|---|---|---
Method | Algorithm | Test Acc. | Com.-Cost | Test Acc. | Com.-Cost
Local Training | - | $64.02~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.48)}$ | - | $64.02~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.48)}$ | -
Non- Personalized | FedAvg | $15.53~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.42)}$ |  | $20.00~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 2.66)}$ |
| FedProx | $13.64~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.23)}$ |  | $16.71~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.79)}$ |
| Scaffold | $12.92~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.46)}$ |  | $16.41~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.14)}$ |
| FedDF | $15.18~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.18)}$ |  | $17.09~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.54)}$ |
Personalized | Per-FedAvg (GM) | $14.47~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.59)}$ |  | $14.61~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.83)}$ |
| Per-FedAvg (LM) | $51.74~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.52)}$ |  | $50.25~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.61)}$ |
| Ditto (LM) | $67.21~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.86)}$ |  | $68.88~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.95)}$ |
| Ditto (GM) | $21.63~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 2.13)}$ |  | $19.16~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 2.34)}$ |
| FedFomo |  74.62 ($\pm 0.42$) | $\bm{\mathrm{3900\times 10^{7}}}$ |  77.56 ($\pm 0.75$) | $\bm{\mathrm{5850}\times 10^{7}}$
|  HypCluster ($c=2$) | $34.11~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 2.63)}$ | $2340\times 10^{7}$ | $28.70~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 3.03)}$ | $3320\times 10^{7}$
|  HypCluster ($c=3$) | $39.17~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 2.64)}$ | $2540\times 10^{7}$ | $41.99~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 2.48)}$ | $3510\times 10^{7}$
|  HypCluster ($c=5$) | $52.51~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.37)}$ | $2930\times 10^{7}$ | $51.92~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.64)}$ | $3900\times 10^{7}$
|  HypCluster ($c=6$) | $65.77~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 2.76)}$ | $3120\times 10^{7}$ | $63.28~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.16)}$ | $4100\times 10^{7}$
|  PerFed-CKT ($c=1$) | $70.70~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.46)}$ | $4.4\times 10^{7}$ | $67.66~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.30)}$ | $6.4\times 10^{7}$
|  PerFed-CKT ($c=2$) | $73.33~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.26)}$ | $4.8\times 10^{7}$ | $70.86~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.73)}$ | $6.8\times 10^{7}$
|  PerFed-CKT ($c=3$) |  74.31 ($\pm 0.40$) | $\bm{\mathrm{5.2\times 10^{7}}}$ |  76.74 ($\pm 0.71$) | $\bm{\mathrm{7.2\times 10^{7}}}$
|  PerFed-CKT ($c=4$) | $72.67~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.31)}$ | $5.6\times 10^{7}$ | $73.52~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.15)}$ | $7.6\times 10^{7}$

$2150\times 10^{7}$

$3120\times 10^{7}$

Table 2: Average test accuracy across all clients and total communication cost (number of parameters communicated per round) for the model-heterogeneous scenario with total number of clients $K=100$. The standard deviation for the test accuracy across random seeds is shown in the parenthesis. |  | $C=0.10$ | $C=0.15$
---|---|---|---
Method | Algorithm | Test Acc. | Com.-Cost | Test Acc. | Com.-Cost
Local Training | - | $59.29~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.05)}$ | - | $59.29~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.05)}$ | -
Non-Personalized | FedDF | $15.25~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.21)}$ | $1780\times 10^{7}$ | $14.16~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 1.11)}$ | $2620\times 10^{7}$
Personalized |  PerFed-CKT ($c=1$) | $70.94~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.46)}$ | $4.4\times 10^{7}$ | $70.77~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.14)}$ | $6.4\times 10^{7}$
|  PerFed-CKT ($c=2$) |  72.25 ($\pm 0.17$) | $\bm{\mathrm{4.8\times 10^{7}}}$ |  76.14 ($\pm 0.73$) | $\bm{\mathrm{6.8\times 10^{7}}}$
|  PerFed-CKT ($c=3$) | $71.84{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.40)}$ | $5.2\times 10^{7}$ |  76.02 ($\pm 0.66$) | $7.2\times 10^{7}$
|  PerFed-CKT ($c=4$) | $70.01~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.31)}$ | $5.6\times 10^{7}$ | $73.14~{}{\color[rgb]{.25,.25,.25}\definecolor[named]{pgfstrokecolor}{rgb}{.25,.25,.25}\pgfsys@color@gray@stroke{.25}\pgfsys@color@gray@fill{.25}(\pm 0.61)}$ | $7.6\times 10^{7}$

### Experimental Results

We demonstrate the efficacy of PerFed-CKT in terms of the average test
accuracy across all clients with the communication cost defined as the total
number of parameters communicated across server/clients during the training
process including uplink and downlink.

#### Test Accuracy and Communication Cost.

In LABEL:fig:table1, we show the performance of PerFed-CKT along with the
performance of other SOTA FL algorithms in regards to the achieved highest
test accuracy and communicated number of parameters between server and client
with different fractions of selected clients $C$. For $C=0.1$, we show that
PerFed-CKT achieves high test accuracy of $74.31\%$ with $c=3$, with small
communication cost compared to other algorithms (saving at maximum $\times
750$). FedFomo achieves a slightly higher test accuracy performance with
$74.62\%$, but the communication cost spent, $3900\times 10^{7}$ parameters,
is significantly larger compared to PerFed-CKT which is $5.2\times 10^{7}$
parameters. Moreover note that algorithms that train a single global model in
the non-personalized FL setting performs worse than personalized algorithms
showing that the traditional FL framework does not perform well to individual
clients in the setting of high data heterogeneity. For $C=0.15$, we also show
that PerFed-CKT is able to achieve a comparable high test accuracy of
$76.74\%$ with only a small communication cost of $7.2\times 10^{7}$
parameters (saving at maximum $\times 812.5$) where FedFomo achieves a
slightly higher accuracy of $77.56\%$ with large communication cost of
$5850\times 10^{7}$.

#### Model Heterogeneity.

We demonstrate the performance of PerFed-CKT where clients have different
models dependent on their dataset size (see Figure 2(b)) in LABEL:fig:table2.
Note that this is a realistic setting of FL where clients can have smaller or
larger models dependent on their dataset size or system capabilities. Only
FedDF is capable for model heterogeneity amongst the SOTA FL algorithms which
we included for comparison. With model heterogeneity, PerFed-CKT achieves high
test accuracy $72.25\%$ for $C=0.1$ and $76.14\%$ for $C=0.15$ with even
smaller communication cost of $4.8\times 10^{7}$ and $6.8\times 10^{7}$
respectively. For $C=0.15$, the test accuracy is close to that of PerFed-CKT
for model homogeneity, showing that allowing model heterogeneity increases
feasibility while not hurting the local performance of clients.

#### Effect of Clustering.

PerFed-CKT performs clustering at the server side to cluster the logit
information received from the clients. We evaluate how the number of clusters
effects the test accuracy and communication cost. For both $C=0.1$ and
$C=0.15$ in LABEL:fig:table1, $c=3$ achieves the best test accuracy
performance. Increasing $c$ from that point actually deteriorates the
performance with higher communication cost. This shows that while clustering
can help to a certain extent, too much clustering can hurt generalization
since we are decreasing the number of clients in each cluster and the
diversity of information within each cluster. Similar behavior is observed in
LABEL:fig:table1 where the best test accuracy is achieved in $c=2$, and then
the accuracy decreases for higher $c$.

##  6 Concluding Remarks

The inherently high data and system heterogeneity across resource-constrained
clients in FL should be considered for devising realistic personalized FL
schemes. However, previous work in personalized FL restricted clients to have
homogeneous models across clients with direct communication of the model
parameters which can incur heavy communication cost. We propose PerFed-CKT
that caters to the data and system heterogeneity across clients by using
clustered knowledge transfer, allowing heterogeneous model deployment without
direct communication of the models. We show that PerFed-CKT achieves
competitive performance compared to other SOTA personalized FL schemes at a
much smaller communication cost. Interesting future directions of this work
include understanding the privacy implications that arise due to the
clustering of clients and communicating their logits instead of the actual
model parameters.

## References

  * [1] H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agøura y Arcas.  Communication-Efficient Learning of Deep Networks from Decentralized Data.  International Conference on Artificial Intelligenece and Statistics (AISTATS), April 2017.
  * [2] Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurelien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D’Oliveira, Salim El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adria Gascon, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaid Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konecny, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrede Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Ozgur, Rasmus Pagh, Mariana Raykova, Hang Qi, Daniel Ramage, Ramesh Raskar, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramer, Praneeth Vepakomma, Jianyu Wang, Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao.  Advances and open problems in federated learning.  arXiv preprint arXiv:1912.04977, 2019.
  * [3] Keith Bonawitz, Hubert Eichner, Wolfgang Grieskamp, Dzmitry Huba, Alex Ingerman, Vladimir Ivanov, Chloe Kiddon, Jakub Konecny, Stefano Mazzocchi, H. Brendan McMahan, Timon Van Overveldt, David Petrou, Daniel Ramage, and Jason Roselander.  Towards Federated Learning at Scale: System Design.  SysML, April 2019.
  * [4] Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly, Deepesh Data, et al.  A field guide to federated optimization.  arXiv preprint arXiv:2107.06917, 2021.
  * [5] Hao Yu, Sen Yang, and Shenghuo Zhu.  Parallel restarted SGD for non-convex optimization with faster convergence and less communication.  The Thirty-Third AAAI Conference on Artificial Intelligence (AAAI-19), 2019.
  * [6] Sebastian U Stich.  Local SGD converges fast and communicates little.  In International Conference on Learning Representations (ICLR), 2019\.
  * [7] Jianyu Wang and Gauri Joshi.  Cooperative SGD: A unified framework for the design and analysis of communication-efficient SGD algorithms.  Journal of Machine Learning Research (JMLR), 2021.
  * [8] Yae Jee Cho, Jianyu Wang, and Gauri Joshi.  Client selection in federated learning: Convergence analysis and power-of-choice selection strategies.  abs/2010.01243, 2020.
  * [9] Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečnỳ, Sanjiv Kumar, and H Brendan McMahan.  Adaptive federated optimization.  In International Conference on Learning Representations (ICLR), 2021\.
  * [10] Farzin Haddadpour and Mehrdad Mahdavi.  On the convergence of local descent methods in federated learning.  arXiv preprint arXiv:1910.14425, 2019.
  * [11] A Khaled, K Mishchenko, and P Richtárik.  Tighter theory for local SGD on identical and heterogeneous data.  In The 23rd International Conference on Artificial Intelligence and Statistics (AISTATS 2020), 2020.
  * [12] Sebastian U Stich and Sai Praneeth Karimireddy.  The error-feedback framework: Better rates for SGD with delayed gradients and compressed communication.  Journal of Machine Learning Research (JMLR), 2020.
  * [13] Blake Woodworth, Kumar Kshitij Patel, Sebastian U Stich, Zhen Dai, Brian Bullins, H Brendan McMahan, Ohad Shamir, and Nathan Srebro.  Is local SGD better than minibatch SGD?  In Proceedings of the 37th International Conference on Machine Learning, 2020.
  * [14] Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian U Stich.  A unified theory of decentralized SGD with changing topology and local updates.  In Proceedings of 37th International Conference on Machine Learning, 2020.
  * [15] Zhouyuan Huo, Qian Yang, Bin Gu, Lawrence Carin, and Heng Huang.  Faster on-device training using new federated momentum algorithm.  arXiv preprint arXiv:2002.02090, 2020.
  * [16] Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu.  FedPD: A federated learning framework with optimal rates and adaptivity to non-IID data.  Asilomar Conference on Signals, Systems, and Cmputers., 2020.
  * [17] Reese Pathak and Martin J Wainwright.  FedSplit: An algorithmic framework for fast federated optimization.  In Advances in Neural Information Processing Systems, 2020.
  * [18] Grigory Malinovsky, Dmitry Kovalev, Elnur Gasanov, Laurent Condat, and Peter Richtárik.  From local SGD to local fixed point methods for federated learning.  In Proceedings of the 37th International Conference on Machine Learning, 2020.
  * [19] Anit Kumar Sahu, Tian Li, Maziar Sanjabi, Manzil Zaheer, Ameet Talwalkar, and Virginia Smith.  Federated optimization for heterogeneous networks.  In Proceedings of the 3rd MLSys Conference, January 2020.
  * [20] Paul Pu Liang, Terrance Liu, Liu Ziyin, , Nicholas B. Allen, Randy P. Auerbach, David Brent, Ruslan Salakhutdinov, and Louis-Philippe Morency.  Think locally, act globally: Federated learning with local and global representations.  In nternational Workshop on Feder-ated Learning for User Privacy and Data Confidentiality inConjunction with NeurIPS 2019 (FL-NeurIPS’19), 2019\.
  * [21] Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran.  An efficient framework for clustered federated learning.  In Advances in Neural Information Processing Systems, 2020.
  * [22] Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang.  Towards personalized federated learning.  CoRR, abs/2103.00710, 2020.
  * [23] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar.  Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach.  In Advances in Neural Information Processing Systems, 2020.
  * [24] M. Zhang, K. Sapra, S. Fidler, S. Yeung, and J. M. Alvarez.  Personalized federated learning with first order model optimization.  In International Conference on Learning Representations (ICLR), 2021\.
  * [25] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith.  Ditto: Fair and robust federated learning through personalization.  In Proceedings of the 38th International Conference on Machine Learning, 2021.
  * [26] Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh.  Three approaches for personalization with applications to federated learning.  arXiV preprint cs.LG 2002.10619, 2020.
  * [27] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  In International Conference on Learning Respresentations (ICLR), 2017.
  * [28] Chaoyang He, Murali Annavaram, and Salman Avestimehr.  Group knowledge transfer: Federated learning of large cnns at the edge.  In Advances in Neural Information Processing Systems, 2020.
  * [29] S. Sodhani, O. Delalleau, M. Assran, K. Sinha, N. Ballas, and M. Rabbat.  A closer look at codistillation for distributed training.  CoRR, abs/2010.02838, 2020.
  * [30] I. Bistritz, A. J. Mann, and N. Bambos.  Distributed distillation for on-device learning.  In Advances in Neural Information Processing Systems, 2020.
  * [31] Tao Lin, Lingjing Kong, Sebastian U. Stich, and Martin Jaggi.  Ensemble distillation for robust model fusion in federated learning.  In Advances in Neural Information Processing Systems, 2020.
  * [32] Rohan Anil, Gabriel Pereyra, Alexandre Passos, Robert Ormandi, George E. Dahl, and Geoffrey E. Hinton.  Large scale distributed neural network training through online distillation.  In International Conference on Learning Representations (ICLR), 2018\.
  * [33] Canh T. Dinh, Nguten H. Tran, and Tuan Dung Nguyen.  Personalized federated learning with moreau envelopes.  In Advances in Neural Information Processing Systems, 2020.
  * [34] Daliang Li and Junpu Wang.  Fedmd: Heterogenous federated learning via model distillation.  In International Workshop on Feder-ated Learning for User Privacy and Data Confidentiality inConjunction with NeurIPS 2019 (FL-NeurIPS’19), 2019.
  * [35] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  ArXiv, March 2015.
  * [36] Jiaxin Ma, Ryo Yonetani, and Zahid Iqbal.  Adaptive distillation for decentralized learning from heterogeneous clients.  In 2020 25th International Conference on Pattern Recognition (ICPR), August 2020.
  * [37] Sohei Itahara, Takayuki Nishio, Yusuke Koda, Masahiro Morikura, and Koji Yamamoto.  Distillation-based semi-supervised federated learning for communication-efficient collaborative training with non-iid private data.  IEEE Transactions on Mobile Computing, August 2020.
  * [38] Lichao Sun and Lingjuan Lyu.  Federated model distillation with noise-free differential privacy.  In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence (IJCAI-21), May 2021.
  * [39] Qinbin Li, Bingsheng He, and Dawn Song.  Practical one-shot federated learning for cross-silo setting.  In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence (IJCAI-21), May 2021.
  * [40] Yanlin Zhou, George Pu, Xiyao Ma, Xiaolin Li, and Dapeng Wu.  Distilled one-shot federated learning.  ArXiv, June 2021.
  * [41] Sangho Lee, Kiyoon Yoo, and Nojun Kwak.  Edge bias in federated learning and its solution by buffered knowledge distillation.  ArXiv, February 2021.
  * [42] Eunjeong Jeong, Seungeun Oh, Hyesung Kim, Jihong Park, Mehdi Bennis, and Seong-Lyun Kim.  International Workshop on Machine Learning on the Phone and other Consumer Devices in Conjunction with NeurIPS (NeurIPS-MLPCD), 2018.
  * [43] X. Lan, X. Zhu, and S. Gong.  Knowledge distillation by on-the-fly native ensemble.  In Proceedings of the 32nd International Conference on Neural Information Processing Systems., pages 7528–7538, 2018.
  * [44] Lingjuan Lyu, Han Yu, Xingjun Ma, Lichao Sun, Jun Zhao, Qiang Yang, and Philip S. Yu.  Privacy and robustness in federated learning: Attacks and defenses.  arXiv preprint arXiv:2012.06337, 2020.
  * [45] Y. Zhang, T. Xiang, T. M. Hospedales, and H. Lu.  Deep mutual learning.  In In IEEE Conf. on Computer Vision and Pattern Recognition (CVPR), page 4320–4328, 2018.
  * [46] Haoran Zhang, Zhenzhen Hu, Wei Qin, Mingliang Xu, and Meng Wang.  Adversarial co-distillation learning for image recognition.  Pattern Recognition, 111:107659, 2021.
  * [47] Xinyang Lin, , Hanting Chen, Yixing Xu, Chao Xu, Xiaolin Gui, Yiping Deng, and Yunhe Wang.  Federated Learning with Positive and Unlabeled Data.  preprint, 2021.
  * [48] J. A. Hartigan and M. A. Wong.  Algorithm as 136: A k-means clustering algorithm.  Journal of the Royal Statistical Society. Series C (Applied Statistics), 28(1):100–108, 1979.
  * [49] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton.  Learning multiple layers of features from tiny images.  CIFAR-10 (Canadian Institute for Advanced Research), 2009.
  * [50] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton.  Cifar-100 (canadian institute for advanced research).  http://www.cs.toronto.edu/~kriz/cifar.html, 2009.
  * [51] Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  In International Conference on Learning Representations (ICLR), 2015\.
  * [52] Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown.  Measuring the effects of non-identical data distribution for federated visual classification.  In International Workshop on Federated Learning for User Privacy and Data Confidentiality in Conjunction with NeurIPS 2019 (FL-NeurIPS’19), December 2019.
  * [53] Tomaso Poggio, Stephen Voinea, and Lorenzo Rosasco.  Online learning, stability, and stochastic gradient descent, 2019.
  * [54] Ya.I. Alber, A.N. Iusem, and M.V. Solodovz.  On the projected subgradient method for nonsmooth convex optimization in a hilbert space.  Mathematical Programming, 81(1):23–25, 1998.
  * [55] S. Ben-David, J. Blitzer, K. Crammer, A. Kulesza, F. Pereira, and J. W. Vaughan.  A theory of learning from different domains.  Machine Learning, 79(1-2):151–175, 2009.

##  Appendix A Proof for Theorem 4.1

In this section we present the proof for Theorem 4.1. We follow the techniques
presented by [30] for the proof. For notational simplicity, we notate all
super subscript $(t,0)$ as $(t)$ throughout the proof, dropping the local
iteration index. We define the following $\sigma$-algebra on the set that
contains the history of the model updates for all clients with
$\mathbf{w}^{(t)}=[\mathbf{w}_{1}^{(t)}...\mathbf{w}_{K}^{(t)}]$ and
$\overline{\mathbf{s}}^{(t)}=[\overline{\mathbf{s}}_{1}^{(t)}...\overline{\mathbf{s}}_{K}^{(t)}]$
as
$\mathcal{H}_{t}=\sigma(\\{\mathbf{w}^{(i)},\overline{\mathbf{s}}^{(i)}\\}~{}|~{}i\leq
t)$. Recall
$\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\triangleq\frac{1}{|\xi_{k}^{(t)}|}\sum_{\xi\in\xi_{k}^{(t)}}\nabla
f(\mathbf{w}_{k}^{(t)},\xi)+\frac{2\lambda}{|\mathcal{P}_{k}^{(t)}|}\sum_{\mathbf{x}\in\mathcal{P}_{k}^{(t)}}\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})^{T}\left(s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x})\right)$.

### Additional Lemmas

We first present useful Lemmas and their proofs which we use for the
intermediate steps in the main proof for Theorem 4.1.

######  Lemma A.1.

The gradient of the second term in
$\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})$ with respect
to $\mathbf{w}_{k}^{(t)}$ is Lipschitz continuous, and therefore with 4.1,
$\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})$ is also a
Lipschitz-smooth function with factor $L_{p}$.

###### Proof.

With dropping the iteration index $t$ for the upper script for simplicity,
let’s define the second term in
$\Phi_{k}(\mathbf{w}_{k};\overline{\mathbf{s}}_{k})$ as
$q(\mathbf{w}_{k};\overline{\mathbf{s}}_{k})\triangleq\frac{\lambda}{|\mathcal{P}|}\sum_{\mathbf{x}\in\mathcal{P}}\|\overline{\mathbf{s}}(\mathbf{x})-s(\mathbf{w}_{k},\mathbf{x})\|_{2}^{2}$.
Then we have that

$\displaystyle\nabla_{\mathbf{w}_{k}}q(\mathbf{w}_{k};\overline{\mathbf{s}}_{k})=\frac{2\lambda}{|\mathcal{P}|}\sum_{\mathbf{x}\in\mathcal{P}}\nabla
s(\mathbf{w}_{k},\mathbf{x})^{T}\left(s(\mathbf{w}_{k},\mathbf{x})-\overline{\mathbf{s}}_{k}(\mathbf{x})\right)$
(17)

For an arbitrary $\mathbf{e}_{k}$ in the domain of
$q(\cdot;\overline{\mathbf{s}}_{k})$ for each $\mathbf{x}\in\mathcal{P}$, we
have that

$\displaystyle\begin{aligned} &\|\nabla
s(\mathbf{w}_{k},\mathbf{x})^{T}\left(s(\mathbf{w}_{k},\mathbf{x})-\overline{\mathbf{s}}_{k}(\mathbf{x})\right)-\nabla
s(\mathbf{e}_{k},\mathbf{x})^{T}(s(\mathbf{e}_{k},\mathbf{x})-\overline{\mathbf{s}}_{k}(\mathbf{x}))\|^{2}\\\
&\leq 3\|\nabla
s(\mathbf{w}_{k},\mathbf{x})^{T}\left(s(\mathbf{w}_{k},\mathbf{x})-\overline{\mathbf{s}}_{k}(\mathbf{x})\right)-\nabla
s(\mathbf{w}_{k},\mathbf{x})^{T}\left(s(\mathbf{e}_{k},\mathbf{x})-\overline{\mathbf{s}}_{k}(\mathbf{x})\right)\|^{2}\\\
&\quad+3\|\nabla
s(\mathbf{w}_{k},\mathbf{x})^{T}(s(\mathbf{e}_{k},\mathbf{x})-\overline{\mathbf{s}}_{k}(\mathbf{x}))-\nabla
s(\mathbf{e}_{k},\mathbf{x})^{T}\left(s(\mathbf{e}_{k},\mathbf{x})-\overline{\mathbf{s}}_{k}(\mathbf{x})\right)\|^{2}\end{aligned}$
(18) $\displaystyle\leq 3\|\nabla
s(\mathbf{w}_{k},x)\|^{2}\|s(\mathbf{w}_{k},x)-s(\mathbf{e}_{k},x)\|^{2}+3\|\nabla
s(\mathbf{w}_{k},x)-\nabla
s(\mathbf{e}_{k},x)\|^{2}\|s(\mathbf{e}_{k},x)-\overline{\mathbf{s}}_{k}(x)\|^{2}$
(19) $\displaystyle\leq
3L_{s}^{4}\|\mathbf{w}_{k}-\mathbf{e}_{k}\|^{2}+6L_{g}^{2}\|\mathbf{w}_{k}-\mathbf{e}_{k}\|^{2}$
(20)
$\displaystyle=(3L_{s}^{4}+6L_{g}^{2})\|\mathbf{w}_{k}-\mathbf{e}_{k}\|^{2}$
(21)

where 18 uses Jensen’s inequality for the $\ell_{2}$-norm for three terms, 19
uses the submultiplicativity of the norm, and the LHS of 20 uses 4.5.
Therefore we can conclude that $\|\nabla
s(\mathbf{w}_{k},\mathbf{x})^{T}(s(\mathbf{w}_{k},\mathbf{x})-\overline{\mathbf{s}}_{k}(\mathbf{x}))-\nabla
s(\mathbf{e}_{k},\mathbf{x})^{T}(s(\mathbf{e}_{k},\mathbf{x})-\overline{\mathbf{s}}_{k}(\mathbf{x}))\|^{2}$
for any $\mathbf{x}$ is Lipschitz-continuous, and hence
$\nabla_{\mathbf{w}_{k}}q(\mathbf{w}_{k};\overline{\mathbf{s}}_{k})$ is also
Lipschitz-continuous. ∎

######  Lemma A.2.

We have that
$\mathbb{E}[\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})|\mathcal{H}_{t}]=\nabla_{\mathbf{w}_{k}^{(t)}}\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})$
and
$\mathbb{E}[\|\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|_{2}^{2}]\leq
2G^{2}+16\lambda^{2}L_{s}^{2}$ and
$\|\nabla_{\mathbf{w}_{k}^{(t)}}\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|$
and $\|\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|$
is each bounded by constant $M_{1}\geq 0$ and $M_{2}\geq 0$.

###### Proof.

By definition of the gradient we have

$\displaystyle\begin{aligned}
&\mathbb{E}[\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})|\mathcal{H}_{t}]\\\
&=\mathbb{E}[\frac{1}{|\xi_{k}^{(t)}|}\sum_{\xi\in\xi_{k}^{(t)}}\nabla
f(\mathbf{w}_{k}^{(t)},\xi)+\frac{2\lambda}{|\mathcal{P}_{k}^{(t)}|}\sum_{\mathbf{x}\in\mathcal{P}_{k}^{(t)}}\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})^{T}(s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x}))|\mathcal{H}_{t}]\end{aligned}$
(22) $\displaystyle\begin{aligned} =\nabla
F_{k}(\mathbf{w}_{k}^{(t)})+\frac{2\lambda}{|\mathcal{P}|}\sum_{\mathbf{x}\in\mathcal{P}}\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})^{T}(s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x}))\end{aligned}$
(23)
$\displaystyle=\nabla_{\mathbf{w}_{k}^{(t)}}\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})$
(24)

finishing the proof for the first part of lemma A.2. Next, we prove the second
part of lemma A.2 showing that

$\displaystyle\begin{aligned}
&\mathbb{E}[\|\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|^{2}]\\\
&=\mathbb{E}[\|\frac{1}{|\xi_{k}^{(t)}|}\sum_{\xi\in\xi_{k}^{(t)}}\nabla
f(\mathbf{w}_{k}^{(t)},\xi)+\frac{2\lambda}{|\mathcal{P}_{k}^{(t)}|}\sum_{\mathbf{x}\in\mathcal{P}_{k}^{(t)}}\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})^{T}(s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x}))\|^{2}]\end{aligned}$
(25) $\displaystyle\leq
2\mathbb{E}[\|\frac{1}{|\xi_{k}^{(t)}|}\sum_{\xi\in\xi_{k}^{(t)}}\nabla
f(\mathbf{w}_{k}^{(t)},\xi)\|^{2}]+2\mathbb{E}[\|\frac{2\lambda}{|\mathcal{P}_{k}^{(t)}|}\sum_{\mathbf{x}\in\mathcal{P}_{k}^{(t)}}\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})^{T}(s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x}))\|^{2}]$
(26)
$\displaystyle\leq\mathbb{E}[\frac{8\lambda^{2}}{|\mathcal{P}_{k}^{(t)}|}\sum_{\mathbf{x}\in\mathcal{P}_{k}^{(t)}}\|\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})^{T}(s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x}))\|^{2}]+2G^{2}$
(27)
$\displaystyle\leq\frac{8\lambda^{2}}{|\mathcal{P}|}\sum_{\mathbf{x}\in\mathcal{P}}\mathbb{E}[\|\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})\|^{2}\|s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x})\|^{2}]+2G^{2}$
(28) $\displaystyle\leq 2G^{2}+16\lambda^{2}L_{s}^{2}$ (29)

where 26 is due to the Cauchy–Schwarz inequality and AM-GM inequality, 27 is
due to 4.4 and Jensen’s inequality, 28 is due to the submultiplicativity of
the norm, and 29 is due to 4.5 and that the maximum $\ell_{2}$-norm distance
between two probability vectors is $\sqrt{2}$.

Moreover,

$\displaystyle\|\nabla_{\mathbf{w}_{k}^{(t)}}\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|=\|\frac{2\lambda}{|\mathcal{P}|}\sum_{\mathbf{x}\in\mathcal{P}}\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})^{T}(s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x}))+\nabla
F_{k}(\mathbf{w}_{k}^{(t)})\|$ (30)
$\displaystyle\leq\frac{2\lambda}{|\mathcal{P}|}\|\sum_{\mathbf{x}\in\mathcal{P}}\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})^{T}(s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x}))\|+\|\nabla
F_{k}(\mathbf{w}_{k}^{(t)})\|$ (31) $\displaystyle\leq
L_{f}+{2\lambda}\sum_{\mathbf{x}\in\mathcal{P}}\|\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})^{T}(s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x}))\|$
(32) $\displaystyle\leq
L_{f}+{2\lambda}\sum_{\mathbf{x}\in\mathcal{P}}\|\nabla
s(\mathbf{w}_{k}^{(t)},\mathbf{x})\|\|(s(\mathbf{w}_{k}^{(t)},\mathbf{x})-\overline{\mathbf{s}}_{k}^{(t)}(\mathbf{x}))\|$
(33) $\displaystyle\leq L_{f}+{2\sqrt{2}\lambda}L_{s}|\mathcal{P}|=M_{1}$ (34)

and therefore
$\left\|\nabla_{\mathbf{w}_{k}^{(t)}}\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\right\|$
is bounded by $M_{1}\geq 0$. With similar steps we can show that
$\|\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|\leq
M_{2}$ for a certain constant $M_{2}\geq 0$. ∎

### Main Proof for Theorem 4.1

Using Lemma A.1 and Lemma A.2 we have that

$\displaystyle\mathbb{E}[\Phi_{k}(\mathbf{w}_{k}^{(t+1)};\overline{\mathbf{s}}_{k}^{(t)})|\mathcal{H}_{t}]=\mathbb{E}[\Phi_{k}(\mathbf{w}_{k}^{(t)}-\eta_{t}\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)}));\overline{\mathbf{s}}_{k}^{(t)})|\mathcal{H}_{t}]$
(35)
$\displaystyle\leq\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})+\frac{\eta_{t}^{2}L_{p}}{2}\mathbb{E}[\|\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|^{2}|\mathcal{H}_{t}]-\eta_{t}\nabla_{\mathbf{w}_{k}^{(t)}}\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})^{T}\mathbb{E}[\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})|\mathcal{H}_{t}]$
(36)
$\displaystyle\leq\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})-\eta_{t}\|\nabla_{\mathbf{w}_{k}^{(t)}}\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|^{2}+\frac{\eta_{t}^{2}L_{p}}{2}(2G^{2}+16\lambda^{2}L_{s}^{2})$
(37)

Assuming $\sum_{t=0}^{\infty}\eta_{t}^{2}<\infty$ and
$\sum_{t=0}^{\infty}\eta_{t}=\infty$, and applying Robbins-Siegmund Theorem
(Theorem B.1. in [53]) on 37, we have that with probability 1,

$\displaystyle\sum_{t=1}^{\infty}\eta_{t}\|\nabla_{\mathbf{w}_{k}^{(t)}}\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|^{2}<\infty$
(38)

Now we can show

$\displaystyle\|\nabla\Phi_{k}(\mathbf{w}_{k}^{(t+1)};\overline{\mathbf{s}}_{k}^{(t)})\|^{2}-\|\nabla\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|^{2}$
(39)
$\displaystyle=(\|\nabla\Phi_{k}(\mathbf{w}_{k}^{(t+1)};\overline{\mathbf{s}}_{k}^{(t)})\|+\|\nabla\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|)(\|\nabla\Phi_{k}(\mathbf{w}_{k}^{(t+1)};\overline{\mathbf{s}}_{k}^{(t)})\|-\|\nabla\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|)$
(40) $\displaystyle\leq
2M_{1}(\|\nabla\Phi_{k}(\mathbf{w}_{k}^{(t+1)};\overline{\mathbf{s}}_{k}^{(t)})\|-\|\nabla\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|)$
(41) $\displaystyle\leq
2M_{1}\|\nabla\Phi_{k}(\mathbf{w}_{k}^{(t+1)};\overline{\mathbf{s}}_{k}^{(t)})-\nabla\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|$
(42) $\displaystyle\leq
2M_{1}\|\eta_{t}\mathbf{g}_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|\leq
2M_{1}M_{2}\eta_{t}$ (43)

Finally, using Proposition 2 in [54] we have that for $t\rightarrow\infty$,
$\|\nabla_{\mathbf{w}_{k}^{(t)}}\Phi_{k}(\mathbf{w}_{k}^{(t)};\overline{\mathbf{s}}_{k}^{(t)})\|\rightarrow
0$ with probability 1.

##  Appendix B Proof for Theorem 4.2

We have that 13 is equal to:

$\displaystyle\widetilde{\mathbf{w}}_{k}=\frac{1}{1+\lambda\nu/\beta}\widehat{\mathbf{w}}_{k}+\frac{1}{1+\beta/\lambda\nu}\sum_{i=1}^{K}\alpha_{k,i}\widehat{\mathbf{w}}_{i}$
(44)

and the Bayes optimal $\mathbf{w}_{k}$ in 12 becomes

$\displaystyle\mathbf{w}_{k}=\frac{1}{1+\sigma^{2}/\beta\upsilon_{k}^{2}}\widehat{\mathbf{w}}_{k}+\frac{A_{k}\sigma^{2}}{\sigma^{2}+A_{k}\beta\upsilon_{k}^{2}}\sum_{i=1}^{K}\frac{1}{\sigma^{2}+\beta\upsilon_{i}^{2}}\widehat{\mathbf{w}}_{i}+\varsigma_{k}$
(45)

where $A_{k}=\left(\sum_{i\in[K],i\neq
k}\frac{1}{\sigma^{2}+\beta\upsilon_{i}^{2}}\right)^{-1}$ and
$\varsigma_{k}\sim\mathcal{N}\left(0,(\frac{\beta}{A_{k}+\beta\upsilon_{k}^{2}}+\frac{\beta}{\sigma^{2}})^{-1}\right)$.
If we aim to find the $\lambda_{k}$ and $\alpha_{k,i},i\in[K]$ that minimizes
$\mathbb{E}[F_{k}(\widetilde{\mathbf{w}}_{k})]$ given
$\widehat{\mathbf{w}}_{k}$ and $\overline{\theta}_{\setminus k}$, in other
words,

$\displaystyle\lambda_{k}^{*},\alpha_{k,i}^{*},i\in[K]$
$\displaystyle=\operatorname*{arg\,min}_{\lambda_{k},\alpha_{k,i},i\in[K]}\mathbb{E}[F_{k}(\widetilde{\mathbf{w}}_{k})|\widehat{\mathbf{w}}_{k},\overline{\theta}_{\setminus
k}]$ (46)
$\displaystyle=\operatorname*{arg\,min}_{\lambda_{k},\alpha_{k,i},i\in[K]}\mathbb{E}[\|\mathbf{X}_{k}\widetilde{\mathbf{w}}_{k}-(\mathbf{X}_{k}\mathbf{w}_{k}+\mathbf{z})\|_{2}^{2}|\widehat{\mathbf{w}}_{k},\overline{\theta}_{\setminus
k}]$ (47)
$\displaystyle=\operatorname*{arg\,min}_{\lambda_{k},\alpha_{k,i},i\in[K]}\mathbb{E}[\|\mathbf{X}_{k}(\widetilde{\mathbf{w}}_{k}-\mathbf{w}_{k})\|_{2}^{2}|\widehat{\mathbf{w}}_{k},\overline{\theta}_{\setminus
k}]$ (48)
$\displaystyle=\operatorname*{arg\,min}_{\lambda_{k},\alpha_{k,i},i\in[K]}\mathbb{E}[\|\widetilde{\mathbf{w}}_{k}-\mathbf{w}_{k}\|_{2}^{2}|\widehat{\mathbf{w}}_{k},\overline{\theta}_{\setminus
k}]$ (49)

then taking 44 and 45 into 49 we have that

$\displaystyle\lambda_{k}^{*}=\sigma^{2}/\upsilon_{k}^{2}\nu$ (50)
$\displaystyle\alpha_{k,i}^{*}=\frac{B_{k}}{\sigma^{2}+\beta\upsilon_{i}^{2}}$
(51)

where
$B_{k}=\frac{A_{k}(\sigma^{2}+\beta\upsilon_{k}^{2})}{\sigma^{2}+A_{k}\beta\upsilon_{k}^{2}}$.

##  Appendix C Further discussion on Theorem 4.2

Theorem 4.2 presents insights on how to set the weights
$\\{\alpha_{k,i}\\}_{i\in[K]}$ and regularization weight $\lambda_{k}^{*}$ for
each client $k\in[K]$ from 1 to improve generalization with PerFed-CKT. While
in the main paper we discussed the implications of the optimal
$\\{\alpha_{k,i}^{*}\\}_{i\in[K]}$ in 16 and the motivation for clustering,
here we continue the discussion in regards to the optimal
$\\{\lambda_{k}^{*}\\}_{k\in[K]}$, i.e., the optimal regularization weight.
Recapping the linear regression setup from Section 4, we have $\theta$
uniformly distributed on $\mathbb{R}^{d}$, and each device $k\in[K]$ has its
data distributed with parameters $\mathbf{w}_{k}=\theta+\zeta_{k}$ where
$\zeta_{k}\sim\mathcal{N}(0,\upsilon_{k}^{2}\mathbf{I}_{d})$ and
$\mathbf{I}_{d}$ is the $d\times d$ identity matrix and $\upsilon_{k}$ is
unique to the client’s task. Suppose we have
$\mathbf{y}_{k}=\mathbf{X}_{k}\mathbf{w}_{k}+\mathbf{z},~{}k\in[K]$ where
$\mathbf{y}_{k}\in\mathbb{R}^{n},~{}\mathbf{X}_{k}\in\mathbb{R}^{n\times d}$,
and $\mathbf{z}\in\mathbb{R}^{n}$ such that
$\mathbf{z}\sim\mathcal{N}(0,\sigma^{2}\mathbf{I}_{d})$.

With PerFed-CKT, the optimal regularization weight is equal to
$\lambda_{k}^{*}=\sigma^{2}/\upsilon_{k}^{2}\nu$ as shown in Theorem 4.2,
where $\sigma^{2}$ and $\nu$ are constant across clients. This shows that
clients with large $\upsilon_{k}$ can improve its generalization performance
by having a smaller regularization weight. Intuitively, clients that have
larger $\upsilon_{k}$ have a higher chance to have larger discrepancy in data
distribution from other clients, and therefore having a smaller $\lambda_{k}$
can prevent from assimilating irrelevant knowledge from the other clients.
Similarly, the opposite also holds where clients with smaller $\upsilon_{k}$
have a higher optimal $\lambda_{k}^{*}$. This result gives insight into how to
set the regularization weight dependent on the client’s data discrepancy to
other clients. Although in our experiments we use identical $\lambda_{k}$ for
$k\in[K]$, interesting future directions include varying $\lambda_{k}$ of
across clients dependent on their data and training progress.

##  Appendix D Generalization Bound for Ensemble Models in Personalization

As defined in Section 3, we have the true data distribution of client $k$
defined as $\mathcal{D}_{k}$, and the empirical data distribution associated
with the client‘s training dataset defined as $\widehat{\mathcal{D}}_{k}$. For
a multi-class classification problem with a finite set of classes, we have
that the data’s domain is defined by the input space
$\mathbf{x}\in\mathcal{X}$ and the output space $y\in\mathcal{Y}$. For the
generalization bound analysis we consider hypotheses that maps
$h:\mathcal{X}\rightarrow\mathcal{Y}$, and $\mathcal{H}$ is defined as the
hypotheses space such that $h\in\mathcal{H}$. The loss function
$l(h(\mathbf{x}),y)$ measures the classification performance of $h$ for a
single data point $(\mathbf{x},y)$ and we define the expected loss over all
data points that follow distribution $\mathcal{D}$ as
$\mathcal{L}_{\mathcal{D}}(h)=\mathbb{E}_{(\mathbf{x},y)\sim\mathcal{D}}[l(h(\mathbf{x}),y)]$.
We assume that $\mathcal{L}(\cdot)$ is convex, and is in the range $[0,1]$. We
define the minimizer of the expected loss over the data that follows the
distribution $\mathcal{D}_{k}$ and $\widehat{\mathcal{D}}_{k}$ as each
$h_{k}=\operatorname*{arg\,min}_{h}\mathcal{L}_{\mathcal{D}_{k}}(h)$ and
$\widehat{h}_{k}=\operatorname*{arg\,min}\mathcal{L}_{\widehat{\mathcal{D}}_{k}}(h)$.
Note that for sufficiently large training dataset, we will have
$h_{k}\simeq\widehat{h}_{k}$.

Our goal is to show the generalization bound for client $k$ such that
$\mathcal{L}_{\mathcal{D}_{k}}\left(\sum_{i=1}^{K}\alpha_{k,i}h_{\widehat{\mathcal{D}}_{i}}\right)$,
where $h_{\widehat{\mathcal{D}}_{i}}$ represents the hypothesis trained from
client $i$’s training dataset and $\alpha_{k,i}$ represents the weight for the
hypothesis of client $i$ for client $k$. For client $i\in[K]$,
$h_{\widehat{\mathcal{D}}_{i}}$ will be the optimal hypothesis with respect to
the training dataset for each client participating in FL, and the
generalization bound for
$\mathcal{L}_{\mathcal{D}_{k}}\left(\sum_{i=1}^{K}\alpha_{k,i}h_{\widehat{\mathcal{D}}_{i}}\right)$
will show how the weighted average of different hypothesis from the other
clients with respect to $\alpha_{k,i},~{}i\in[K]$ helps the generalization of
an individual client $k$ with respect to its true data distribution. Before
presenting the generalization bound, we present several useful lemmas.

######  Lemma D.1 (Domain adaptation [55]).

With two true distributions $\mathcal{D}_{A}$ and $\mathcal{D}_{B}$, for
$\forall~{}\delta\in(0,1)$ and hypothesis $\forall h\in\mathcal{H}$, with
probability at least $1-\delta$ over the choice of samples, there exists:

$\displaystyle\mathcal{L}_{\mathcal{D}_{A}}(h)\leq\mathcal{L}_{\mathcal{D}_{B}}(h)+\frac{1}{2}d(\mathcal{D}_{A},\mathcal{D}_{B})+\nu$
(52)

where $d(\mathcal{D}_{A},\mathcal{D}_{B})$ measures the distribution
discrepancy between two distributions [55] and
$\nu=\inf_{h}\mathcal{L}_{\mathcal{D}_{A}}(h)+\mathcal{L}_{\mathcal{D}_{B}}(h)$.

######  Lemma D.2 (Generalization with limited training samples).

For $\forall~{}k\in[K]$, with probability at least $1-\delta$ over the choice
of samples, there exists:

$\displaystyle\mathcal{L}_{\mathcal{D}_{k}}(h_{\widehat{\mathcal{D}}_{k}})\leq\mathcal{L}_{\widehat{\mathcal{D}}_{k}}(h_{\widehat{\mathcal{D}}_{k}})+\sqrt{\frac{\log{2/\delta}}{2m_{k}}}$
(53)

where $m_{k}$ is the number of training samples of client $k$. This lemma
shows that for small number of training samples, i.e., small $m_{k}$, the
generalization error increases due to the discrepancy between
$\mathcal{D}_{k}$ and $\widehat{\mathcal{D}}_{k}$.

###### Proof.

We seek to bound the gap between
$\mathcal{L}_{\mathcal{D}_{k}}(h_{\hat{\mathcal{D}}_{k}})$ and
$\mathcal{L}_{\hat{\mathcal{D}}_{k}}(h_{\hat{\mathcal{D}}_{k}})$. Observe that
$\mathcal{L}_{\mathcal{D}_{k}}(h_{\hat{\mathcal{D}}_{k}})=\mathbb{E}\left[\mathcal{L}_{\hat{\mathcal{D}}_{k}}(h_{\hat{\mathcal{D}}_{k}})\right]$,
where the expectation is taken over the randomness in the sample draw that
generates $\hat{\mathcal{D}}_{k}$, and that
$\mathcal{L}_{\hat{\mathcal{D}}_{k}}(h_{\hat{\mathcal{D}}_{k}})$ is an
empirical mean over losses $l(h(x),y)$ that lie within $[0,1]$. Since we are
simply bounding the difference between a sample average of bounded random
variables and its expected value, we can directly apply Hoeffding’s inequality
to obtain

$\displaystyle\mathbb{P}\left[\mathcal{L}_{\hat{\mathcal{D}}_{k}}(h_{\hat{\mathcal{D}}_{k}})-\mathcal{L}_{\mathcal{D}_{k}}(h_{\hat{\mathcal{D}}_{k}})\geq\epsilon\right]\leq
2e^{-2m\epsilon^{2}}.$ (54)

Setting the right hand side to $\delta$ and rearranging gives the desired
bound with probability at least $1-\delta$ over the choice of samples:

$\displaystyle\mathcal{L}_{\mathcal{D}_{k}}(h_{\hat{\mathcal{D}}_{k}})$
$\displaystyle\leq\mathcal{L}_{\hat{\mathcal{D}_{k}}}(h_{\hat{\mathcal{D}}_{k}})+\sqrt{\frac{\log
2/\delta}{2m_{k}}}.$

∎

We now present the generalization bound for
$\mathcal{L}_{\mathcal{D}_{k}}\left(\sum_{i=1}^{K}\alpha_{k,i}h_{\widehat{\mathcal{D}}_{i}}\right)$
as follows:

$\displaystyle\mathcal{L}_{\mathcal{D}_{k}}\left(\sum_{i=1}^{K}\alpha_{k,i}h_{\widehat{\mathcal{D}}_{i}}\right)\operatorname*{\leq}_{(c)}\sum_{i=1}^{K}\alpha_{k,i}\mathcal{L}_{\mathcal{D}_{k}}(h_{\widehat{\mathcal{D}}_{i}})\operatorname*{\leq}_{(d)}\sum_{i=1}^{K}\alpha_{k,i}[\mathcal{L}_{\mathcal{D}_{i}}(h_{\widehat{\mathcal{D}}_{i}})+\frac{1}{2}d(\mathcal{D}_{i},\mathcal{D}_{k})+\nu_{i}]$
(55)

where
$\nu_{i}=\inf_{h}\mathcal{L}_{\mathcal{D}_{i}}(h)+\mathcal{L}_{\mathcal{D}_{k}}(h)$,
(c) is due to the convexity of $\mathcal{L}$, and (d) is due to lemma D.1. We
can further bound 55 using lemma D.2 as

$\displaystyle\mathcal{L}_{\mathcal{D}_{k}}\left(\sum_{i=1}^{K}\alpha_{k,i}h_{\widehat{\mathcal{D}}_{i}}\right)\leq\sum_{i=1}^{K}\alpha_{k,i}\mathcal{L}_{\widehat{\mathcal{D}}_{i}}(h_{\widehat{\mathcal{D}}_{i}})+\sum_{i=1}^{K}\alpha_{k,i}\sqrt{\frac{\log{2/\delta}}{2m_{k}}}+\frac{1}{2}\sum_{i=1}^{K}\alpha_{k,i}d(\mathcal{D}_{i},\mathcal{D}_{k})+\sum_{i=1}^{K}\alpha_{k,i}\nu_{i}$
(56)
$\displaystyle=\sum_{i=1}^{K}\alpha_{k,i}\mathcal{L}_{\widehat{\mathcal{D}}_{i}}(h_{\widehat{\mathcal{D}}_{i}})+\sqrt{\log{\delta^{-1}}}\sum_{i=1}^{K}\frac{\alpha_{k,i}}{\sqrt{m_{k}}}+\frac{1}{2}\sum_{i=1}^{K}\alpha_{k,i}d(\mathcal{D}_{i},\mathcal{D}_{k})+\sum_{i=1}^{K}\alpha_{k,i}\nu_{i}$
(57)

From 57, with
$\mathcal{L}_{\widehat{\mathcal{D}}_{i}}(h_{\widehat{\mathcal{D}}_{i}})$ in
general being small for $\forall i\in[K]$ as it is the minimum loss, and
$m_{i}$ being similar to other $m_{i^{\prime}},~{}i^{\prime}\in[K]$, the only
way to minimize the generalization error of
$\mathcal{L}_{\mathcal{D}_{k}}\left(\sum_{i=1}^{K}\alpha_{k,i}h_{\widehat{\mathcal{D}}_{i}}\right)$
is to set the weights $\alpha_{k,i}$ so that the third term
$\frac{1}{2}\sum_{i=1}^{K}\alpha_{k,i}d(\mathcal{D}_{i},\mathcal{D}_{k})$ is
minimized. Note that it is difficult to know the value of $\nu_{i}$, making it
impractical to minimize the fourth term in practice. This generalization
results strengthens our motivation to use to find the weights
$\alpha_{k,i},~{}i\in[K]$ that minimizes
$\frac{1}{|\mathcal{P}|}\sum_{\mathbf{x}\in\mathcal{P}}\|\sum_{i=1}^{K}\alpha_{k,i}s_{i}(\mathbf{w}_{i},\mathbf{x})-s(\mathbf{w}_{k},\mathbf{x})\|_{2}^{2}$
in regards to the objective function we have in 1.

##  Appendix E Details of Experimental Setup

Codes for the results in the paper are presented in the supplementary
material.

### Description for Toy Example - Figure 1

For Figure 1, we design a linear regression problem where the true local model
for each client is generated as $\mathbf{w}_{i}=\theta+\zeta_{i},i\in[3]$
where $\theta\in\mathbb{R}^{2\times 1}$ is a non-informative prior which
elements are uniformly distributed $\mathcal{U}(-10,10)$ and the elements of
$\zeta_{i}$ follows the normal distribution
$\mathcal{N}(0,\sigma_{i}),~{}\sigma_{1}=2,\sigma_{2}=5,\sigma_{3}=200$. The
discrepancy across the variance denotes the data-heterogeneity across the
clients. The range for $\mathbf{x}$ is $[-10,10]$ for all elements. We assume
all clients have identical dataset size. For implementing PerFed-CKT for the
toy example, we set the public data range as identical to the input data
range, and set $\lambda=50$. For KT w/o clustering, the co-distillation term
uses a simple average of all the logits from the clients for regularizing
while for KT with clustering the weights are set so that clients with similar
true local models have higher weights for each other. This setting is also
consistent with the generalization analysis presented in Section 4. The code
for the linear regression toy example is presented in the supplementary
material.

### Description for CIFAR10 Experiments.

#### Data Partitioning.

We experiment with three different seeds for the randomness in the dataset
partition across clients and present the averaged results across the seeds
with the standard deviation. The partitioning of each individual client’s data
to training/validation/test dataset is done as follows: after partitioning the
entire dataset by the Dirichlet distribution $\text{Dir}_{K}(\alpha)$ with
$\alpha=0.01$ across clients, we partition each client dataset by a
$\\{0.1,0.3,0.4\\}/0.1/0.5$ ratio where the ratio for the training dataset is
chosen by random from $\\{0.1,0.3,0.4\\}$ for each client. Such partitioning
simulates a more realistic FL setting where individual clients may not have
sufficient labeled data samples for training that represents the test
dataset’s distribution. For all experiments we run 500 communication rounds
which we have observed allows convergence for all experiments.

#### Local Training and Hyperparameters.

For the local-training hyperparameters, we do a grid search over the learning
rate: $\eta\in\\{0.1,0.05,0.01,0.005,0.001\\}$, batchsize:
$b\in\\{32,64,128\\}$, and local iterations: $\tau\in\\{10,30,50\\}$ to find
the hyper-parameters with the highest test accuracy for each benchmark. For
all benchmarks we use the best hyper-parameter for each benchmark after doing
a grid search over feasible parameters referring to their source codes that
are open-sourced. For the knowledge distillation server-side hyperparameters,
we do a grid search over the public batch size:
$b^{\prime}\in\\{32,64,128\\}$, regularization weight
$\lambda\in\\{0.05,0.1,0.5,1,2,4\\}$ to find the best working hyperparameters.
The best hyperparameters for PerFed-CKT we use is
$\eta=0.001,b=64,\tau=50,b^{\prime}=128,\lambda=2$.

#### Model Setup.

For the model configuration, for the CNN we have a self-defined convolutional
neural network with 2 convolutional layers with max pooling and 4 hidden fully
connected linear layers of units $[120,100,84,50]$. The input is the flattened
convolution output and the output is consisted of 10 units each of one of the
0-9 labels. For the VGG, we use the open-sourced VGG net from Pytorch with
torchvision ver.0.4.1 presented in Pytorch without pretrained as False and
batchnorm as True.

#### Platform.

All experiments are conducted with clusters equipped with one NVIDIA TitanX
GPU. The number of clusters we use vary by $C$, the fraction of clients we
select. The machines communicate amongst each other through Ethernet to
transfer the model parameters and information necessary for client selection.
Each machine is regarded as one client in the federated learning setting. The
algorithms are implemented by PyTorch.