# Scalable Collaborative Learning via Representation Sharing

Frédéric Berdoz
MIT, EPFL
<EMAIL_ADDRESS>Abhishek Singh
MIT
<EMAIL_ADDRESS>Martin Jaggi
EPFL
<EMAIL_ADDRESS>Ramesh Raskar
MIT
<EMAIL_ADDRESS>
###### Abstract

Privacy-preserving machine learning has become a key conundrum for multi-party
artificial intelligence. Federated learning (FL) and Split Learning (SL) are
two frameworks that enable collaborative learning while keeping the data
private (on device). In FL, each data holder trains a model locally and
releases it to a central server for aggregation. In SL, the clients must
release individual cut-layer activations (smashed data) to the server and wait
for its response (during both inference and back propagation). While relevant
in several settings, both of these schemes have a high communication cost,
rely on server-level computation algorithms and do not allow for tunable
levels of collaboration. In this work, we present a novel approach for
privacy-preserving machine learning, where the clients collaborate via online
knowledge distillation using a contrastive loss (contrastive w.r.t. the
labels). The goal is to ensure that the participants learn similar features on
similar classes without sharing their input data. To do so, each client
releases averaged last hidden layer activations of similar labels to a central
server that only acts as a relay (i.e., is not involved in the training or
aggregation of the models). Then, the clients download these last layer
activations (feature representations) of the ensemble of users and distill
their knowledge in their personal model using a contrastive objective. For
cross-device applications (i.e., small local datasets and limited
computational capacity), this approach increases the utility of the models
compared to independent learning and other federated knowledge distillation
(FD) schemes, is communication efficient and is scalable with the number of
clients. We prove theoretically that our framework is well-posed, and we
benchmark its performance against standard FD and FL on various datasets using
different model architectures.

##  1 Introduction

Motivated by concerns such as data privacy, large scale training and others,
Machine Learning (ML) research has seen a rise in different types of
collaborative ML techniques. Collaborative ML is typically characterized by an
orchestrator algorithm that enables training ML model(s) over data from
multiple owners without requiring them to share their sensitive data with
untrusted parties. Some of the well known algorithms include Federated
Learning (FL) [35], Split Learning (SL) [16] and Swarm Learning [51]. While
the majority of the works in collaborative ML rely upon a centralized
coordinator, in this work, we design a decentralized learning framework where
the server plays a secondary role and could easily be replaced by a peer-to-
peer network. As we show in the rest of the paper, the main benefit of our
decentralized approach over centrally coordinated ML is the increased
flexibility among clients in controlling the information flow across different
aspects such as communication, privacy, computational capacity, data
heterogeneity, etc.

Most existing collaborative learning schemes are built upon FL, where the
training algorithm and/or the model architecture is usually imposed to the
clients to match the computational capacity of the weaker participant (or the
server). We refer to this property as _non-tunable_ collaboration, since the
degree of collaboration is mostly imposed by the server. While this is less of
a problem in cross-silo applications (small number of data owners, large local
datasets) since the participants can easily find consensus over the best
training parameters, it can constitute a strong barrier for participation in
cross-device applications (large number of users with small local datasets and
low/heterogeneous computational capacity). Indeed, finding consensus becomes
harder in that scenario. To alleviate this, we propose a new framework for
_tunable_ collaboration, i.e., where each participant has near total control
over its data release, its model architecture and how the knowledge of other
users should be integrated in its own personalized model.

Our main idea is to share learned feature representations of each class among
users and to use these representations cleverly during local training. Since
the clients can choose which features are aggregated and shared, our framework
enables the clients to assign different privacy levels to different samples.
Our decentralized approach also ensures that the overall system remains
asynchronous and functions as expected even if all but two clients are offline
in the whole system. Finally, our framework makes it convenient to account for
model heterogeneity and model personalization, since every user can select a
subset of peers based on their goals of generalization and personalization.
While some of these advantages have been introduced in recent FL based
schemes, our framework allows natural integration of such several ideas. Our
work is different from the body of work done in FL due to its decentralized
design. Specifically, our system does not use server-level computation
algorithms that are directly involved in the training of the model.
Nevertheless, for completeness, we experimentally compare it with FL in
Section 3.

Our contribution can be summarized as follows:

  * •

We present a new _tunable_ collaborative learning algorithm based on
contrastive representation learning and online knowledge distillation.

  * •

We prove theoretically that our local objective is well-posed.

  * •

We show empirically the advantages of our framework against other similar
schemes in terms of utility, communication, and scalability.

##  2 Related Work

#### Federated Learning

FL is considered to be the first formal framework for collaborative learning.
In their initial paper, McMahan et al. [35] introduce a new algorithm called
FedAvg, in which each client performs several optimization steps on their
local private dataset before sending the updated model back to the server for
aggregation using weight averaging. While this approach alleviates the
communication cost of the baseline collaborative optimization algorithm
FedSGD, it also decreases the personalization capacity of the global model due
to the naive model averaging, especially in heterogeneous environments.
Several algorithms have been proposed to address these limitations, in
particular FedProx [29], FedPer [4], FedMa [49], FedDist [40], FedNova [50],
Scaffold [22] and VRL-SGD [31]. While improving the convergence rate, these
algorithms suffer from the same constraints as FedAvg, i.e., homogeneous model
architecture for every client, high communication overhead and _non-tunable_
collaboration, all potential barriers for participation.

#### Fully Decentralized Learning

The use of a central server in traditional FL constitutes a single point of
failure and can also become a bottleneck when the number of clients grows, as
shown by Lian et al. [30]. To alleviate these issues, Vanhaesebrouck et al.
[47] formalize a new framework where each client participates in the learning
task via a peer-to-peer network using gossip algorithms [41, 13]. In this
configuration, there is no global solution and each client has its own
personalized model, which enables both personalization and generalization. On
the other hand, it creates other challenges about convergence, practical
implementation and privacy [21]. Moreover, as in FL, the entire model must be
released at every communication round, which can constitute a barrier for
participation for the same reasons.

#### Knowledge Distillation

(a) Offline KD

(b) Online KD

(c) Ensemble offline KD

Figure 1: Different types of response-based (green) and feature-based (blue)
knowledge distillation.

The concept of knowledge distillation (KD) originated in Bucila et al. [8] as
a way of compressing models, and was later generalized by Hinton et al. [18]
(see Gou et al. [14] for an overview of the field). The standard use case for
KD is that of a Teacher-Student (or offline) configuration, in which the
teacher model (usually a large and well-trained model) transfers its knowledge
to the student model (usually a smaller model) by sharing its last layer
activations on a given _transfer_ dataset (see Fig. 1(a)). The knowledge is
then _distilled_ into the student model using a divergence loss between the
teacher and student models outputs (response-based KD) or intermediate layers
(feature-based KD) on the transfer dataset. Traditional KD schemes use a
transfer set that is similar (or identical) to the teacher training dataset,
but some recent work has focused on _data-free_ (or _zero-shot_) KD. This can
be achieved either by looking at some of the teacher model statistics to
generate synthetic transfer data [34, 37, 6], or by training a GAN in parallel
[36, 10, 2]. It has also been shown that positive results can be obtained
using mismatched or random unlabeled data for the distillation [26, 38]. A key
concept for our framework is the idea of online KD (or co-distillation [3],
see Fig. 1(b)), where each model is treated as both a teacher and a student,
meaning that the KD is performed synchronously with the training of each model
(rather than after the training of the teacher model) [15, 54, 42]. Finally,
ensemble KD (see Fig. 1(c)) refers to the setup where the knowledge is
distilled from an ensemble of teacher (offline) or teacher-student (online)
models.

#### Knowledge Distillation for Collaborative Learning

A growing body of literature has recently investigated ways of using online KD
for collaborative learning in order to alleviate the need of sharing model
updates (FL) or individual smashed data (SL). Jeong et al. [20] present
Federated Knowledge Distillation (FD), where each client uploads its mean (per
class) logits to a central server, who aggregates and broadcasts them back.
These soft labels are then used as the teacher output for the KD loss during
local training. A closely related idea is to compute the mean logits on a
common public dataset [28, 19, 9], but we argue that selecting this dataset
can induce bias and is not always feasible, since additional trust is needed
for its selection, and sufficient relevant data might be lacking. Besides FD,
KD can enable collaborative learning in various ways: In an attempt to
decrease the communication cost, Wu et al. [52] introduce FedKD, in which each
client trains a (large) teacher network on their private data and transfer
locally its knowledge to a smaller student model, which is then used in a
standard FL algorithm. On the other hand, Lin et al. [32] and Chen and Chao
[11] use KD on an unlabeled synthetic or public dataset to make the FL
aggregation algorithm more robust. Our approach differs significantly from
these schemes, as it does not rely on traditional FL algorithms.

##  3 Methods

Figure 2: The proposed framework: clients exchange averaged
($\mathcal{A}_{s}$) feature representations for each class and use these
representations in their own local training, but keep their raw data private
(on device). The ensemble of participants (network) constitutes the teacher,
and each single participant is a student. The server acts mainly as a relay as
it does not take part in the training or aggregation of the models. At each
communication round, the student downloads a subset of representations and
uploads some of its own.

#### Preamble and Motivation

Consider any classification task on $d$-dimensional raw inputs with $C$
distinct classes and a set of $N$ participating users $\\{u\\}_{u=1}^{N}$,
each with a local private dataset $\mathcal{D}_{u}$ and a model $f_{u}$:

$\mathcal{D}_{u}:=\\{(\mbf{x}_{i},y_{i})\overset{iid}{\sim}p_{u}\\}_{i=1}^{n_{i}},\qquad
p_{u}(\mbf{x},y):=p_{\mbf{X},Y|U=u}(\mbf{x},y),\qquad
f_{u}=\tau_{u}\circ\phi_{u},$

where $p_{\mbf{X},Y,U}$ represents the joint probability of choosing a user
$U$ and drawing a sample $(\mbf{X},Y)$ from its distribution, and $\tau_{u}$,
$\phi_{u}$ represent the linear classifier and neural network (up to last
hidden layer) of user $u$, respectively (with potentially different
architectures across users). More precisely, let $d^{\prime}$ be the output
dimension of $\phi_{u}$ and
$\mbf{w}_{u}=\\{{\bm{\theta}}_{u},\mbf{W}_{u},\mbf{b}_{u}\\}$ be model weights
of user $u$. We have:

$\displaystyle\phi_{u}:\mathbb{R}^{d}$
$\displaystyle\to\mathbb{R}^{d^{\prime}}$
$\displaystyle\tau_{u}:\mathbb{R}^{d^{\prime}}$
$\displaystyle\to\mathbb{R}^{C}$ $\displaystyle\mbf{x}$
$\displaystyle\mapsto\mbf{x}^{\prime}:=\phi_{u}(\mbf{x};{\bm{\theta}}_{u})$
$\displaystyle\mbf{x}^{\prime}$
$\displaystyle\mapsto\mbf{z}:=\tau_{u}(\mbf{x}^{\prime};\mbf{W}_{u},\mbf{b}_{u})=\mbf{W}_{u}\mbf{x}^{\prime}+\mbf{b}_{u}$

where $\mbf{x},\mbf{x}^{\prime}$ and $\mbf{z}$ are the raw input, the feature
representation and the logits, respectively. We motivate our approach as
follows. Assuming no sharing of raw data $\mbf{x}$ (for privacy concerns), any
collaborative learning framework falls in one of two buckets: weight sharing
or activation sharing. As mentioned in Section 2, sharing weights (i.g., FL)
comes with strong constraints (communication, model architecture, etc.) and
might not always be suitable. Concerning activation sharing (i.g., SL), it can
be done at any layer (hidden or output). However, since activations are
usually strongly correlated to the raw input [48], sharing single sample
activations can raise privacy concerns. Instead, sharing averaged activations
can easily be met with differential privacy guarantees (see privacy in Section
3). Due to the high non-linearity of neural networks, sharing averaged
activations only makes sense at the output layer or last hidden layer (since
the classifier $\tau_{u}$ is linear). Sharing only averaged outputs (averaged
over samples of the same class) like in FD has been shown to have limited
success as the quantity of shared information is restricted to the dimension
of the output $C$, and we argue in this paper that sharing the feature
representations (i.e., outputs of the last hidden layer, also averaged over
samples of the same class) is more flexible and leads to better results. In
this light, our objective is to collaboratively learn the best feature
representation for each class, using contrastive (w.r.t. classes)
representation learning [43] and feature-based knowledge distillation [14]. In
other words, we want to learn collaboratively (i.e., only once) the structure
of the feature representation space so that each client does not need to find
it on its own with its limited amount of data and/or computational capacity.

#### Contrastive Objective

In this section, we introduce a contrastive objective function for private
online knowledge distillation (i.e., when users synchronously learn
personalized models without sharing their raw input data and by collaborating
via online distillation). This new objective function and its derivation are
partly inspired by the offline contrastive loss presented in Tian et al. [43]
and van den Oord et al. [45], with a few important differences. In the offline
non-private setting, both the teacher and student models $\phi_{t}$ and
$\phi_{s}$ have access to the same dataset $\\{x_{i}\\}_{i=1}^{n}$. In that
scenario, the representations $\phi_{s}(x_{i})$ and $\phi_{t}(x_{j})$ are
_pulled apart_ if $i\neq j$ and are _pulled together_ if $i=j$. However, in
the private setting, $\phi_{s}$ does not have access to the raw input data
that was used to train the teacher model, but has its own private dataset. At
a given communication round and from the perspective of user $u$, we define
$\phi_{u}$ as the student model and $\phi_{U}$ as the teacher model, where
$U\sim p_{U}$ is a user selected at random. Consider the following procedure:
$u$ samples $Y\sim p_{Y|U=u}$ from its own data distribution and then samples
either jointly (i.e., from the same observation of $Y$) or independently (from
two independent observations of $Y$) the two random vectors ${\bm{\Phi}}_{s}$
and ${\bm{\Phi}}_{t}$ defined as follows:

$\displaystyle\mbf{X}$ $\displaystyle\sim p_{\mbf{X}|Y,U=u}$
$\displaystyle{\bm{\Phi}}_{s}$ $\displaystyle:=\phi_{u}(\mbf{X}),$ (1)
$\displaystyle U\sim p_{U},\quad(\mbf{X}_{1},...,\mbf{X}_{n_{avg}})$
$\displaystyle\overset{iid}{\sim}p_{\mbf{X}|Y,U},$
$\displaystyle{\bm{\Phi}}_{t}$
$\displaystyle:=\frac{1}{n_{avg}}\sum_{i=1}^{n_{avg}}\phi_{U}(\mbf{X}_{i}).$
(2)

The parameter $n_{avg}$ defines over how many samples we take the average,
which in turn defines the concentration of the distribution of
${\bm{\Phi}}_{t}$. From the (student) perspective of client $u$ and in the
spirit of collaboration, the goal is to maximize the mutual information
$\mathcal{I}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})$. Still from the perspective of
$u$, let $p_{s,t}$, $p_{s}$ and $p_{t}$ be the joint and marginal
distributions of ${\bm{\Phi}}_{s}$ and ${\bm{\Phi}}_{t}$, respectively, and
let $I$ be a Bernoulli random variable indicating if a tuple
$({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})$ has been drawn from the joint distribution
$p_{s,t}$ or from the product of marginals $p_{s}p_{t}$. Finally, let
$q(\mbf{s},\mbf{t},i)$ be the joint distribution of
$({\bm{\Phi}}_{s},{\bm{\Phi}}_{t},I)$ such that
$q(\mbf{s},\mbf{t}|i=1)=p_{s,t}(\mbf{s},\mbf{t})$ and
$q(\mbf{s},\mbf{t}|i=0)=p_{s}(\mbf{s})p_{t}(\mbf{t})$ and suppose that the
prior $q(i)$ satisfy $q(i=1)=\frac{1}{K+1}$ and $q(i=0)=\frac{K}{K+1}$, i.e.,
for each sample from the distribution $p_{s,t}$, we draw $K$ samples from the
distribution $p_{s}p_{t}$. We can show the following bound.

######  Theorem 1.

Using the above notation, let $h(i,\mbf{s},\mbf{t})$ be any estimate of
$q(i|\mbf{s},\mbf{t})$ with Bernoulli parameter $\hat{h}(\mbf{s},\mbf{t})$,
and let $\mathcal{L}_{disc}(h,\phi_{u})$ be defined as follows:

$\mathcal{L}_{disc}(h,\phi_{u}):=-\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
p_{s,t}}\left[\log\hat{h}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]-K\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
p_{s}p_{t}}\left[\log(1-\hat{h}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t}))\right].$ (3)

The mutual information $\mathcal{I}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})$ can be
bounded as

$\mathcal{I}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\geq\log(K)-\mathcal{L}_{disc}(h,\phi_{u}),$
(4)

with equality iff $h=q$ (better estimates lead to tighter bounds). The proof
is joined in the supplementary material.

Hence, by minimizing $\mathcal{L}_{disc}$ in Eq. (4), we optimize a lower
bound on the mutual information
$\mathcal{I}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})$. Taking advantage of the
classifier $\tau_{u}$, a natural choice for $h$ is the discriminator $h_{u}$
with Bernoulli parameter

$\hat{h}_{u}(\mbf{s},\mbf{t};\mbf{W}_{u},\mbf{b}_{u})=\langle\mbox{softmax}(\tau_{u}(\mbf{s})),\mbox{softmax}(\tau_{u}(\mbf{t}))\rangle.$
(5)

With this choice, $\hat{h}_{u}(\mbf{s},\mbf{t})$ can be interpreted as the
estimated probability that the features $\mbf{s}$ and $\mbf{t}$ come from the
same class. Note that in their work, Tian et al. [43] train an external
discriminator (i.e., that is not defined using the model classifier
$\tau_{u}$).

#### Final Objective

Intuitively, from the perspective of $u$, minimizing $\mathcal{L}_{disc}$
ensures that its classifier $\tau_{u}$ can distinguish if two feature
representations, one local and one from another user, come from the same
class. To improve the algorithm convergence and to ensure that $\tau_{u}$ can
classify the feature representation ${\bm{\Phi}}_{t}$ of another client
(similar as in Invariant Risk Minimization [5]), we also introduce a classical
feature-based KD term $\mathcal{L}_{KD}$. This term minimizes the $L_{2}$
distance between the local and global feature representations of a same class
(we define the global representation of class $c$ as the expected value
$\mathbb{E}_{(X,U)\sim p_{X,U|Y=c}}[\phi_{U}(\mbf{X})]$). An important
distinction between $\mathcal{L}_{KD}$ and $\mathcal{L}_{disc}$ is that the
first one uses an _inter_ -client averaged representation, whereas the second
one uses an _intra_ -client averaged representation. Combining
$\mathcal{L}_{KD}$ and $\mathcal{L}_{disc}$ with the standard cross-entropy
loss $\mathcal{L}_{CE}(\tau_{u},\phi_{u})$ using the meta parameters
$\lambda_{KD}$ and $\lambda_{disc}$, the final optimization problem of $u$
becomes:

$\mbox{Find}\quad{\bm{\theta}}_{u}^{\star},\mbf{W}^{\star}_{u},\mbf{b}^{\star}_{u}=\underset{{\bm{\theta}}_{u},\mbf{W}_{u},\mbf{b}_{u}}{\mbox{argmin}}\,\mathcal{L}_{CE}(\tau_{u},\phi_{u})+\lambda_{KD}\mathcal{L}_{KD}(\phi_{u})+\lambda_{disc}\mathcal{L}_{disc}(\hat{h}_{u},\phi_{u}).$
(6)

#### Implementation

With the standard assumption that the datasets are composed of IID samples
drawn from their corresponding distributions, the expected losses
$\mathcal{L}_{CE},\mathcal{L}_{KD}$ and $\mathcal{L}_{disc}$ can be
approximated by their unbiased mini-batch estimators $L_{CE},L_{KD}$ and
$L_{disc}$, respectively. For one given sample, the loss functions are given
by $\ell_{CE}(\mbf{z},c):=-\log\left(\mbox{softmax}(\mbf{z})\right)_{c}$,
$\ell_{KD}(\mbf{x}^{\prime},\mbf{x}^{\prime\prime}):=\|\mbf{x}^{\prime}-\mbf{x}^{\prime\prime}\|^{2}$
and finally

$\displaystyle\ell_{disc}:\\{0,1\\}\times\mathbb{R}^{d^{\prime}}\times\mathbb{R}^{d^{\prime}}$
$\displaystyle\to\mathbb{R}_{+}$ $\displaystyle(i,\mbf{s},\mbf{t})$
$\displaystyle\mapsto-i\log\left(\hat{h}(\mbf{s},\mbf{t})\right)-(1-i)\log\left(1-\hat{h}(\mbf{s},\mbf{t})\right).$
(7)

Finally, for the sampling procedure of the contrastive objective, we use the
most intuitive scheme: for every training sample $(\mbf{x}_{i},y_{i})$ with
feature representation $\mbf{s}_{i}=\phi_{u}(\mbf{x}_{i})$, we use one
observation of ${\bm{\Phi}}_{t}$ sampled using $y_{i}$ (i.e., $I=1$), and one
observation of ${\bm{\Phi}}_{t}$ sampled using each $c\neq y_{i}$ (i.e.,
$I=0$). Thus, we have $K=C-1$. With these considerations, a detailed
description of our collaborative learning framework can be found in the
supplementary material.

#### Communication

In terms of communication, the uplink and downlink volumes are of order
${\mathcal{O}((M_{\uparrow}+1)Cd^{\prime})}$ and
$\mathcal{O}(N(M_{\downarrow}+1)Cd^{\prime})$ per round, where $M_{\uparrow}$
and $M_{\downarrow}$ represent the number of ${\bm{\Phi}}_{t}$ realizations
(per class) that are uploaded and downloaded by the clients, respectively
(these parameters can be tuned to match the communication capacity of the
network). For comparison, these volumes are of order $\mathcal{O}(D)$ and
$\mathcal{O}(ND)$ for FL (with $D$ the model size) and
$\mathcal{O}(nd^{\prime})$ and $\mathcal{O}(Nnd^{\prime})$ for SL. Because in
most scenarios $D\gg n\gg d^{\prime}\gg C$ and since $M_{\uparrow}$ and
$M_{\downarrow}$ are tunable, we observe that our framework is communication
efficient compared to FL and SL.

#### Relaxation to peer-to-peer

We emphasize here that our contrastive objective $\mathcal{L}_{disc}$ is fully
compatible with a peer-to-peer configuration, since the server only acts as a
relay for the observations of ${\bm{\Phi}}_{t}$. For the traditional feature-
based KD objective $\mathcal{L}_{KD}$, we use one global representation per
class (i.e., one for the whole network), which in theory can only be computed
by a central entity. One way to alleviate this could be to use the average of
all the observations of ${\bm{\Phi}}_{t}$ that were downloaded by user $u$ as
a proxy for the global feature representations. However, to focus solely on
the effectiveness of the proposed objectives rather than the topology of the
network, we only present experiments in which a central entity computes the
global representations.

#### Privacy

Our framework naturally attains a certain degree of privacy by sharing
representations from the penultimate layer, which, as per information
bottleneck principle [44], has minimal mutual information from the original
data. Although such a level of privacy may not be enough from a membership
attack point of view, we emphasize that the privacy is amplified by averaging
out these representations, hence amplifying the privacy of individual samples.
To further make our framework formally private, we propose using recent
progress in differentially private mean estimation [33, 7] to obtain a
differentially private mean before it gets shared with other clients.
Furthermore, our framework naturally dovetails with recent works [24, 1] in
variable privacy levels for different samples. The users can assign variable
levels of privacy to different samples before they get aggregated in a privacy
preserving manner.

##  4 Experiments

Figure 3: Ablation study for $\lambda_{KD}$ and $\lambda_{disc}$. Average test
accuracy improvement [$\%$] w.r.t. IL (upper left corner) when different
combinations of $\lambda_{KD}$ and $\lambda_{disc}$ are used (MNIST/LeNet5
experiment with 5 users and 100 communication rounds). The red square
indicates the value of $\lambda_{KD}$ and $\lambda_{disc}$ used in all our
experiments.

#### Datasets, models and training

We run several experiments with the MNIST [12], Fashion-MNIST [53] and CIFAR10
[25] datasets. For MNIST, we use a simple convolutional neural network (CNN)
similar to LeNet5 ($\approx 30K$ parameters) [27] and for Fashion-MNIST and
CIFAR10, we use ResNet9 ($\approx 2,4M$ parameters) and ResNet18 ($\approx
11.3M$ parameters) architectures [17], respectively. For the dimension of
feature representation space, we set $d^{\prime}=84$ for LeNet5,
$d^{\prime}=128$ for ResNet9 and $d^{\prime}=256$ for ResNet18. In order to
simulate a scenario where the data is sparse, we only select a fraction of the
train dataset (1200 samples for MNIST, 6000 for Fashion-MNIST and 10000 for
CIFAR10) that we split uniformly at random across $N\in\\{2,5,10\\}$ users.
For the validation, we use the entire test dataset for each task (10000
samples). In order to have fair comparisons, we train all the models for the
same number of communication rounds, and we stop the training as soon as
framework has reasonably converged ($r=100$ for MNIST/LeNet5, $r=20$ for
Fashion-MNIST/ResNet9 and CIFAR10/ResNet18). We compare our framework with
centralized learning (or CL, i.e., with $N=1$ and
$\lambda_{KD}=\lambda_{disc}=0$), independent learning (or IL, i.e., with
$\lambda_{KD}=\lambda_{disc}=0$), federated learning using FedAvg (FL) and
federated knowledge distillation (FD). We use the default learning rate
$\eta=10^{-3}$ for all experiments and we perform 1 local epoch of training
per communication round. For CL, IL, FD and our framework, we use the Adam
stochastic optimization algorithm [23]. Finally, supported by Fig. 3, we set
$\lambda_{KD}=10$ and $\lambda_{disc}=1$ in our final objective (Eq. (6)).

#### Network emulation

For our contrastive objective $\mathcal{L}_{disc}$, we use $n_{avg}=10$ (i.e.,
for every class, each user selects $n_{avg}$ samples of that class, computes
and averages their feature representations, uploads them to the server and
downloads the representations of another user chosen at random). For the
feature-based KD objective $\mathcal{L}_{KD}$, each client average the feature
representation of all the samples of a same class and uploads these averaged
representations to the server, which in turn averages them to obtain one
global representation per class. For simplicity, we use
$M_{\uparrow}=M_{\downarrow}=1$ (clients upload and download one observation
of ${\bm{\Phi}}_{t}$ for each class).

#### Software and hardware

We use the CUDA enabled PyTorch library [39] with a GTX 1080 Ti (11GB) GPU and
an Intel(R) Xeon(R) CPU (2.10GHz).

##  5 Discussion

#### Utility

As seen in Table 1, our framework outperforms every other framework for when a
small model is used (MNIST/LeNet5), especially when the number of clients
grows, which is typically the kind of configuration that would be relevant for
a cross-device application). In that setup, FL performs particularly poorly as
it struggles to find a low-capacity model that matches the data distribution
of each client. This is particularly visible on Fig. 4(b). Although in that
configuration, FD shows similar performance for small number of clients
($N=2,5$), it still exhibits a lower rate of convergence (compare Fig. 4(c)
and Fig. 4(d)). Our framework even outperforms centralized learning (CL) when
$N=2$, suggesting that the added objectives can also be seen as regularizers.
For the Fashion-MNIST dataset, our framework shows significant improvement
over IL and FD (i.e., frameworks where there are no global model), but is not
able to compete with FL anymore, which in that case even outperforms CL.
However, the comparison between FL and our framework is unfair for large
models due to the amounts of shared information (i.g., for ResNet9, one
communication round of our framework communicates $\approx 1000$ times fewer
bits than standard FL. This ratio increases to $\approx 2000$ for ResNet19).
Similar conclusions can be drawn for the CIFAR10 experiments. However, in that
case, even IL outperforms our framework for $N=10$ after $r=20$ rounds.
Indeed, since our objective function is highly complex (the global minimum
depends on the data of other peers), the algorithm struggles to converge to
the optimal model when the number of parameters is very large. However, we
argue that this is not critical since training ResNet18 models on low capacity
devices is not particularly adapted to cross-device applications, as they
usually require powerful GPUs and large datasets to converge.

(a) IL

(b) FL

(c) FD

(d) Ours

Figure 4: Comparison of the train and test accuracy during 100 communication rounds of IL, FL, FD and our framework on the MNIST dataset with LeNet5 architectures and $N=5$ users. The shaded areas represent $\pm$ the standard deviation of the metric across clients. For the validation, each model is tested using the entire test dataset. Table 1: Average test accuracy over clients $[\%]$ of the different frameworks after $r$ communication rounds when the same amount of data is divided uniformly at random between $N$ users. We use 1200 training samples for MNIST, 6000 for Fashion-MNIST and 10000 for CIFAR10. The validation is done using the full test dataset (10000 samples for each task). | MNIST ($r=100$) | Fashion-MNIST ($r=20$) | CIFAR10 ($r=20$)
---|---|---|---
CL | 94.00 | 87.77 | 66.15
|  $N=2$ |  $N=5$ |  $N=10$ |  $N=2$ |  $N=5$ |  $N=10$ |  $N=2$ |  $N=5$ |  $N=10$
FL | 92.64 | 86.79 | 70.06 | 89.79 | 89.28 | 88.21 | 67.99 | 59.18 | 51.05
IL | 91.46 | 85.26 | 72.86 | 86.04 | 83.61 | 80.52 | 59.85 | 46.46 | 38.51
FD | 94.45 | 90.55 | 77.90 | 87.17 | 83.32 | 79.44 | 56.75 | 44.91 | 31.43
Ours | 94.19 | 90.63 | 82.07 | 87.91 | 84.44 | 80.77 | 63.49 | 47.28 | 37.78

#### Effect of $\mathcal{L}_{KD}$

The feature-based KD term $\mathcal{L}_{KD}$ makes sure that similar feature
representations are learned by each client. Fig. 5 provides insight into the
evolution of the latent feature space during training. In FL, the structure of
the latent feature space is imposed from the start to each client, whereas in
our framework, clients have more freedom to learn the best structure that fits
their data, since they are only _guided_ towards a global representation
structure. This clearly improves convergence, as we can see that clusters are
already formed at round 25 with our algorithm, whereas with FL the latent
feature space is still relatively chaotic at that point. On the other hand, we
observe that for IL, the learned latent feature space is totally different for
each client as they learn it by themselves, and the convergence rate is also
notably slower (compare round 25 for instance).

Figure 5: T-distributed Stochastic Neighbor Embedding [46] of the feature
representation at different rounds in a two party configuration for the
MNIST/LeNet5 experiment. Clients are distinguished by color and classes by
opacity.

#### Effect of $\mathcal{L}_{disc}$

The effect of the contrastive objective is more subtle, but can still improve
the test accuracy significantly if tuned properly. In fact, whereas
$\mathcal{L}_{KD}$ can be used without $\mathcal{L}_{disc}$, the reciprocal is
not true as observed on Fig. 3. This is because we have defined the
discriminator of user $u$ (i.e., $\hat{h}_{u}$) using its feature
representation classifier $\tau_{u}$ (see Eq. (5)). Hence, in order to
generate useful gradients for back propagation (i.e., that do not look like
noise), $\tau_{u}$ must be able to classify relatively well the feature
representations that are generated by other models. We have also tested
external discriminators (like in [43]), but the results were not as positive.

#### Limitations and Future work

As shown by the results, our approach is particularly relevant when the model
capacity is relatively small, but struggles when the model complexity grows,
especially when compared with FL. However, the comparison is unfair in that
scenario, as the amount of information that is shared between clients at each
round is larger by several orders of magnitudes for FL. We also posit that the
convergence gap could be reduced by adding sophistication into the algorithm
(i.g., by having a local algorithm that selects the most relevant observed
representations from other users, by only sharing a fraction of the feature
representation dimensions, by introducing a more complex discriminator $h$,
etc.). Another important limitation is the fact that our framework cannot be
used for supervised tasks other than classification, as the feature
representations need to be sorted by class. This is harder to alleviate since
the whole system’s architecture is build upon this requirement. Future work
should also investigate how the number of class and their distributions might
influence the performance of the algorithm, and also if (and how) it could be
used for unsupervised learning.

##  6 Conclusion

We introduce a new collaborative learning algorithm that enables tunable
collaboration in cross-device applications and whose uplink and downlink
communication does not scale with the model size (as in FL) or the dataset
size (as in SL). We prove that our objective is well posed from the point of
view of collaboration, as it maximizes a lower bound on the mutual information
between the feature representations of different users across the network.
Then, we show empirically that it is particularly relevant in setups where the
number of clients is large and when each of them have limited computational
resources. In that scenario, it converges faster than IL, FD and even standard
FL. Due to its simplicity and modularity with other potential requirements
(such as differential privacy, server free/peer-to-peer configurations,
heterogeneous model architectures, etc.), we posit that our algorithm could be
implemented in numerous settings where FL (or other similar frameworks) is not
suitable, as long as careful task-specific deployment considerations are made.

## References

  * Acharya et al. [2020] Jayadev Acharya, Kallista Bonawitz, Peter Kairouz, Daniel Ramage, and Ziteng Sun.  Context aware local differential privacy.  In _International Conference on Machine Learning_ , pages 52–62. PMLR, 2020.
  * Addepalli et al. [2020] Sravanti Addepalli, Gaurav Kumar Nayak, Anirban Chakraborty, and Venkatesh Babu Radhakrishnan.  Degan: Data-enriching gan for retrieving representative samples from a trained classifier.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pages 3130–3137, 2020.
  * Anil et al. [2018] Rohan Anil, Gabriel Pereyra, Alexandre Passos, Róbert Ormándi, George E. Dahl, and Geoffrey E. Hinton.  Large scale distributed neural network training through online distillation.  _CoRR_ , abs/1804.03235, 2018.  URL http://arxiv.org/abs/1804.03235.
  * Arivazhagan et al. [2019] Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary.  Federated learning with personalization layers.  _CoRR_ , abs/1912.00818, 2019.  URL http://arxiv.org/abs/1912.00818.
  * Arjovsky et al. [2019] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz.  Invariant risk minimization, 2019.
  * Bhardwaj et al. [2019] Kartikeya Bhardwaj, Naveen Suda, and Radu Marculescu.  Dream distillation: A data-independent model compression framework, 2019\.  URL https://arxiv.org/abs/1905.07072.
  * Brown et al. [2021] Gavin Brown, Marco Gaboardi, Adam Smith, Jonathan Ullman, and Lydia Zakynthinou.  Covariance-aware private mean estimation without private covariance estimation.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Bucila et al. [2006] Cristian Bucila, Rich Caruana, and Alexandru Niculescu-Mizil.  Model compression.  In _Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining_ , KDD ’06, page 535–541, New York, NY, USA, 2006. Association for Computing Machinery.  ISBN 1595933395.  doi: 10.1145/1150402.1150464.  URL https://doi.org/10.1145/1150402.1150464.
  * Chang et al. [2019] Hongyan Chang, Virat Shejwalkar, Reza Shokri, and Amir Houmansadr.  Cronus: Robust and heterogeneous collaborative learning with black-box knowledge transfer, 2019.  URL https://arxiv.org/abs/1912.11279.
  * Chen et al. [2019] Hanting Chen, Yunhe Wang, Chang Xu, Zhaohui Yang, Chuanjian Liu, Boxin Shi, Chunjing Xu, Chao Xu, and Qi Tian.  Data-free learning of student networks.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 3514–3522, 2019.
  * Chen and Chao [2020] Hong-You Chen and Wei-Lun Chao.  Feddistill: Making bayesian model ensemble applicable to federated learning.  _CoRR_ , abs/2009.01974, 2020.  URL https://arxiv.org/abs/2009.01974.
  * Deng [2012] Li Deng.  The mnist database of handwritten digit images for machine learning research.  _IEEE Signal Processing Magazine_ , 29(6):141–142, 2012.
  * Dimakis et al. [2010] Alexandros G. Dimakis, Soummya Kar, José M. F. Moura, Michael G. Rabbat, and Anna Scaglione.  Gossip algorithms for distributed signal processing.  _Proceedings of the IEEE_ , 98(11):1847–1864, 2010.  doi: 10.1109/JPROC.2010.2052531.
  * Gou et al. [2021] Jianping Gou, Baosheng Yu, Stephen J Maybank, and Dacheng Tao.  Knowledge distillation: A survey.  _International Journal of Computer Vision_ , 129(6):1789–1819, 2021.
  * Guo et al. [2020] Qiushan Guo, Xinjiang Wang, Yichao Wu, Zhipeng Yu, Ding Liang, Xiaolin Hu, and Ping Luo.  Online knowledge distillation via collaborative learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , June 2020.
  * Gupta and Raskar [2018] Otkrist Gupta and Ramesh Raskar.  Distributed learning of deep neural network over multiple agents.  _Journal of Network and Computer Applications_ , 116:1–8, 2018.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * Hinton et al. [2015] Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al.  Distilling the knowledge in a neural network.  _arXiv preprint arXiv:1503.02531_ , 2(7), 2015.
  * Itahara et al. [2020] Sohei Itahara, Takayuki Nishio, Yusuke Koda, Masahiro Morikura, and Koji Yamamoto.  Distillation-based semi-supervised federated learning for communication-efficient collaborative training with non-iid private data.  _CoRR_ , abs/2008.06180, 2020.  URL https://arxiv.org/abs/2008.06180.
  * Jeong et al. [2018] Eunjeong Jeong, Seungeun Oh, Hyesung Kim, Jihong Park, Mehdi Bennis, and Seong-Lyun Kim.  Communication-efficient on-device machine learning: Federated distillation and augmentation under non-iid private data.  _CoRR_ , abs/1811.11479, 2018.  URL http://arxiv.org/abs/1811.11479.
  * Kairouz et al. [2021] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al.  Advances and open problems in federated learning.  _Foundations and Trends® in Machine Learning_ , 14(1–2):1–210, 2021.
  * Karimireddy et al. [2020] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh.  Scaffold: Stochastic controlled averaging for federated learning.  In _International Conference on Machine Learning_ , pages 5132–5143. PMLR, 2020.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kotsogiannis et al. [2020] Ios Kotsogiannis, Stelios Doudalis, Sam Haney, Ashwin Machanavajjhala, and Sharad Mehrotra.  One-sided differential privacy.  In _2020 IEEE 36th International Conference on Data Engineering (ICDE)_ , pages 493–504. IEEE, 2020.
  * [25] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton.  Cifar-10 (canadian institute for advanced research).  URL http://www.cs.toronto.edu/~kriz/cifar.html.
  * Kulkarni et al. [2017] Mandar Kulkarni, Kalpesh Patil, and Shirish Subhash Karande.  Knowledge distillation using unlabeled mismatched images.  _CoRR_ , abs/1703.07131, 2017.  URL http://arxiv.org/abs/1703.07131.
  * LeCun et al. [1989] Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel.  Backpropagation Applied to Handwritten Zip Code Recognition.  _Neural Computation_ , 1(4):541–551, 12 1989\.  ISSN 0899-7667.  doi: 10.1162/neco.1989.1.4.541.  URL https://doi.org/10.1162/neco.1989.1.4.541.
  * Li and Wang [2019] Daliang Li and Junpu Wang.  Fedmd: Heterogenous federated learning via model distillation.  _CoRR_ , abs/1910.03581, 2019.  URL http://arxiv.org/abs/1910.03581.
  * Li et al. [2020] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith.  Federated optimization in heterogeneous networks.  _Proceedings of Machine Learning and Systems_ , 2:429–450, 2020.
  * Lian et al. [2017] Xiangru Lian, Ce Zhang, Huan Zhang, Cho-Jui Hsieh, Wei Zhang, and Ji Liu.  Can decentralized algorithms outperform centralized algorithms? a case study for decentralized parallel stochastic gradient descent.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Liang et al. [2019] Xianfeng Liang, Shuheng Shen, Jingchang Liu, Zhen Pan, Enhong Chen, and Yifei Cheng.  Variance reduced local SGD with lower communication complexity.  _CoRR_ , abs/1912.12844, 2019.  URL http://arxiv.org/abs/1912.12844.
  * Lin et al. [2020] Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi.  Ensemble distillation for robust model fusion in federated learning.  _Advances in Neural Information Processing Systems_ , 33:2351–2363, 2020.
  * Liu et al. [2021] Xiyang Liu, Weihao Kong, Sham Kakade, and Sewoong Oh.  Robust and differentially private mean estimation.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Lopes et al. [2017] Raphael Gontijo Lopes, Stefano Fenu, and Thad Starner.  Data-free knowledge distillation for deep neural networks.  _CoRR_ , abs/1710.07535, 2017.  URL http://arxiv.org/abs/1710.07535.
  * McMahan et al. [2017] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In _Artificial intelligence and statistics_ , pages 1273–1282. PMLR, 2017.
  * Micaelli and Storkey [2019] Paul Micaelli and Amos J Storkey.  Zero-shot knowledge transfer via adversarial belief matching.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Nayak et al. [2019] Gaurav Kumar Nayak, Konda Reddy Mopuri, Vaisakh Shaj, Venkatesh Babu Radhakrishnan, and Anirban Chakraborty.  Zero-shot knowledge distillation in deep networks.  In _International Conference on Machine Learning_ , pages 4743–4751. PMLR, 2019.
  * Nayak et al. [2021] Gaurav Kumar Nayak, Konda Reddy Mopuri, and Anirban Chakraborty.  Effectiveness of arbitrary transfer sets for data-free knowledge distillation.  In _Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)_ , pages 1430–1438, January 2021.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems 32_ , pages 8024–8035. Curran Associates, Inc., 2019.  URL http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf.
  * Sannara et al. [2021] EK Sannara, François Portet, Philippe Lalanda, and VEGA German.  A federated learning aggregation algorithm for pervasive computing: Evaluation and comparison.  In _2021 IEEE International Conference on Pervasive Computing and Communications (PerCom)_ , pages 1–10. IEEE, 2021.
  * Shah [2009] Devavrat Shah.  Gossip algorithms.  _Found. Trends Netw._ , 3(1):1–125, jan 2009\.  ISSN 1554-057X.  doi: 10.1561/1300000014.  URL https://doi.org/10.1561/1300000014.
  * Sodhani et al. [2020] Shagun Sodhani, Olivier Delalleau, Mahmoud Assran, Koustuv Sinha, Nicolas Ballas, and Michael G. Rabbat.  A closer look at codistillation for distributed training.  _CoRR_ , abs/2010.02838, 2020.  URL https://arxiv.org/abs/2010.02838.
  * Tian et al. [2020] Yonglong Tian, Dilip Krishnan, and Phillip Isola.  Contrastive representation distillation.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=SkgpBJrtvS.
  * Tishby and Zaslavsky [2015] Naftali Tishby and Noga Zaslavsky.  Deep learning and the information bottleneck principle.  In _2015 ieee information theory workshop (itw)_ , pages 1–5. IEEE, 2015.
  * van den Oord et al. [2018] Aäron van den Oord, Yazhe Li, and Oriol Vinyals.  Representation learning with contrastive predictive coding.  _CoRR_ , abs/1807.03748, 2018.  URL http://arxiv.org/abs/1807.03748.
  * van der Maaten and Hinton [2008] Laurens van der Maaten and Geoffrey Hinton.  Visualizing data using t-sne.  _Journal of Machine Learning Research_ , 9(86):2579–2605, 2008.  URL http://jmlr.org/papers/v9/vandermaaten08a.html.
  * Vanhaesebrouck et al. [2016] Paul Vanhaesebrouck, Aurélien Bellet, and Marc Tommasi.  Decentralized collaborative learning of personalized models over networks.  _CoRR_ , abs/1610.05202, 2016.  URL http://arxiv.org/abs/1610.05202.
  * Vepakomma et al. [2020] Praneeth Vepakomma, Abhishek Singh, Otkrist Gupta, and Ramesh Raskar.  Nopeek: Information leakage reduction to share activations in distributed deep learning.  In _2020 International Conference on Data Mining Workshops (ICDMW)_ , pages 933–942. IEEE, 2020.
  * Wang et al. [2020a] Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris S. Papailiopoulos, and Yasaman Khazaeni.  Federated learning with matched averaging.  _CoRR_ , abs/2002.06440, 2020a.  URL https://arxiv.org/abs/2002.06440.
  * Wang et al. [2020b] Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor.  Tackling the objective inconsistency problem in heterogeneous federated optimization.  _Advances in neural information processing systems_ , 33:7611–7623, 2020b.
  * Warnat-Herresthal et al. [2021] Stefanie Warnat-Herresthal, Hartmut Schultze, Krishnaprasad Lingadahalli Shastry, Sathyanarayanan Manamohan, Saikat Mukherjee, Vishesh Garg, Ravi Sarveswara, Kristian Händler, Peter Pickkers, N Ahmad Aziz, et al.  Swarm learning for decentralized and confidential clinical machine learning.  _Nature_ , 594(7862):265–270, 2021.
  * Wu et al. [2022] Chuhan Wu, Fangzhao Wu, Lingjuan Lyu, Yongfeng Huang, and Xing Xie.  Communication-efficient federated learning via knowledge distillation".  _Nat Commun_ , 13(1):2032, apr 2022.
  * Xiao et al. [2017] Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  _CoRR_ , abs/1708.07747, 2017.  URL http://arxiv.org/abs/1708.07747.
  * Zhang et al. [2018] Ying Zhang, Tao Xiang, Timothy M Hospedales, and Huchuan Lu.  Deep mutual learning.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 4320–4328, 2018.

##  Appendix A Appendix

###  A.1 Notation

  * •

$N$: Number of participating users/clients/peers/data owners.

  * •

$d$: Raw input data dimension (e.g., $d=3072$ for $32\times 32$ RBG images).

  * •

$d^{\prime}$: Latent feature (or feature representation) space dimensionality
(i.e., width of last hidden layer).

  * •

$C$: Number of class.

  * •

$(\mbf{x},y)\in\mathbb{R}^{d}\times\\{0,...,C-1\\}$: Labeled data sample.

  * •

$p_{\mbf{X},Y,U}:\mathcal{\mathbb{R}}^{d}\times\\{1,...,C\\}\times\\{1,...,U\\}\to\mathbb{R}_{+}$:
Joint probability density function across users.

  * •

$p_{u}(\mbf{x},y):=p_{\mbf{X},Y|U=u}(\mbf{x},y)$: Data distribution of user
$u$.

  * •

$\mathcal{D}_{u}:=\\{(\mbf{x}_{i},y_{i})\overset{iid}{\sim}p_{u}\\}_{i=1}^{n_{i}}$:
Dataset of user $u$.

  * •

$\mbf{w}_{u}:=\\{{\bm{\theta}}_{u},{\bm{W}}_{u},\mbf{b}_{u}\\}$ with
${\bm{\theta}}_{u}\in\Theta_{u},{\bm{W}}_{u}\in\mathbb{R}^{C\times
d^{\prime}},\mbf{b}_{u}\in\mathbb{R}^{C}$: Parameters of the neural network of
$u$, with $\Theta_{u}$ the achievable model parameters for user $u$.

  * •

$\phi_{u}:\mathbb{R}^{d}\to\mathbb{R}^{d^{\prime}},\
\mbf{x}\mapsto\mbf{x}^{\prime}=\phi(\mbf{x};{\bm{\theta}}_{u})$: Neural
network of $u$ (or similar parameterized function) that maps a raw input into
a latent feature space.

  * •

$\tau_{u}:\mathbb{R}^{d^{\prime}}\to\mathbb{R}^{C},\
\mbf{x}^{\prime}\mapsto\mbf{z}=\tau(\mbf{x}^{\prime};\mbf{W}_{u},\mbf{b}_{u}):=\mbf{W}_{u}\mbf{x}^{\prime}+\mbf{b}_{u}$:
Linear classifier of $u$.

  * •

$\lambda_{KD},\lambda_{disc},n_{avg},M_{\uparrow},M_{\downarrow}$:
Hyperparameters.

  * •

$\eta$: Learning rate.

  * •

$\ell_{CE},\ell_{KD},\ell_{disc}$: Cross-entropy, feature-based KD and
discriminator loss functions, respectively.

  * •

$\mathcal{L}_{CE},\mathcal{L}_{KD},\mathcal{L}_{disc}$: Expected value (over
the data) of $\ell_{CE},\ell_{KD}$ and $\ell_{disc}$, respectively.

  * •

$L_{CE},L_{KD},L_{disc}$: Mini-batch estimates of
$\mathcal{L}_{CE},\mathcal{L}_{KD}$ and $\mathcal{L}_{disc}$, respectively.

  * •

${\bm{\Phi}}_{s},{\bm{\Phi}}_{t}$: Random vectors (feature representations) of
the student (user $u$) and the teacher (random user $U$).

  * •

$p_{s,t},p_{s},p_{t}$: Joint and marginal distributions of ${\bm{\Phi}}_{s}$
and ${\bm{\Phi}}_{t}$, respectively.

  * •

$q$: Joint distribution of ${\bm{\Phi}}_{s},{\bm{\Phi}}_{t}$ and $I$, where
$I$ is a binary random variable indicating if
${\bm{\Phi}}_{s},{\bm{\Phi}}_{t}$ has been drawn from $p_{s,t}$ ($I=1$) or
$p_{s}p_{t}$ ($I=0$).

  * •

$h$: Binary discriminator with Bernoulli parameter $\hat{h}$ (i.e., learnable
estimate of $q_{I|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t}}$).

  * •

$\mbf{s},\mbf{t}$: Observation/realization of ${\bm{\Phi}}_{s}$ and
${\bm{\Phi}}_{t}$, respectively.

  * •

$\bar{\mbf{t}}^{c},\mbf{t}^{c}$: Global (i.e., using all the samples across
users) and local (i.e., using $n_{avg}$ samples) average feature
representations of class $c$, respectively.

###  A.2 Proof of Theorem 1

Recall that $q(\mbf{s},\mbf{t},i)$ is the joint distribution of
$({\bm{\Phi}}_{s},{\bm{\Phi}}_{t},I)$ such that
$q(\mbf{s},\mbf{t}|i=1)=p_{s,t}(\mbf{s},\mbf{t})$ and
$q(\mbf{s},\mbf{t}|i=0)=p_{s}(\mbf{s})p_{t}(\mbf{t})$ and suppose that the
prior $q(i)$ satisfy $q(i=1)=\frac{1}{K+1}$ and $q(i=0)=\frac{K}{K+1}$, i.e.,
for each sample from the distribution $p_{s,t}$, we draw $K$ samples from the
distribution $p_{s}p_{t}$. We have:

$\displaystyle\mathcal{I}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})$
$\displaystyle=\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
p_{s,t}}\left[-\log\frac{p_{s}({\bm{\Phi}}_{s})p_{t}({\bm{\Phi}}_{t})}{p_{s,t}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})}\right]$
(8) $\displaystyle=\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
p_{s,t}}\left[-\log\left(K\frac{p_{s}({\bm{\Phi}}_{s})p_{t}({\bm{\Phi}}_{t})}{p_{s,t}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})}\right)\right]+\log(K)$
(9) $\displaystyle\geq\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
p_{s,t}}\left[-\log\left(1+K\frac{p_{s}({\bm{\Phi}}_{s})p_{t}({\bm{\Phi}}_{t})}{p_{s,t}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})}\right)\right]+\log(K)$
(10) $\displaystyle=\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
p_{s,t}}\left[\log q(i=1|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]+\log(K)$ (11)

where the last equality is obtained using the Bayes’ rule on the posterior
$q(i=1|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})$:

$\displaystyle q(i=1|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})$
$\displaystyle=\frac{q({\bm{\Phi}}_{s},{\bm{\Phi}}_{t}|i=1)q(i=1)}{q({\bm{\Phi}}_{s},{\bm{\Phi}}_{t}|i=0)q(i=0)+q({\bm{\Phi}}_{s},{\bm{\Phi}}_{t}|i=1)q(i=1)}$
(12)
$\displaystyle=\frac{p_{s,t}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})}{Kp_{s}({\bm{\Phi}}_{s})p_{t}({\bm{\Phi}}_{t})+p_{s,t}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})}$
(13)
$\displaystyle=\left(1+K\frac{p_{s}({\bm{\Phi}}_{s})p_{t}({\bm{\Phi}}_{t})}{p_{s,t}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})}\right)^{-1}.$
(14)

Hence, by optimizing $\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
p_{s,t}}\left[\log q(i=1|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]$ with respect
to the model parameters ${\bm{\theta}}_{u}$ of the student, we optimize a
lower bound on the mutual information between
${\bm{\Phi}}_{s},{\bm{\Phi}}_{t}$. By noting that $\log
q(i=0|{\bm{\Phi}}_{s},{\bm{\Phi}}_{s})\leq 0$, we can further bound the
expectation term in (11) as follows:

$\displaystyle\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
p_{s,t}}\left[\log q(i=1|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]$
$\displaystyle\geq\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
q|I=1}\left[\log q(i=1|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]$
$\displaystyle\quad+K\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
q|I=0}\left[\log q(i=0|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]$ (15)
$\displaystyle=(K+1)\sum_{i}q(i)\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
q|I=i}\left[\log q(i|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]$ (16)
$\displaystyle=(K+1)\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t},I)\sim
q}\left[\log q(I|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]$ (17)
$\displaystyle=(K+1)\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
q}\left[\mathbb{E}_{I\sim q|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t}}\left[\log
q(I|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]\right]$ (18)

However, similar to Tian et al. [43], the Bernoulli distribution
$q(i|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})$ is unknown and must therefore be
approximated by training a discriminator
$h:\\{0,1\\}\times\mathbb{R}^{d^{\prime}}\times\mathbb{R}^{d^{\prime}}\to[0,1]$.
Using Gibbs’ inequality, we obtain

$-\mathbb{E}_{I\sim q|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t}}\left[\log
q(I|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]\leq-\mathbb{E}_{I\sim
q|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t}}\left[\log
h(I,{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right],$ (19)

where the right-hand term is the expected negative log-likelihood loss of the
discriminator for a particular set $({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})$. Hence,
Eq. (18) is proportional to minus the expected loss of the discriminator. Let
$\hat{h}(\mbf{s},\mbf{t})\in[0,1]$ denote the Bernoulli parameter of $h$ given
the data $(\mbf{s},\mbf{t})$ (i.e.,
${h(i,\mbf{s},\mbf{t})=\hat{h}(\mbf{s},\mbf{t})^{i}(1-\hat{h}(\mbf{s},\mbf{t}))^{1-i}}$),
we define our learning objective for the discriminator as follows:

$\displaystyle-\mathcal{L}_{disc}(h,\phi_{u})$
$\displaystyle:=(K+1)\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
q}\left[\mathbb{E}_{I\sim q|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t}}\left[\log
h(I,{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]\right]$ (20)
$\displaystyle=(K+1)\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t},I)\sim
q}\left[\log h(I|{\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]$ (21)
$\displaystyle=(K+1)\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t},I)\sim
q}\left[\log\left(\hat{h}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})^{I}(1-\hat{h}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t}))^{(1-I)}\right)\right]$
(22) $\displaystyle=+(K+1)\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
q|I=1}\left[\log\hat{h}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]q(i=1)$
$\displaystyle\quad(K+1)\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
q|I=0}\left[\log\left(1-\hat{h}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right)\right]q(i=0)$
(23) $\displaystyle=\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
p_{s,t}}\left[\log\hat{h}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\right]+K\mathbb{E}_{({\bm{\Phi}}_{s},{\bm{\Phi}}_{t})\sim
p_{s}p_{t}}\left[\log(1-\hat{h}({\bm{\Phi}}_{s},{\bm{\Phi}}_{t}))\right],$
(24)

which concludes the derivation.

###  A.3 Algorithms

Input: Server $S$, $N$ users with local datasets
$\\{\mathcal{D}_{u}\\}_{u=1}^{N}$.

$S$ initializes randomly $\\{\bar{\mbf{t}}^{c}\\}_{c=1}^{C}$ and random
observations $\\{\\{\mbf{t}_{m}^{c}\\}_{c=1}^{C}\\}_{m=1}^{N\cdot
M_{\uparrow}}$;

Each client $u$ initializes
$\mbf{w}_{u}^{0}:=\\{{\bm{\theta}}^{0}_{u},\mbf{W}^{0}_{u},\mbf{b}^{0}_{u}\\}$;

$r\leftarrow 0$;

while  _training:_

$r\leftarrow r+1$;

for  _each client $u$:_

$u$ downloads $\\{\bar{\mbf{t}}^{c}\\}_{c=1}^{C}$ and $M_{\downarrow}$ random
observations $\\{\\{\mbf{t}_{m}^{c}\\}_{c=1}^{C}\\}_{m=1}^{M_{\downarrow}}$;

$\mbf{w}_{u}^{r}\leftarrow\mbox{{LocalUpdate}}(\mathcal{D}_{u},\mbf{w}_{u}^{r-1},\\{\bar{\mbf{t}}^{c}\\}_{c=1}^{C},\\{\\{\mbf{t}_{m}^{c}\\}_{c=1}^{C}\\}_{m=1}^{M_{\downarrow}}$);

$u$ computes and uploads its local averaged representations
$\\{\bar{\mbf{t}}_{u}^{c}\\}_{c=1}^{C}$;

$u$ computes (using $n_{avg}$) and uploads $M_{\uparrow}$ observations
$\\{\\{\mbf{t}_{u,m}^{c}\\}_{c=1}^{C}\\}_{m=1}^{M_{\downarrow}}$;

$S$ aggregates $\\{\\{\bar{\mbf{t}}_{u}^{c}\\}_{c=1}^{C}\\}_{u=1}^{N}$ to
obtain $\\{\bar{\mbf{t}}^{c}\\}_{c=1}^{C}$;

$S$ stores (and shuffles)
$\\{\\{\mbf{t}_{u,m}^{c}\\}_{c=1}^{C}\\}_{m=1}^{M_{\downarrow}}$ in their
corresponding class buffers;

return $\\{\mbf{w}_{u}^{r}\\}_{u=1}^{N}$

Algorithm 1 GlobalUpdate

Input: Local dataset $\mathcal{D}_{u}$, model parameters
$\mbf{w}_{u}=\\{{\bm{\theta}}_{u},\mbf{W}_{u},\mbf{b}_{u}\\}$, global features
$\\{\bar{\mbf{t}}^{c}\\}_{c=1}^{C}$ and observations
$\\{\\{\mbf{t}_{m}^{c}\\}_{c=1}^{C}\\}_{m=1}^{M_{\downarrow}}$, number of
local training rounds $E$, averaging parameter $n_{a}vg$.

Buffer initialization (per class):
$\\{{\bm{\hat{\Phi}}}_{u}^{c}\leftarrow\mbf{0}\\}_{c=1}^{C}$;

for  _$e\in[1,...,E]$ :_

for  _mini-batch $\mathcal{B}\in\mathcal{D}_{u}$:_

$L_{KD},L_{CE},L_{disc}\leftarrow 0$;

for  _$(\mbf{x}_{i},y_{i})\in\mathcal{B}$ :_

$\mbf{s}_{i}\leftarrow\phi_{u}(\mbf{x}_{i})$;

$L_{CE}\leftarrow
L_{CE}+\frac{1}{|\mathcal{B}|}\ell_{CE}(\tau_{u}(\mbf{s}_{i}),y_{i})$;

$L_{KD}\leftarrow
L_{KD}+\frac{1}{|\mathcal{B}|}\ell_{KD}(\mbf{s}_{i},\bar{\mbf{t}}^{y_{i}})$;

Sample $m\sim\mbox{{Uniform}}(1,...,M_{\downarrow})$;

$L_{disc}\leftarrow
L_{disc}+\frac{1}{|\mathcal{B}|}\left(\ell_{disc}(1,\mbf{s}_{i},\mbf{t}_{m}^{y_{i}})+\sum_{c\neq
y_{i}}\ell_{disc}(0,\mbf{s}_{i},\mbf{t}_{m}^{c})\right)$;

$L\leftarrow L_{CE}+\lambda_{KD}L_{KD}+\lambda_{disc}L_{disc}$;

$\mbf{w}_{u}\leftarrow\mbf{w}_{u}-\eta\nabla_{\mbf{w}_{u}}L$;

return $\mbf{w}_{u}$

Algorithm 2 LocalUpdate