# Personalization Disentanglement for Federated Learning: An explainable
perspective

1st Peng Yan  Faculty of Engineering and IT
University of Technology Sydney
Sydney, Australia
<EMAIL_ADDRESS>2nd Guodong Long  Faculty of Engineering and IT
University of Technology Sydney
Sydney, Australia
<EMAIL_ADDRESS>
###### Abstract

Personalized federated learning (PFL) jointly trains a variety of local models
through balancing between knowledge sharing across clients and model
personalization per client. This paper addresses PFL via explicit
disentangling latent representations into two parts to capture the shared
knowledge and client-specific personalization, which leads to more reliable
and effective PFL. The disentanglement is achieved by a novel Federated Dual
Variational Autoencoder (FedDVA), which employs two encoders to infer the two
types of representations. FedDVA can produce a better understanding of the
trade-off between global knowledge sharing and local personalization in PFL.
Moreover, it can be integrated with existing FL methods and turn them into
personalized models for heterogeneous downstream tasks. Extensive experiments
validate the advantages caused by disentanglement and show that models trained
with disentangled representations substantially outperform those vanilla
methods.

###### Index Terms:

Federated Learning, Disentanglement, Variational Autoencoder

##  I Introduction

With increasing attention to privacy protection, federated learning (FL) [1]
has recently been a hot topic in machine learning. Vanilla FL tasks require
learning a global model collaboratively by clients while keeping their data
decentralized and private. Various methods have been proposed under this
constraint and advanced in capturing universal knowledge from data on
different clients [2, 3, 4, 5, 6]. Meanwhile, samples in FL contain universal
knowledge applicable all over the federation and demonstrate their host
client’s bias as personalized knowledge. Then personalized federated learning
(PFL) is proposed to learn many local models simultaneously by leveraging
universal knowledge and client biases [7].

In a PFL task, clients will first update local models with entangled raw
sample representations, which contain both the universal and personalized
knowledge, and then eliminate personalized impacts to update a global model to
share universal knowledge. Most works perform these operations by designing
new model architectures or adapting optimization strategies [8, 9, 10, 11,
12]. However, how to solve the PFL challenge from representation perspectives
still need to be studied. One can encode a sample as two disentangled
representations, each capturing one type of the above knowledge. Hence other
FL algorithms for downstream tasks may learn from them separately,
facilitating in extracting and sharing of universal knowledge. On the other
hand, disentangled representation of personalized knowledge will help identify
essential knowledge constituting a client’s personality, which can support a
better understanding of the locally learned models. An example of the
disentangled sample representations is in Fig.1.

Figure 1: Suppose samples (digits) on different clients are entangled with
client-related features (sinusoidal and elliptical marks). Our goal is to
learn disentangled representations that each capture one type of feature.

To this end, we develop a federated dual variational autoencoding framework
(FedDVA), where clients in the federation share two encoders inferring the
above representations. The two encoders are trained collaboratively by
fundamental FL algorithms like FedAvg [1]. Clients update the encoders locally
by maximizing client-specific evidence lower bound (ELBO). Then a server
collects local updates and aggregates them by averaging parameters. Moreover,
the two encoders are cascaded and constrained by different prior knowledge so
that each encoder will capture only one type of knowledge mentioned above.

We evaluate the performance of FedDVA from the perspective of disentanglement
and PFL tasks. First, we explore the disentangled representations through
reconstructed samples, demonstrating two different data manifolds
corresponding to universal and personalized knowledge. In addition, we train
off-the-shelf FL models on the disentangled representations for classification
tasks and show they will converge fast and achieve better accuracy despite
various client personalities.

The main contributions of this work are summarized as follows:

  * •

We propose a novel FedDVA method to solve the PFL challenge from
representation perspectives. It infers disentangled representations of
universal knowledge and personalized knowledge in FL.

  * •

We introduce a client-specific ELBO to optimize FedDVA and analyze its
capability in personalization.

  * •

Experiments on real-world datasets validate FedDVA’s effectiveness in
disentanglement and show that FL models will converge fast and achieve
competitive classification performance when trained on disentangled
representations.

##  II Related Work

Representation Disentanglement Disentanglement has been extensively studied in
unsupervised learning. It refers to learning a representation where a change
in one dimension corresponds to a change in one factor of variation of the
sample while being relatively invariant to changes in other factors [13].
Variational autoencoder (VAE) and its variations are popular frameworks for
learning disentangled representations [14, 15]. It is attractive for elegant
theoretical backgrounds and high computation efficiency. However, [16] proved
that unsupervised learning of disentangled representations is fundamentally
impossible without inductive biases on both the models and the data. Then
existing VAE models will fail in disentangling universal knowledge and client
preferences in FL since there needs to be supervised information or ad-hoc
inductive biases to distinguish between them.

Personalized FL Most methods solve the PFL challenge by designing new model
architecture or adapting optimization strategies. [17] designed a central
hypernetwork model, which is trained to generate a set of models, one model
for each client. [12] proved that steps as simple as fine-tuning the client’s
local data would improve the performance of a global model. [11] adapts the
idea from MAML [18], where clients co-learn an initialization for local models
through the FedAvg. [10] proposed the Ditto, where clients share global
parameters to constrain the learning process of personalized models.

##  III Personalization Disentanglement for Federated Learning

Figure 2: The architecture of FedDVA. An encoder $f(x)$ (Blue) will first
infer the posterior $q(z|x)$, and then another encoder $h(x,z)$ (red) will
infer the conditional posterior $q(c|x,z)$. The decoder $g(z,c)$ (white) will
try to reconstruct $x$ from $z$ and $c$.

###  III-A Background and Notation

Federated Learning In a federated learning system with $K$ clients, each
client is indexed by $k$ with its local data denoted as $\mathcal{D}_{k}$. The
objective of a federated learning task is to find the optimal parameters
$\theta^{*}$ sharing with all clients, which is formulated as

$\displaystyle\theta^{*}=\operatorname*{arg\,min}_{\theta}\sum_{k=1}^{K}w_{k}\mathcal{L}_{k}(\theta;\mathcal{D}_{k})$
(1)

where
$\mathcal{L}_{k}(\theta;\mathcal{D}_{k})=\sum_{x\in\mathcal{D}_{k}}\ell(\theta;x)$
is the learning objective on the $k$-th client,
$w_{k}=|\mathcal{D}_{k}|/\sum_{k=1}^{K}|\mathcal{D}_{k}|$ is an importance
weight for the client, and $\ell(\theta;x)$ is a loss function

VAE Framework VAE assumes any sample $x$ corresponds to a latent
representation $z$ from the prior $p(z)=\mathcal{N}(z;0,I)$. It learns an
encoder to infer the variational posterior $q_{\theta}(z|x)$ and a decoder to
reconstruct the sample $x$ from $z\sim q_{\theta}(z|x)$. In general, the
encoder is a neural network whose outputs are the mean and covariance of the
variational posterior $q_{\theta}(z|x)$, that is,
$q_{\theta}(z|x)=\mathcal{N}(z;\mu(x),\Sigma(x))$. The covariance matrix
$\Sigma$ is assumed to be diagonal for computation simplicity. The decoder is
another neural network generating $x$ by maximizing the log-likelihood $\log
p_{\varphi}(x|z)$. $\theta$ and $\varphi$ are learnable parameters of the
encoder and the decoder. They can be optimized by maximizing the below ELBO

$\displaystyle EB(x)=\mathbb{E}_{q_{\theta}(z|x)}[\log
p_{\varphi}(x|z)]-D_{KL}(q_{\theta}(z|x)||p(z))$ (2)

The first term on the RHS of Eq.2 measures the reconstruction performance of
latent representation $z$, and the second term measures the $KL$-divergence
between the posterior $q_{\theta}(z|x)$ and the prior $p(z)$. Gradient-based
optimization methods apply with the reparameterization trick [19]. In the rest
of the paper, we will use $\ell(\theta,\varphi;x)$ to denote the negative ELBO
and omit the subscripts $\theta$ and $\varphi$ for notational simplicity,
i.e.,

$\displaystyle\ell(\theta,\varphi;x)=-\mathbb{E}_{q(z|x)}[\log
p(x|z)]+D_{KL}(q(z|x)||p(z))$ (3)

###  III-B Methodology

Problem Formulations The goal of our method is to learn disentangled sample
representations for universal knowledge and personalized knowledge. We denote
them as $z$ and $c$. Since $z$ is irrelevant to clients, we assume samples on
each client have the same prior distribution $p(z)=\mathcal{N}(z;0,I)$.
Meanwhile, since samples in FL are private and distributed, the prior
distribution of $c$ is unknown and varies among clients. We denote it as
$p_{k}(c)$, where $k$ is the index of the $k$-th client. But we can not make
assumptions about the $p_{k}(c)$ as we can not ensure that the relationship
between the assumed one is consistent with the relationship between clients.
For example, we can only allocate the same $p_{k}(c)$ to two clients after
disclosing that they have similar personalities. Alternatively, we assume the
prior distribution of the data in the federation is standard Gaussian, or
equivalently, the mixture distribution of local priors
$q(c)=\sum_{k=1}^{K}w_{k}p_{k}(c)$ is $\mathcal{N}(c;0,I)$.

Dual Encoders As illustrated in Fig.2, the proposed FedDVA learns the above
representations through two encoders. For any sample $x$, an encoder first
infers variational posterior $q(z|x)=\mathcal{N}(z;\mu(x),\Sigma(x))$ for the
universal knowledge $z$. Then another encoder infers variational posterior
$q(c|x,z)=\mathcal{N}(c;\hat{\mu}(x,z),\hat{\Sigma}(x,z))$, conditioned on
both the sample $x$ and the representation $z$, for the impacts of
personalities. In addition, a client-specific local decoder will evaluate the
reconstruction performance of representations $z$ and $c$. It is implemented
by a neural network maximizing the client-specific log-likelihood $\log
p_{k}(x|z,c)$. The negative ELBO optimizing FedDVA is in Eq.4

$\displaystyle\ell^{cva}(\theta,\varphi_{k};x)=-\mathbb{E}_{q(z|x)}[\mathbb{E}_{q(c|x,z)}[\log
p_{k}(x|z,c)]$ (4)
$\displaystyle+\beta\mathcal{R}_{c}(q(c|x,z))]+\alpha\mathcal{R}_{z}(q(z|x))$

$\theta$ in the Eq.4 denotes parameters of the shared encoders, $\varphi_{k}$
denotes the parameters of the local decoder specific to the $k$-th client,
$\mathcal{R}_{z}(q(z|x))$ and $\mathcal{R}_{c}(q(c|x,z))$ denote the
regularizers for the posterior $q(z|x)$ and $q(c|x,z)$, $\alpha$ and $\beta$
are their importance weights.

Similar to traditional VAE models, the posterior $q(z|x)$ is regularized by
$D_{KL}(q(z|x)||p(z))$, which enforces the distribution of the representation
$z$ to approximate the standard Gaussian distribution. But it would be
challenging to regularize the representation $c$ without prior knowledge about
the distribution $p_{k}(c)$. FedDVA handles the problem by a slack regularizer
$D_{KL}(q(c|x,z)||q(c))$ combing with a constraint that

$\displaystyle
D_{KL}(q(c|x,z)||q(c))-D_{KL}(q(c|x,z)||\bar{p}_{k}(c))\geq\xi_{k}$ (5)

where
$\bar{p}_{k}(c)=\frac{1}{|\mathcal{D}_{k}|}\sum_{x\in\mathcal{D}_{k}}q(c|x,z)$
is the mixture distribution of $q(c|x,z)$ of samples on the $k$-th client, and
$\xi_{k}>0$ is a hyperparameter. Intuitively, $\bar{p}_{k}(c)$ is an estimator
of $p_{k}(c)$, and the Ineuqation.5 requires $q(c|x,z)$ to be at least
$\xi_{k}$ closer to $\bar{p}_{k}(c)$ than $q(c)$. We will discuss it in Sec.IV
and show it helps the representation $c$ to capture client personalities.
Combining the $KL$-divergence and the constraint in Inequation.5, regularizers
$\mathcal{R}_{z}(q(z|x))$ and $\mathcal{R}_{c}(q(c|x,z))$ of Eq.4 are

$\displaystyle\mathcal{R}_{z}(q(z|x))=D_{KL}(q(z|x)||p(z))$ (6)
$\displaystyle\mathcal{R}_{c}(q(c|x,z))=\max($
$\displaystyle\xi_{k}+D_{KL}(q(c|x,z)||\bar{p}_{k}(c))$ (7)
$\displaystyle,D_{KL}(q(c|x,z)||q(c)))$

They can be computed and differentiated without estimation (see Appendix B).
Accordingly, the learning problem of FedDVA can be solved by gradient-based
methods.

Optimization To learn the encoders collaboratively, we formulate the learning
objective of FedDVA as follows:

$\displaystyle\operatorname*{arg\,min}_{\theta,\varphi_{1}...\varphi_{K}}\sum_{k=1}^{K}w_{k}\mathcal{L}_{k}(\theta,\varphi_{k};\mathcal{D}_{k})$
(8)

where
$\mathcal{L}_{k}(\theta,\varphi_{k};\mathcal{D}_{k})=\sum_{x\in\mathcal{D}_{k}}\ell^{cva}(\theta,\varphi_{k};x)$.
Then gradient steps optimizing Eq.8 consist of the following two parts

$\displaystyle\varphi_{k}^{\prime}=\varphi_{k}-\eta\nabla_{\varphi_{k}}\mathcal{L}_{k}(\theta,\varphi_{k};\mathcal{D}_{k}),1\leq
k\leq K$ (9)
$\displaystyle\theta^{\prime}=\theta-\lambda\sum_{k=1}^{K}w_{k}\nabla_{\theta}\mathcal{L}_{k}(\theta,\varphi_{k};\mathcal{D}_{k})$
(10)

where $\eta$ and $\lambda$ are their learning rates. Eq.9 updates the client-
specific decoders and is processed by each client independently. Eq.10 updates
the shared encoders shared in the federation. Most FL algorithms like FedAvg
can optimize it. Concretely,
$\theta^{\prime}=\sum_{k=1}^{K}w_{k}\theta^{\prime}_{k}$, where

$\displaystyle\theta^{\prime}_{k}=\theta_{k}-\lambda\nabla_{\theta}\mathcal{L}_{k}(\theta_{k},\varphi_{k};\mathcal{D}_{k}),1\leq
k\leq K$ (11)

, and Eq.11 is performed by each client independently. But it is worth noting
that the optimization steps of Eq.9 and Eq.11 are asynchronous. As only a
subset of clients will participate in the optimization process in each
communication round [1], client-specific decoders may not coincide with the
shared encoders. A client needs to update $\varphi_{k}$ first and later the
$\theta$. Complete pseudo-codes of the optimization process are in
Algorithm.1.

Algorithm 1 FedDVA

Input: $m$: number of clients sampled each round; $b$: batch size; $\lambda$
and $\eta$: learning rates; $\xi_{k}$: the constraint threshold in Inequation
(5).
Server executes:

Initialize $\theta^{(1)}\leftarrow\theta$

for each round $r=1,2,...$ do

Sample a set of $m$ clients $\mathbb{C}$

for each client $k\in\mathbb{C}$ parallel do

$\theta_{k}^{(r+1)}\leftarrow$ClientUpdate$(k,\theta^{(r)})$

end for

$\theta^{(r+1)}\leftarrow\sum_{k=1}^{K}w_{k}\theta_{k}^{(r+1)}$

end for

ClientUpdate$(k,\theta)$:

Initialize $\theta_{k}\leftarrow\theta$, $\varphi_{k}\leftarrow\varphi$

for batch $\mathcal{B}\subset\mathcal{D}_{k}$ do

update $\varphi_{k}$ by Eq.9

end for

$\varphi\leftarrow\varphi_{k}$

for batch $\mathcal{B}\subset\mathcal{D}_{k}$ do

update $\theta_{k}$ by Eq.11

end for

return $\theta_{k}$

##  IV Theoretical Analysis

In this section, we discuss the ELBO corresponding to Eq.4 and show that it
has the capability to capture client personalities.

From the perspective of variational inference, the optimal posteriors $q(z|x)$
and $q(c|x,z)$ are the ones maximizing the following EBLOs jointly

$\displaystyle ELBO_{z}(x,k)=$ $\displaystyle\mathbb{E}_{q(z|x)}[\log
p_{k}(x|z)]$ (12) $\displaystyle-D_{KL}(q(z|x)||p(z))$ $\displaystyle
ELBO_{c}(x,z,k)=$ $\displaystyle\mathbb{E}_{q(c|x,z)}[\log p(x|z,c)]$ (13)
$\displaystyle-D_{KL}(q(c|x,z)||p_{k}(c))$

where the subscript $k$ means the distribution is specific to the $k$-th
client. Ideally, $\log p(x|z,c)$ is a client irrelevant log-likelihood
modeling the sample generating process, that is, $p_{k}(x)=\iint
p(x|z,c)p(z)p_{k}(c)\text{d}z\text{d}c$ (Details of the derivation are given
in Appendix A.1 and A.2). But Eq.13 is hard to compute in practice. Besides
the unknown prior knowledge $p_{k}(c)$, the client irrelevant log-likelihood
$\log p(x|z,c)$ is unavailable in FL. For example, sharing $\log p(x|z,c)$ in
the federation risks privacy leakage as it has the capability to generate
samples.

As an alternative, FedDVA optimizes the posterior $q(c|x,z)$ by maximizing the
ELBO in Eq.14

$\displaystyle ELBO^{\prime}_{c}(x,z,k)=$
$\displaystyle\mathbb{E}_{q(c|x,z)}[\log p_{k}(x|z,c)]$ (14) $\displaystyle-
D_{KL}(q(c|x,z)||q(c))$

which is equivalent to Eq.13, except for that the slack regularizer
$D_{KL}(q(c|x,z)||q(c))$ degenerates the capability of capturing differences
between clients. Specifically, the overall $KL$-divergence between $q(c|x,z)$
and $q(c)$ of samples on the same client is

$\displaystyle\mathbb{E}_{p_{k}(x)}[\mathbb{E}_{q(z|x)}[-H(q(c|x,z))]]+H(\bar{p}_{k}(c),q(c))$
(15)

which requires the distribution of representation $c$ to be close to $q(c)$
wherever the samples are. Inequation.5 helps resolve the problem by
introducing an inductive bias that the posterior $q(c|x,z)$ of samples on the
same client is closer to $p_{k}(c)$ than to $q(c)$, with which
$D_{KL}(\bar{p}_{k}(c)||q(c))\geq\xi_{k}$ holds (Details of the derivation are
in Appendix A.3). Finally, replacing $p_{k}(x|z)$ in Eq.12 with Eq.14, we have
the loss function described in Eq.4, and the hyperparameter $\xi_{k}$ helps
determinate the degree of ’penalization’ representation $c$ captured. The
larger the $\xi_{k}$ is, the more personalized representation $c$ is learned.

##  V Experiment

In this section, we verify FedDVA’s disentanglement effectiveness by
demonstrating manifolds of the sample reconstructed from the disentangled
representations $z$ and $c$. Then, we evaluate classification performance
based on representations from FedDVA. Details and full results are provided in
Appendix.11 and codes are also available on the
Github111https://github.com/pysleepy/FedDVA.

###  V-A Personalization Disentanglement

We empirically study FedDVA’s disentanglement capability on real-world data
sets with different personalization settings.

MNIST222https://yann.lecun.com/exdb/mnist/ is a benchmark dataset of
handwritten digits with 60,000 training images and 10,000 testing images. We
uniformly allocate them to a set of clients and synthesize them with client-
specific marks. An example is in Fig.3(a).

Figure 3: a) Digits in each quadrant denote the samples from the same client.
They are synthesized with client-specific marks, i.e., horizontal sinusoids,
ellipses, vertical sinusoids and plain digits. b) Reconstructions will vary
along with the changes in general representation $z$ (rows) while invariant to
the personalized representation $c$ (columns). Vice versa for client-specific
marks.

We train the FedDVA model by Algorithm.1 and visualize manifolds of data
reconstructed from the learned representations. Fig.3(b) shows that
representations $z$ and $c$ are disentangled. Digits in the images will vary
along with the changes in general representation $z$ (rows) while invariant to
the personalized presentation $c$ (columns). Similarly, client-specific marks
will vary along with the changes in $c$ and remain unchanged when $z$ changes.

CelebA333https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html is a large-scale
face dataset containing 202,599 face images of celebrities. We allocate them
to clients according to face attributes so that images on the same client will
demonstrate a bias towards some attributes, e.g., hairstyles. Examples of
personalized face images are shown in Fig.4(a).

Figure 4: a) The faces in each quadrant correspond to the data from the same
client. They share similar attributes regarding their hairstyles, i.e.,
wearing hats, black hair, blond, and bald. b) Images generated from the same
general representation $z$ (rows) will have similar faces and vary in
personalized attributes regarding personalized representation $c$ (columns).

We can find that major face attributes and hairstyles are disentangled. Faces
generated from the same general representation $z$ are similar and will vary
in hairstyles when personalized representation $c$ changes. Meanwhile, other
significant attributes like background colors will also vary along with $z$,
while miscellaneous attributes like face angles are implied as personalized
knowledge in $c$.

In addition, we also visualize the distribution of the representations learned
by FedDVA. We embed them into the 2-dimension space by the t-SNE and visualize
them by scatter plots (Fig.5). It can be found that distributions of the
general representation $z$ (left) are mixed and client-irrelevant. The
counterparts of the personalized representation $c$ (right) are clustered
regarding their clients.

Figure 5: Representation distributions regarding $z$ (Left) and $c$ (Right).
Each dot denotes a sample’s latent representation, and the color denotes the
client it is sampled from. It can be found that distributions of the general
representation $z$ are client irrelevant, and the counterparts of the
personalized representation $c$ are clustered regarding their clients

###  V-B Personalized Classification

We evaluate the classification performance of representations learned by
FedDVA. We tune the dual encoders along with a classification head and compare
their performance with vanilla FL algorithms FedAvg [1], FedAvg+Fine Tuning
[12] and DITTO [10]. Two personalization settings are applied. 1)
heterogeneous inputs: digits on clients are synthesized with client-specific
marks as in Fig.3(a); 2) heterogeneous outputs: we allocate samples to clients
so that they will vary in label distributions. Two benchmark datasets, MNIST
and CIFAR-10444https://www.cs.toronto.edu/ kriz/cifar.html, are applied.
Results in Fig.6 and Fig.7 show that a model based on the disentangled
representations will converge fast and achieve competitive performance to
those vanilla FL methods.

Figure 6: Accuracy on clients where digits are synthesized with marks. FedDVA
(blue) will achieve better accuracy (vertical) at the same communication round
(horizontal).

Figure 7: Accuracy on MNIST and CIFAR0-10. FedDVA will achieve higher accuracy
each communication round and have less variance regarding the accuracy among
clients (shadows).

##  VI Conclusion

In conclusion, this paper proposes a novel FedDVA method to disentangle
general and personalized representations for PFL. Empirical studies validate
FedDVA’s disentanglement capability and show that disentangled representations
will improve convergence and classification performance.

## References

  * [1] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas, “Communication-efficient learning of deep networks from decentralized data,” in _Artificial intelligence and statistics_. PMLR, 2017, pp. 1273–1282.
  * [2] T. Li, A. K. Sahu, M. Zaheer, M. Sanjabi, A. Talwalkar, and V. Smith, “Federated optimization in heterogeneous networks,” _arXiv preprint arXiv:1812.06127_ , 2018.
  * [3] P. Kairouz, H. B. McMahan, B. Avent, A. Bellet, M. Bennis, A. N. Bhagoji, K. Bonawitz, Z. Charles, G. Cormode, R. Cummings _et al._ , “Advances and open problems in federated learning,” _arXiv preprint arXiv:1912.04977_ , 2019.
  * [4] D. Rothchild, A. Panda, E. Ullah, N. Ivkin, I. Stoica, V. Braverman, J. Gonzalez, and R. Arora, “Fetchsgd: Communication-efficient federated learning with sketching,” in _International Conference on Machine Learning_. PMLR, 2020, pp. 8253–8265.
  * [5] S. Reddi, Z. Charles, M. Zaheer, Z. Garrett, K. Rush, J. Konečnỳ, S. Kumar, and H. B. McMahan, “Adaptive federated optimization,” _arXiv preprint arXiv:2003.00295_ , 2020.
  * [6] T. H. Hsu, H. Qi, and M. Brown, “Measuring the effects of non-identical data distribution for federated visual classification,” _CoRR_ , vol. abs/1909.06335, 2019. [Online]. Available: http://arxiv.org/abs/1909.06335
  * [7] Y. Mansour, M. Mohri, J. Ro, and A. T. Suresh, “Three approaches for personalization with applications to federated learning,” _arXiv preprint arXiv:2002.10619_ , 2020.
  * [8] T. Lin, L. Kong, S. U. Stich, and M. Jaggi, “Ensemble distillation for robust model fusion in federated learning,” _arXiv preprint arXiv:2006.07242_ , 2020\.
  * [9] H. Wang, M. Yurochkin, Y. Sun, D. Papailiopoulos, and Y. Khazaeni, “Federated learning with matched averaging,” _arXiv preprint arXiv:2002.06440_ , 2020\.
  * [10] T. Li, S. Hu, A. Beirami, and V. Smith, “Ditto: Fair and robust federated learning through personalization,” in _International Conference on Machine Learning_. PMLR, 2021, pp. 6357–6368.
  * [11] A. Fallah, A. Mokhtari, and A. Ozdaglar, “Personalized federated learning: A meta-learning approach,” _arXiv preprint arXiv:2002.07948_ , 2020.
  * [12] G. Cheng, K. Chadha, and J. Duchi, “Fine-tuning is fine in federated learning,” _arXiv preprint arXiv:2108.07313_ , 2021.
  * [13] Y. Bengio, A. Courville, and P. Vincent, “Representation learning: A review and new perspectives,” _IEEE transactions on pattern analysis and machine intelligence_ , vol. 35, no. 8, pp. 1798–1828, 2013.
  * [14] K. Gregor, G. Papamakarios, F. Besse, L. Buesing, and T. Weber, “Temporal difference variational auto-encoder,” _arXiv preprint arXiv:1806.03107_ , 2018.
  * [15] R. T. Chen, X. Li, R. Grosse, and D. Duvenaud, “Isolating sources of disentanglement in vaes,” in _Proceedings of the 32nd International Conference on Neural Information Processing Systems_ , 2019, pp. 2615–2625.
  * [16] F. Locatello, S. Bauer, M. Lucic, G. Raetsch, S. Gelly, B. Schölkopf, and O. Bachem, “Challenging common assumptions in the unsupervised learning of disentangled representations,” in _international conference on machine learning_. PMLR, 2019, pp. 4114–4124.
  * [17] A. Shamsian, A. Navon, E. Fetaya, and G. Chechik, “Personalized federated learning using hypernetworks,” _arXiv preprint arXiv:2103.04628_ , 2021.
  * [18] C. Finn, P. Abbeel, and S. Levine, “Model-agnostic meta-learning for fast adaptation of deep networks,” in _International Conference on Machine Learning_. PMLR, 2017, pp. 1126–1135.
  * [19] D. P. Kingma and M. Welling, “Auto-encoding variational bayes,” _arXiv preprint arXiv:1312.6114_ , 2013.

##  Appendix A Evidence Lower Bounds

###  A-A ELBO optimizing $q(z|x)$

Suppose $p(z|x)$ is the true posterior of $z$, $q(z|x)$ is the variational
posterior approximating $p(z|x)$ and samples on the same client are
independent and identical distributed (iid.), then the learning task on the
$k$-th client is to minimize $D_{KL}(q(z|x)||p(z|x))$, which is

$\displaystyle D_{KL}(q(z|x)||p(z|x))=$ $\displaystyle\int
q(z|x)\log\frac{q(z|x)}{p(z|x)}\text{d}z$ (16) $\displaystyle=$
$\displaystyle\int q(z|x)\log\frac{q(z|x)p_{k}(x)}{p(z|x)p_{k}(x)}\text{d}z$
$\displaystyle=$ $\displaystyle\log p_{k}(x)+\int
q(z|x)\log\frac{q(z|x)}{p_{k}(x|z)p(z)}\text{d}z$ $\displaystyle=$
$\displaystyle\log p_{k}(x)-\mathbb{E}_{q(z|x)}[\log
p_{k}(x|z)]+D_{KL}(q(z|x)||p(z))$

or equivalently,

$\displaystyle\log p_{k}(x)\geq ELBO_{z}(x,k)=\mathbb{E}_{q(z|x)}[\log
p_{k}(x|z)]-D_{KL}(q(z|x)||p(z))$ (17)

###  A-B ELBO optimizing $q(z|x,c)$

Suppose $p(c|x,z)$ is the true posterior of $c$, $q(c|x,z)$ is the variational
posterior approximating $p(c|x,z)$ and samples on the same client are iid.,
then the learning task on the $k$-th client is to minimize
$D_{KL}(q(c|x,z)||p(c|x,z))$, which is

$\displaystyle D_{KL}(q(c|x,z)||p(c|x,z))=$ $\displaystyle\int
q(c|x,z)\log\frac{q(c|x,z)}{p(c|x,z)}\text{d}c$ (18) $\displaystyle=$
$\displaystyle\int
q(c|x,z)\log\frac{q(c|x,z)p_{k}(x|z)}{p(c|x,z)p_{k}(x|z)}\text{d}c$
$\displaystyle=$ $\displaystyle\log p_{k}(x|z)+\int
q(c|x,z)\log\frac{q(c|x,z)}{p_{k}(x,c|z)}\text{d}c$ $\displaystyle=$
$\displaystyle\log p_{k}(x|z)-\mathbb{E}_{q(c|x,z)}[\log
p_{k}(x,c|z)]-H(q(c|x,z))$

Ideally, there is a client-irrelevant likelihood $p(x|z,c)$ modeling the
sample generating process, that is $p_{k}(x)=\iint
p(x|z,c)p(z)p_{k}(c)\text{d}z\text{d}c$, where the personality of a client
lies on $p_{k}(c)$. Then we have

$\displaystyle\log p_{k}(x)\geq ELBO_{c}(x,z,k)=\mathbb{E}_{q(c|x,z)}[\log
p(x|z,c)]-D_{KL}(q(c|x,z)||p_{k}(c))$ (19)

which is equivalent to

$\displaystyle\log p_{k}(x)\geq
ELBO^{\prime}_{c}(x,z,k)=\mathbb{E}_{q(c|x,z)}[\log
p_{k}(x|z,c)]-D_{KL}(q(c|x,z)||q(c))$ (20)

###  A-C Difference between $D_{KL}(q(c|x,z)||q(c))$ and
$D_{KL}(q(c|x,z)||\bar{p}_{k}(c))$

For samples on the same client and suppose they are independent and identical
distributed, according to Eq.5, there is

$\displaystyle\mathbb{E}_{p_{k}(x)}[\mathbb{E}_{q(z|x)}[D_{KL}(q(c|x,z)||q(c))-D_{KL}(q(c|x,z)||\bar{p}_{k}(c))]]$
(21) $\displaystyle=$
$\displaystyle\mathbb{E}_{p_{k}(x)}[\mathbb{E}_{q(z|x)}[\mathbb{E}_{q(c|x,z)}[\log\bar{p}_{k}(c)-\log
q(c)]]]$ $\displaystyle=$
$\displaystyle\mathbb{E}_{\bar{p}_{k}(c)}[\log\bar{p}_{k}(c)-\log q(c)]$
$\displaystyle=$ $\displaystyle D_{KL}(\bar{p}_{k}(c)||q(c))$
$\displaystyle\geq$ $\displaystyle\xi_{k}$

##  Appendix B Computation of the $KL$-Divergence

###  B-A $KL$-Divergence between two Gaussian distributions

For the $i$-th sample $x_{i}$ and $j$-th sample $x_{j}$, variational
posteriors inferring representation $c$ are
$q(c|x_{i},z_{i})=\mathcal{N}(c;\mu_{i},\Sigma_{i})$ and
$q(c|x_{j},z_{j})=\mathcal{N}(c;\mu_{j},\Sigma_{j})$, where $c$ is a
$d$-dimensional vector and convariance matrices of $\Sigma_{i}$ and
$\Sigma_{j}$ are diagonal. Then we have

$\displaystyle\int q(c|x_{i},z_{i})\log q(c|x_{i},z_{i})\text{d}c=$
$\displaystyle\int\mathcal{N}(c;\mu_{i},\Sigma_{i})\log\mathcal{N}(c;\mu_{i},\Sigma_{i})\text{d}c$
(22) $\displaystyle=$
$\displaystyle-\frac{1}{2}(d\log(2\pi)+\log|\Sigma_{i}|+d)$

and

$\displaystyle\int q(c|x_{i},z_{i})\log q(c|x_{j},z_{j})\text{d}c=$
$\displaystyle\int\mathcal{N}(c;\mu_{i},\Sigma_{i})\log\mathcal{N}(c;\mu_{j},\Sigma_{j})\text{d}c$
(23) $\displaystyle=$
$\displaystyle-\frac{1}{2}(d\log(2\pi)+\log|\Sigma_{j}|+Tr(\Sigma_{j}^{-1}\Sigma_{i})+(\mu_{i}-\mu_{j})^{T}\Sigma_{j}^{-1}(\mu_{i}-\mu_{j}))$

Combining Eq.22 and Eq.23, the $KL$-Divergence between $q(c|x_{i},z_{i})$ and
$q(c|x_{j},z_{j})$ is

$\displaystyle D_{KL}(q(c|x_{i},z_{i})||q(c|x_{j},z_{j}))=$ $\displaystyle\int
q(c|x_{i},z_{i})(\log q(c|x_{i},z_{i})-\log q(c|x_{j},z_{j}))\text{d}c$ (24)
$\displaystyle=$
$\displaystyle\int\mathcal{N}(c;\mu_{i},\Sigma_{i})(\log\mathcal{N}(c;\mu_{i},\Sigma_{i}))\text{d}c-\int\mathcal{N}(c;\mu_{i},\Sigma_{i})(\log\mathcal{N}(c;\mu_{j},\Sigma_{j}))\text{d}c$
$\displaystyle=$
$\displaystyle\frac{1}{2}[(\mu_{i}-\mu_{j})^{T}\Sigma_{j}^{-1}(\mu_{i}-\mu_{j})-\log|\Sigma_{j}^{-1}\Sigma_{i}|+Tr(\Sigma_{j}^{-1}\Sigma_{i})-d]$
$\displaystyle=$
$\displaystyle\frac{1}{2}\sum_{l=1}^{d}[(\frac{\mu_{i}^{(l)}-\mu_{j}^{(l)}}{\sigma_{j}^{(l)}})^{2}-\log(\frac{\sigma_{i}^{(l)}}{\sigma_{j}^{(l)}})^{2}+(\frac{\sigma_{i}^{(l)}}{\sigma_{j}^{(l)}})^{2}-1]$

where $l$ denotes the $l$-th element and $\sigma_{i}^{(l)}$ denotes the
positive root of the $l$-th element on the diagonal of covariance matrix
$\Sigma_{i}$.

###  B-B Computation of $D_{KL}(q(c|x,z)||\bar{p}_{k}(c))$

Let $x_{i}$ and $x_{j}$ denotes the $i$-th and $j$-th sample in dataset
$\mathcal{D}_{k}$ with size is $n_{k}$

$\displaystyle D_{KL}(q(c|x_{i},z_{i})||\bar{p}_{k}(c))=$
$\displaystyle\mathbb{E}_{q(c|x_{i},z_{i})}[\log
q(c|x_{i},z_{i})-\log\frac{1}{n_{k}}\sum_{j=1}^{n_{k}}[q(c|x_{j},z_{j})]$ (25)
$\displaystyle\leq$ $\displaystyle\mathbb{E}_{q(c|x_{i},z_{i})}[\log
q(c|x_{i},z_{i})-\frac{1}{n_{k}}\sum_{j=1}^{n_{k}}\log q(c|x_{j},z_{j})]$
$\displaystyle=$
$\displaystyle\frac{1}{n_{k}}\sum_{j=1}^{n_{k}}\mathbb{E}_{q(c|x_{i},z_{i})}[\log
q(c|x_{i},z_{i})-\log q(c|x_{j},z_{j})]$

Bringing Eq.24 we have

$\displaystyle
D_{KL}(q(c|x_{i},z_{i})||\bar{p}_{k}(c))\leq\frac{1}{2n_{k}}\sum_{j=1}^{n_{k}}\sum_{l=1}^{d}[(\frac{\mu_{i}^{(l)}-\mu_{j}^{(l)}}{\sigma_{j}^{(l)}})^{2}-\log(\frac{\sigma_{i}^{(l)}}{\sigma_{j}^{(l)}})^{2}+(\frac{\sigma_{i}^{(l)}}{\sigma_{j}^{(l)}})^{2}-1]$
(26)

where $\mu_{i}$, $\sigma_{i}$ and $\mu_{j}$, $\sigma_{j}$ are outputs of
neural networks and they can be differentiated and optimized by gradient based
optimization methods.

##  Appendix C Experiments

###  C-A Personalization settings

We follow the work in [6] to allocate samples to 20 clients, with each client
having random fractions regarding classes. As described in Fig.8, each column
denotes fractions of classes on a client, and each color corresponds to a
class. The longer a bar is, the more significant the fraction of the class is
on that client.

Figure 8: Class distributions

###  C-B Model architecture

FedDVA: The encoder for representation $z$ consists of a 4-layer CNN backbone
and two fully connected embedding layers; The encoder for representation $c$
first combines $z$ and $x$ with an FC layer and then forwards the embedding of
$[x;z]$ through a 4-layer CNN backbone and two fully connected embedding
layers; The decoder is the reverse of the encoding modules. The model
architecture is illustrated in Fig.9, and codes are uploaded along with the
supplementary materials.

Figure 9: Model architecture

CNN: To classify samples from MNIST and CIFAR-10, we implemented a CNN
consisting of a 4-layer CNN backbone and a 2-FC layer classification head.

###  C-C Hyperparameters

On all clients, the batch size is 256 and the learning rate is fixed as 0.001.
We trained the global model by 200 communication rounds and 5 epochs during
each round. The dimensions of $z$ and $c$ are set to be 4 respectively for
reconstruction tasks and 8 for classification tasks. $\xi_{k}$ is set to be 8
times the dimension of $c$. $\alpha$ is 1 and $\beta$ is 0.75.

###  C-D Manifolds of reconstructions

Figure 10: Reconstructions of MNIST. Each quadrant corresponds to a client.
General representations $z$ vary along the vertical axis and personalized
representations $c$ vary along the horizontal axis. We can see that $z$ and
$c$ are disentangled as digits vary along with changes in $z$ and personalized
marks changes with $c$.

Figure 11: Reconstructions of CIFAR-10. Each quadrant corresponds to a client.
General representations $z$ vary along the vertical axis and personalized
representations $c$ vary along the horizontal axis. We can see that $z$ and
$c$ are disentangled as faces vary along with changes in $z$ and personalized
attributes like hairstyles and skins change with $c$.