Disentangling data distribution for Federated Learning

## Abstract

Abstract Federated Learning (FL) facilitates collaborative training of a global model whose performance is boosted by private data owned by distributed clients, without compromising data privacy. Yet the wide applicability of FL is hindered by entanglement of data distributions across different clients.
This paper demonstrates for the first time that by disentangling data distributions FL can in principle achieve efficiencies comparable to those of distributed systems, requiring only one round of communication.
To this end, we propose a novel FedDistr algorithm, which employs stable diffusion models to decouple and recover data distributions. Empirical results on the CIFAR100 and DomainNet datasets show that FedDistr significantly enhances model utility and efficiency in both disentangled and near-disentangled scenarios while ensuring privacy, outperforming traditional federated learning methods.

## 1 Introduction

Figure: (a) Disentangled
Refer to caption: extracted/5919221/imgs/framework/intro_a.png

Despite the extensive research in Federated Learning (FL), its practical application remains limited. A key challenge is to achieve high efficiency while preserving both model utility and privacy (McMahan et al., ; Kairouz et al., ).
There is a consensus that this inefficiency stems from the entanglement of data distribution across clients, where many rounds of communications are required to ensure the convergence of the global model (Zhao et al., ; X. Li et al., ; Tian et al., ). The opposite of entanglement implies that client data distributions at different dimensions are
disentangled as shown in Figure 1(a).

We believe that the ideal federated learning algorithm can achieve efficiency levels comparable to those of ideal distributed systems that attain full parallelism (Sunderam, ), provided that the data distributions across clients can be entirely disentangled, as illustrated in Fig. 1(a) (see details in Def. 2). As shown by Theorem 1, it is actually a sufficient condition for achieving ideal efficiency of federated learning with only one round of communication being needed. Moreover, this ideal condition is often approximately fulfilled in practice e.g., when millions of mobile device clients participate in federated learning and most client data distributions are different from each other (Sattler et al., ; Tian et al., ). In other words, under such a near-disentangled condition (see Def. 2 and Fig. 1(b)), there exists a federated learning algorithm capable of achieving global model utility with only one communication round, while ensuring that the utility loss remains within a tolerable threshold (see Theorem 2).

In order to take full advantage of the described disentangled and near–disentangled cases (see Fig. 1), we propose to first disentangle data distributions into distinct components such that each client can launch their respective learning task independently. This disentanglement will lead to significant reduction of required communication rounds as revealed by Theorem 1. Technology-wise, there is a large variety of methods that can be utilized to decompose data distribution, ranging from subspace decomposition (Abdi and Williams, ; Von Luxburg, ) to dictionary learning (Tošić and Frossard, ) etc. (see Section 2.2 for detailed review).
In this work, we propose an algorithm, called FedDistr, which leverages the stable diffusion model technique (Croitoru et al., ) due to its robust ability to extract and generate data distributions effectively. In our approach, clients locally extract data distributions via stable diffusion model, and then upload these decoupled distributions to the server. The server actively identifies the orthogonal or parallel between the base distributions uploaded by clients and aggregate the orthogonal distribution once.

Previous work (Zhang et al., ) has shown that it is impossible to achieve the optimal results among the utility, privacy and efficiency. The proposed FedDistr offers a superior balance between model utility and efficiency under the disentangled and near-disentangled scenarios, while still ensuring the privacy-preserving federated learning: (1) By decoupling client data distributions into the different base distributions, the server actively aligns the base distribution across different client, while FedAvg does not perform this decoupling, leading to a decline in global performance, especially for the disentangled case; (2) The proposed FedDistr requires only one round of communication, and the amount of transmitted distribution parameters is much smaller than that of model gradients; (3) The proposed FedDistr transmits a minimal amount of data distribution parameters, thereby mitigating the risk of individual data privacy leakage to some extent. Furthermore, privacy mechanisms such as differential privacy (DP) can be integrated into FedDistr, offering additional protection for data privacy.

## 2 Related Work

### 2.1 Communication Efficient Federated Learning

One of the earliest approaches to reducing communication overhead is the FedAvg algorithm (McMahan et al., ). FedAvg enables multiple local updates at each client before averaging the models at the server, thereby significantly decreasing the communication frequency. To further mitigate communication costs, compression techniques such as quantization and sparsification (Reisizadeh et al., ) and client selection strategies (Lian et al., ; Liu et al., ) have been implemented in federated learning. However, all of these methods are less effective in Non-IID scenarios, where the model accuracy is significantly impacted.

To address the Non-IID problem, numerous methods have been proposed (T. Li et al., ; Arivazhagan et al., ), which introduce constraints on local training to ensure that models trained on heterogeneous data do not diverge excessively from the global model. However, most of these methods require numerous communication rounds, which is impractical, particularly in wide-area network settings.

Furthermore, some methods propose a one-shot federated learning approach (Guha et al., ; Q. Li et al., ), which requires only a single communication round by leveraging techniques such as knowledge distillation or consistent voting. However, significant challenges arise in this context under Non-IID settings (Diao et al., ).

### 2.2 Distribution Generation

This paper focuses on transferring distribution instead of models. There are several distribution generation methods:
Parametric Methods (Reynolds et al., ) assume that the data follows a specific distribution with parameters that can be estimated from the data. Gaussian Mixture Model (GMM) represents a distribution as a weighted sum of multiple Gaussian components.
Generative Models (Goodfellow et al., ; Croitoru et al., )
aim to learn the underlying distribution of the data so that new samples can be generated from the learned distribution. For example, GANs Goodfellow et al. (); Karras et al. () consist of two networks, a generator $G$ and a discriminator $D$.
The generator learns to produce data similar to the training data, while the discriminator tries to distinguish between real and generated data. Principal Component Analysis (PCA) (Abdi and Williams, ) is a linear method for reducing dimensionality and capturing the directions of maximum variance. It decomposes the data covariance matrix $\Sigma$ into eigenvectors and eigenvalues. Dictionary Learning and Sparse Coding (Tošić and Frossard, ) seeks to represent data as a sparse linear combination of basis vectors (dictionary atoms).

**Table 1: Table of Notation**
| Notation | Meaning |
| --- | --- |
| $\mathcal{D}_{k}$ | Dataset of Client $k$ |
| $K$ | the number of clients |
| $\mathcal{D}$ | $\mathcal{D}_{1}\cup\cdot\cup\mathcal{D}_{K}$ |
| $S_{k}$ and $S$ | the distribution of $\mathcal{D}_{k}$ and $\mathcal{D}$ |
| $\epsilon_{u}$ | Utility loss in Eq. (4) |
| $F_{\omega}$ | the federated model |
| $\{P_{i}\}_{i=1}^{m}$ | $m$ independent sub-distributions |
| $\vec{\pi}_{k}=(\pi_{1,k},\cdots,\pi_{m,k})$ | the probability vector for client $k$ on $\{P_{i}\}_{i=1}^{m}$ |
| $\xi$ | entangled coefficient |
| $s_{k_{1},k_{2}}$ | Entangled coefficient between client $k_{1}$ and $k_{2}$ |
| $p$ | learnable prompt embedding |
| $T_{c}$ | communication rounds |

## 3 Framework

In this section, we first formulate the problem via distribution transferring and then provide the analysis on what conditions clients can directly communicate once through a distribution entanglement perspective.

### 3.1 Problem Formulation

Consider a Horizontal Federated Learning (HFL) consisting of $K$ clients who collaboratively train a HFL model $\omega$ to optimize the following objective:

$$ $\min_{\omega}\sum_{k=1}^{K}\sum_{i=1}^{n_{k}}\frac{f(\omega;(x_{k,i},y_{k,i})) }{n_{1}+\cdots+n_{K}},$ (1) $$

where $f$ is the loss, e.g., the cross-entropy loss, $\mathcal{D}_{k}=\{(x_{k,i},y_{k,i})\}_{i=1}^{n_{k}}$ is the dataset with size $n_{k}$ owned by client $k$. Denote $\mathcal{D}=\mathcal{D}_{1}\cup\cdots\cup\mathcal{D}_{K}$ and $\mathcal{D}$ follows the distribution $S$. The goal of the federated learning is to approach the centralized training with data $\mathcal{D}$ as:

$$ $\min_{\omega}\mathbb{E}_{(x,y)\in S}f(\omega;(x,y)).$ (2) $$

Numerous studies (Zhao et al., ; X. Li et al., ; Tian et al., ) have demonstrated that Eq. (1) converges to Eq. (2) under IID settings. However, significant discrepancies arise in Non-IID scenarios, leading to an increased number of communication rounds and a decline in model utility for Eq. (1). While data heterogeneity in extreme Non-IID settings was considered the main cause for FL algorithms having to compromise model performances for reduced rounds of communication, this paper discloses that
direct optimization of Eq. (2) can actually be achieved at the cost of a single round of communication. Both theoretical analysis and empirical studies show that data heterogeneity is a blessing rather than a curse, as long as data distributions among different clients can be completely disentangled (see Theorem 1). Moreover, we aim to learn the distribution $S$ in a single communication by transferring either the distribution or its parameters, thereby mitigating the risk of individual data privacy leakage (Xiao and Devadas, ) as follows:

$$ $\min_{S}\sum_{k=1}^{K}\lambda_{k}\tilde{f}(S;\mathcal{D}_{k}),\quad\text{s.t. }T_{c}=1$ (3) $$

where $\tilde{f}$ represents the loss function used to evaluate whether $S$ accurately describes the dataset $\mathcal{D}_{k}$ (^1^11For various distribution estimation methods, refer to Section 2.2), $T_{c}$ is the communication round, $\lambda_{k}$ denotes the coefficient for client $k$, and it is required that $\sum_{k=1}^{K}\lambda_{k}=1$.

In this framework, all clients first collaborate to estimate the distribution $S$ of the total dataset $\mathcal{D}=\mathcal{D}_{1}\cup\cdots\cup\mathcal{D}_{K}$.
The following section discusses the conditions under which clients can directly interact once to estimate $S$ through a distribution entangled perspective.

We also define the utility loss for the estimated distribution $\hat{S}$ as:

$$ $\epsilon_{u}=\mathbb{E}_{(x,y)\in S}f(\hat{\omega}^{*};(x,y))-\mathbb{E}_{(x,y )\in S}f(\omega^{*};(x,y)),$ (4) $$

where $\omega^{*}$ is the optimal parameter of Eq. (2) and $\hat{\omega}^{*}=argmin_{\omega}\mathbb{E}_{(x,y)\in\hat{S}}f(\omega^{*};(x,y))$.

Threat Model. We assume the server might be semi-honest adversaries such that they do not submit any malformed messages but may launch privacy attacks on exchanged information from other clients to infer clients’ private data.

### 3.2 Distribution Entanglement Analysis

#### 3.2.1 ξ 𝜉 \xi italic_ξ -entangled

Figure: (a) $\xi=0.11$
Refer to caption: extracted/5919221/imgs/framework/entangled.png

Assume the data $\mathcal{D}$ is drawn from $m$ distinct base distributions. Therefore, we can decompose the complex distribution $S$ into simpler components, each representing a portion of the overall distribution (Reynolds et al., ; Abdi and Williams, ):

$$ $S=\sum_{i=1}^{m}\pi_{i}P_{i}.$ (5) $$

where $P_{i}$ represents the $i$-th base distribution (e.g., DomaiNet has the data with different label and domain, the base distribution represents the distribution followed by the data with a single label and domain),
$\pi_{i}$ is the weight or probability assigned to the $i$-th base distribution, such that $\sum_{i=1}^{m}\pi_{i}=1$ and $\pi_{i}>0$. In the case of Gaussian Mixture Models (GMMs) (Reynolds et al., ), $P_{i}$ is the Gaussian distribution $\mathcal{N}(X|\mu_{i},\Sigma_{i})$. For each client’s dataset $\mathcal{D}_{k}$ following a distribution $S_{k}$ can be represented as:

$$ $S_{k}=\sum_{i=1}^{m}\pi_{i,k}P_{i},$ (6) $$

where $\pi_{i,k}$ is the weight or probability assigned to the $i$-th base distribution, such that $\sum_{i\in\mathcal{C}_{k}}\pi_{i,k}=1$ and $\pi_{i,k}\geq 0$. Noted that $\pi_{i,k}=0$ if client $k$ doesn’t have the distribution $S_{i}$.
We define the Entangled coefficient between two clients based on the probability $\pi_{i,k}$:

###### Definition 1 .

Let $\vec{\pi}_{k}=(\pi_{1,k},\cdots,\pi_{m,k})$ denote the probability vector for client $k$.
The Entangled coefficient $s_{k_{1},k_{2}}$ between client $k_{1}$ and $k_{2}$ is defined as:

$$ $s_{k_{1},k_{2}}=\frac{<\vec{\pi}_{k_{1}},\vec{\pi}_{k_{2}}>}{\|\vec{\pi}_{k_{1 }}\|_{2}\|\vec{\pi}_{k_{2}}\|_{2}}$ (7) $$

where $<a,b>$ represents the inner product of two vectors and $\|\cdot\|_{2}$ is the $\ell_{2}$ norm.

Definition 1 quantifies the overlap in data distribution between clients by identifying the difference between probability on the base distribution. Obviously, $0\leq s_{k_{1},k_{2}}\leq 1$. We define two distributions to be orthogonal if $s_{k_{1},k_{2}}=0$.

According to the Entangled coefficient, we define $\xi$-entangled across clients’ data distribution in federated learning according to the entangled coefficient.

###### Definition 2 .

We define the distributions across $K$ clients as $\xi$-entangled if the Entangled coefficient $s_{k_{1},k_{2}}\leq\xi$ for any two distinct clients $k_{1},k_{2}\in[K]$. Moreover, when $\xi=0$, we define the distribution across $K$ clients to be disentangled.

It is important to note that Entangled coefficient is closely related to the degree of Non-IID data in federated learning (FL) (X. Li et al., ). Specifically, under IID conditions, where each client follows the same independent and identical distribution, this implies that $\pi_{i,k_{1}}=\pi_{i,k_{2}}>0$ for any $k_{1},k_{2}$ in Eq. (6). Consequently, $\xi$ equals $1$. In contrast, under extreme Non-IID conditions, where each client has entirely distinct data distributions (e.g., client $k_{1}$ possesses a dataset consisting solely of ’cat’ images, whereas client $k_{2}$ possesses a dataset exclusively of ’dog’ images), the Entanglement Coefficient, $\xi$, equates to zero.

###### Definition 1 .

###### Definition 2 .

#### 3.2.2 Analysis on Disentangled and Near-disentangled case

Numerous studies, as referenced in the literature (Kairouz et al., ; X. Li et al., ), have demonstrated that the performance of the Federated Averaging (FedAvg) algorithm (McMahan et al., ) is significantly influenced under disentangled conditions. It will be demonstrated herein that there exists an algorithm capable of preserving model utility and efficiency within the disentangled scenario, with the distinct advantage of necessitating merely a single communication round.

###### Theorem 1 .

If $f$ is $L$-lipschitz, a data distribution across clients being disentangled is a sufficient condition for the existence of a privacy-preserving federated algorithm that requires only a single communication round and achieves a utility loss of less than $\epsilon$ with a probability of at least $1-\exp(-\frac{\min\{n_{1},\cdots,n_{K}\}\epsilon^{2})}{2L^{2}}$, i.e.,

$$ $Pr(\epsilon_{u}\leq\epsilon)\geq 1-\exp(-\frac{\min\{n_{1},\cdots,n_{K}\} \epsilon^{2})}{2L^{2}}$ (8) $$

Theorem 1 demonstrates that if the data distribution across clients is disentangled, it is possible to achieve an expected loss of $\epsilon$ with a probability of $1-\exp(-\frac{\min\{n_{1},\cdots,n_{K}\}\epsilon^{2})}{2L^{2}}$
with only a single communication round. It is noteworthy that as the number of datasets $n_{k}$ or the expected loss tends to infinity, this probability approaches one. The implication of Theorem 1 is that when the distributions of individual clients are disentangled, clients can effectively transfer their distributions to the server. Subsequently, the server can aggregate the disentangled distributions uploaded by the clients in a single step to obtain the overall dataset distribution $\mathcal{D}$. Additionally, we provide an analysis of the small $\xi$ condition, referred to as near-disentangled.

###### Theorem 2 .

When the distributions of across $K$ clients satisfy near-disentangled condition, specifically, $\xi<\frac{1}{(K-1)^{2}}$, if $f$ is $L$-lipschitz, then there exists a privacy-preserving federated algorithm that requires only one communication round and achieves a utility loss of less than $\epsilon$ with a probability of at least $1-\exp(-\frac{(1-(K-1)\sqrt{\xi})n\epsilon^{2}}{2mL^{2}})$, i.e.,

$$ $Pr(\epsilon_{u}\leq\epsilon)\geq 1-\exp(-\frac{(1-(K-1)\sqrt{\xi})n\epsilon^{2 }}{2mL^{2}})$ (9) $$

Theorem 2 demonstrates that, within the framework of the near-disentangled scenario, the probability of achieving an expected loss of $\epsilon$ is given by
$1-\exp(-\frac{(1-(K-1)\sqrt{\xi})n\epsilon^{2}}{2mL^{2}})$.
Specifically, as the entanglement coefficient $\xi$ increases or data number $n$ decreases, the probability decreases (see proof in Appendix).

###### Theorem 1 .

###### Theorem 2 .

## 4 The Proposed Algorithm

This section introduces the proposed algorithm, called FedDistr, achieving one communication round while maintaining the model utility by leveraging the Latent Diffusion Model (LDM) Ho and Salimans () due to its fast inference speed and the wide availability of open-source model parameters. FedDistr is mainly divided into the following three steps (see Fig. 3 and Algo. 1):

Figure: Figure 3: Overview of the proposed algorithm FedDistr.
Refer to caption: x1.png

### 4.1 Distribution Disentangling

The first step is to disentangle the local data distribution into several base distribution for each client by leveraging the autoencoder part of the LDM. The autoencoder Van Den Oord et al. (); Agustsson et al. () consists of an encoder, $E$, and a decoder, $\mathcal{D}$. The encoder, $E$, maps input images $x$ to a spatial latent code, denoted as $z=E(x)$, while the decoder $\mathcal{D}$ is trained to perform the inverse mapping of the encoder, achieving an approximation $x\approx\mathcal{D}(E(x))$.

Specifically, each client $k$ initially encodes the private data into the latent feature embeddings as:

$$ $z_{k,i}=E(x_{k,i}),x_{k,i}\in\mathcal{D}_{k}$ (10) $$

To disentangle the data distribution for client $k$, the data is segregated based on these latent feature embeddings. Furthermore, client $k$ performs clustering on the set ${z_{k,i}}_{i=1}^{n_{k}}$ to derive $m_{k}$ clusters by optimizing the objective function:

$$ $\min\sum_{z_{k,i}\in\mathcal{C}_{j}}\sum_{j=1}^{m_{k}}\|z_{k,i}-\bar{e}_{j}\|,$ (11) $$

where $\mathcal{C}_{j}$ represents the set of points assigned to cluster $j$ and $\bar{e}_{j}$ is the centroid (mean) of cluster $j$. Thus, $\{z_{k,i}\}_{i=1}^{n_{k}}$ is separated into $m_{k}$ datasets as $\{B_{k,i}\}_{i=1}^{m_{k}}$. According to the separation result based on the clustering, we learn the disentangled distribution in the next section for each separated dataset.

### 4.2 Distribution Generation

To estimate the base distribution $P$, each client learns their distribution parameters $v(P)$ by leveraging the embedding inversion technique in the diffusion model of LDM Ho and Salimans (); Liang et al. (). Specifically, a fixed prompt, such as “a photo of”, is encoded using a frozen text encoder into a word embedding $p$. A learnable distribution parameter (prompt embedding) $v$ is then randomly initialized and concatenated with $p$ to form the guiding condition [$p$; $v$]. This combined condition is employed to compute the LDM’s loss function. Our optimization goal to learn the distribution parameter $v_{k,i}$ of $i_{th}$ base distribution for client $k$ is formulated as:

$$ $\begin{split}v_{k,i}=&\mathop{\arg\min}\limits_{v}\mathbb{E}_{z\sim E(B_{k,i}) ,p,\epsilon\sim\mathcal{N}(0,1),t}[\|\epsilon-\\ &\epsilon_{\theta}(\sqrt{\alpha_{t}}z+\sqrt{1-\alpha_{t}}\epsilon,t,[p;v])\|_{ 2}^{2}],\\ \end{split}$ (12) $$

where $z$ is the latent code of the input image $B_{k,i}$ generated by the encoder $E$, $\epsilon$ refers to noise sampled from the standard normal distribution $\mathcal{N}(0,1)$, $t$ denotes the timestep in the diffusion process, $\alpha_{t}$ is a hyperparameter related to $t$ and $\epsilon_{\theta}$ is the model that predicts the noise. Therefore, we derive the $m_{k}$ data representations $v_{k,i}$ by solving the optimization problem in Eq. (12). These representations are then fed into the diffusion model to generate the corresponding datasets for each disentangled base distribution $i$. In the federated learning setting, each client uploads their respective set of representations ${v_{k,i}}_{i=1}^{m_{k}}$ and corresponding data number to the server.

### 4.3 Active Distribution Alignment

After each client $k$ uploads the distribution parameters $\{v_{k,i}\}_{i=1}^{m_{k}}$, the server first distinguishes between orthogonal and parallel data base distributions for different clients based on these parameters. Specifically, when the distance between distribution parameters, such as $\|v_{k_{1},i}-v_{k_{2},j}\|_{2}^{2}$, is small, it indicates that clients $k_{1}$ and $k_{2}$ share the same data base distribution.

Moreover, since the server does not have access to the sequence of the distribution parameters, matching these parameters across different clients presents a challenge. This matching problem can be framed as an assignment problem in a bipartite graph Dulmage and Mendelsohn (). To address this assignment problem, we utilize the Kuhn-Munkres algorithm (KM) Zhu et al. (), which is designed to find the maximum-weight matching in a weighted bipartite graph, or equivalently, to minimize the assignment cost. Specifically, the goal is to optimize the following expression for matching prompt embeddings between client $1$ and $k$:

$$ $\min\sum_{i_{1}=1}^{m_{k_{1}}}\sum_{i_{2}=1}^{m_{k_{2}}}\|v_{k_{1},i_{1}}-v_{k _{2},i_{2}}\|e^{k_{1},k_{2}}_{i_{1},i_{2}}$ (13) $$

where $e^{k_{1},k_{2}}_{i_{1},i_{2}}$ represents the binary assignment variable, which equals $1$ if cluster $i_{1}$ is assigned to cluster $i_{2}$, and 0 0 otherwise. Therefore, the server obtains multiple parallel sets $\mathcal{I}_{p}=\{(k,i)|e^{k_{1},k_{2}}_{i_{1},i_{2}}=1,k\in[K],i\in[m_{k}]\}$. The orthogonal sets $\mathcal{I}_{o}=\{(k,i)|k\in[K],i\in[m_{k}]\}-\mathcal{I}_{p}$ consist of other pairs that are not part of the parallel sets.

Next, for the orthogonal sets, the server simply concatenates the parameters as follows
$v_{o}=\bigcup_{k\in\mathcal{I}_{o}}v_{k,i}$.

For the parallel sets, the server selects one of the dominant distribution parameters (trained with the data number is the maximum) to represent all parameters in that set, which only requires a single communication round.

Finally, the server distributes the consolidated distribution parameters $v=[v_{p},v_{o}]$ to all clients.

Figure: Algorithm 1 FedDistr

## 5 Experiment Results

In this section, We present empirical studies to compare FedDistr with existing methods on utility, privacy and communication efficiency. Due to the page limit, please see the ablation study on more clients and large $\xi$ in Appendix B.

### 5.1 Experimental settings

Models & Datasets.
We conduct experiments on two datasets:
CIFAR100 has the 20 superclass and each superclass has 5 subclass (Krizhevsky et al., ), thus, total 100 subclass; DomainNet (Z. Wu et al., ) has the 345 superclass (label) and each superclass has 5 subclass (style), thus, total 1725 subclass. We adopt ResNet (LeCun et al., ) for conducting the classification task to distinguish the superclass(^2^22each client only know the label of the superclass but doesn’t know the label of subclass) on CIFAR100 and DomainNet. Please see details in Appendix A.

Federated Learning Settings.
We simulate a horizontal federated learning system with K = 5, 10, 20 clients in a stand-alone machine with 8 Tesla V100-SXM2 32 GB GPUs and 72 cores of Intel(R) Xeon(R) Gold 61xx CPUs. For DomainNet and CIFAR10, we regard each subclass follow one sub-distribution ($P_{i}$). For the disentangled extent, we choose the averaged entangled coefficient $s=\frac{2}{K(K-1)}\sum_{k_{1},k_{2}}d_{k_{1},k_{2}}$ over all clients as $0,0.003,0.057$. The detailed experimental hyper-parameters are listed in Appendix A.

Baseline Methods.
Four existing methods i.e., FedAvg (McMahan et al., ): FedProx (T. Li et al., ), SCAFFOLD (Karimireddy et al., ), MOON (Q. Li et al., ) and the proposed method FedDistr are compared in terms of following metrics.

Evaluation metric.
We use the (1 - model accuracy) to represent the utility loss, the rounds of required communication and the number of transmitted parameters to represent communication consumption, and the privacy budget in (Abadi et al., ) to represent the privacy leakage.

### 5.2 Comparison with Other Methods

We first evaluate the tradeoff between utility loss and communication rounds in Fig. 4. Our analysis leads to two conclusions: 1) The proposed FedDistr achieves the best tradeoff between utility loss and communication rounds compared to other methods across various entangled cases ($\xi=0,0.003,0.057$). For instance, in the disentangled case, FedDistr achieves a 40% utility loss with only one communication round, whereas MOON incurs approximately a 60% utility loss while requiring 100 communication rounds on CIFAR-100. 2) As the entanglement coefficient increases, other methods, such as FedProx, demonstrate improved performance, while FedDistr remains stable, consistently achieving a 40% utility loss on CIFAR-100 across different entangled scenarios.

Figure: (a) Disentangled for CIFAR100
Refer to caption: extracted/5919221/imgs/exp/MainExp/cifar100-e-noniid.png

Figure: (a) Disentangled
Refer to caption: extracted/5919221/imgs/exp/DPExp/Domiannet-DP-enoniid.png

Moreover, we add random noises (Abadi et al., ) to protect transmitted parameters to evaluate the tradeoff between privacy leakage and utility loss in Fig. 5
. From this analysis, we can draw the following two conclusions: 1) The proposed FedDistr achieves the best tradeoff between utility loss and privacy leakage; for instance, FedDistr (red line) is positioned in the lower-left region compared to other methods on both DomainNet and CIFAR-100. 2) As the noise level (privacy leakage) increases, the utility loss of different methods also increases.

Finally, we evaluate communication consumption, encompassing both the number of communication rounds and the amount of transmitted parameters, for five existing methods alongside our proposed method under disentangled and various $\xi$-entangled scenarios. The results summarized in Table 2 lead us to the following two conclusions: 1) Under near-disentangled or disentangled settings, the communication rounds required by FedDistr is only one, while FedAvg necessitates 80 rounds; 2) The number of transmitted parameters for the proposed FedDistr is merely 30K, in contrast to the 11.7M required by FedAvg for model transmission.

**Table 2: Comparison on communication consumption until convergence (including communication rounds and number of the transmission parameters) for different methods under different $\xi$-entangled scenario.**
| $\xi$ | Method | Communication rounds | Transmission parameters |
| --- | --- | --- | --- |
| 0 (Disentangled) | FedAvg | 150 | 11.7M |
| FedDistr | 1 | 30K |  |
| 0.003 | FedAvg | 100 | 11.7M |
| FedDistr | 1 | 30K |  |
| 0.057 | FedAvg | 80 | 11.7M |
| FedDistr | 1 | 30K |  |

### 5.3 Visualization for distribution Disentangling

We also present the distribution disentangling result (clustering $z_{k,i}$ illustrated in Sect. 4.1) in Fig. 6. It shows that clustering on encoded embedding can separate the data with different base distribution well. Therefore, our disentangling method has a good effectiveness.

Figure: (a) Flip Flops
Refer to caption: extracted/5919221/imgs/ClipVisual/flip_flops_clip_visual.png

## 6 Discussion and Conclusion

In this study, we have addressed the critical inefficiencies inherent in Federated Learning (FL) due to the entanglement of client data distributions. By analyzing the distribution entanglement, we demonstrated that achieving a disentangled data structure significantly enhances both model utility and communication efficiency. Our theoretical analysis shows that under near-disentangled conditions, FL can achieve optimal performance with a single communication round, thus aligning closely with the efficiency of traditional distributed systems.

Furthermore, we propose the FedDistr algorithm by leveraging the diffusion model. The integration of stable diffusion models for data decoupling proves to be a robust solution, paving the way for future explorations in privacy-preserving machine learning. Ultimately, this work contributes to the advancement of FL by providing a clear pathway toward more efficient and effective federated learning.

With advancements in computational capabilities during the era of large language models (LLMs) (Kasneci et al., ), the time consumption for local training has significantly decreased. Consequently, our focus is on enhancing communication efficiency. Moreover, the communication efficiency is especially important particularly in Wide Area Network scenarios (McMahan et al., ; Palmieri and Pardi, ).

Whether transferring data distribution leaks privacy is also an intriguing problem. Some studies (Xiao and Devadas, ) regard the underlying data distribution as a non-sensitive attribute, assuming that sharing insights into the distribution of data across users does not compromise individual privacy. However, other work indicates that using generative models to estimate the distribution still poses a risk of privacy leakage (Y. Wu et al., ).

## Appendix A Experimental Setting

This section provides detailed information on our experimental settings.

### A.1 Dataset & Model Architectures

Models & Datasets.
We conduct experiments on two datasets: CIFAR100 Krizhevsky et al. () and DomainNet Z. Wu et al. (). For CIFAR100, we select 10 out of 20 superclasses, and 2 out of 5 subclasses in each superclass. For DomainNet Z. Wu et al. (), we select 10 out of 345 superclasses (labels), and 3 out of 5 subclasses (domains) in each superclass. We adopt ResNet LeCun et al. () for conducting the classification task to distinguish the superclass(^3^33each client only know the label of the superclass but doesn’t know the label of subclass) on CIFAR100 and DomainNet.

Federated Learning Settings.
We simulate a horizontal federated learning system with K = 5, 10, 20 clients in a stand-alone machine with 8 NVIDIA A100-SXM4 80 GB GPUs and 56 cores of dual Intel(R) Xeon(R) Gold 6348 CPUs. For DomainNet and CIFAR10, we regard each subclass follow one sub-distribution ($P_{i}$). For the disentangled extent, we choose the averaged entangled coefficient $s=\frac{2}{K(K-1)}\sum_{k_{1},k_{2}}d_{k_{1},k_{2}}$ over all clients as $0,...$. The detailed experimental hyper-parameters are listed in Appendix A.

Baseline Methods.
Four existing methods i.e., FedAvg McMahan et al. (): FedProx T. Li et al. (), SCAFFOLD Karimireddy et al. (), MOON Q. Li et al. () and the proposed method FedDistr are compared in terms of following metrics.

Evaluation metric.
We use the model accuracy on the main task to represent the utility, the communication rounds and transmission parameters to represent communication efficiency, and the standard derivation of noise level to represent the privacy leakage.

## Appendix B Ablation Study

We evaluate the different methods for a large $\xi=0.385$ in Figure 7. It shows that even for a large $\xi$, FedDistr achieves the best tradeoff between the communication rounds and utility loss. Moreover, compared to small $\xi$, other methods such as MOON converges to a better utility while it still requires beyond 100 communication rounds. Finally, we test the robustness of FedDistr for different number of clients in Figure 8. It illustrates FedDistr still achieve the better utlity than FedAvg with 20 clients.

Figure: (a) DomainNet
Refer to caption: extracted/5919221/imgs/exp/MainExp/cifar100-1.0.png

Figure: (a) DomainNet
Refer to caption: extracted/5919221/imgs/exp/more_clients.png

## Appendix C Theoretical Analysis

### C.1 Proof for Theorem 1

Consider $\mathcal{D}\sim S$ has total $n$ samples $\{X_{1},\cdots,X_{n}\}$.
We first demonstrate the estimated error of the distribution parameter on different data number $n$. In simplify for the analysis, we assume the $S$ is the Truncated normal distribution $\phi(\mu,\sigma,a,b)$ and only need to estimated parameter is $\mu$.

###### Lemma 1 .

Hoeffding inequality.
If $X_{1},X_{2},\ldots,X_{n}$ are independent random variables such that $X_{i}$ is bounded by [$a$, $b$] (i.e., $a\leq X_{i}\leq b$), and let $\bar{X}=\frac{1}{n}\sum_{i=1}^{n}X_{i}$ be the sample mean, then for any $\epsilon>0$:

$$ $P\left(\bar{X}-\mu\geq\epsilon\right)\leq\exp\left(-\frac{2n\epsilon^{2}}{(a-b )^{2}}\right)$ $$

where $\mu=E[X_{i}]$ is the expected value of the random variables.

According to Hoeffding inequality, the probability of the estimated error for the distribution parameter $\mu$ is smaller than $\exp\left(-\frac{2n\epsilon^{2}}{(a-b)^{2}}\right)$. And this probability tends to zero if the $n$ tends to infinity.

We estimate the distribution mean $\mu$ via $\hat{\mu}$ according to the $\mathcal{D}$. Let the $\hat{\mu}$ be the mean of the estimated distribution $\hat{S}$. Define the loss function $f(\omega,z)$ on the model $\omega$ and data $z$. Define expectation loss function $F(\omega)$ and $F^{\prime}(\omega)$ on data $S$ and $\hat{S}$ as:

$$ $F(\omega)=\mathbb{E}_{z\in S}f(\omega,z)\quad\text{and}\quad\hat{F}(\omega)= \mathbb{E}_{z\in\hat{S}}f(\omega,z)$ (1) $$

###### Lemma 2 .

Define $\omega^{*}=argmin_{\omega}F(\omega)$ and $\hat{\omega}^{*}=argmin_{\omega}\hat{F}(\omega)$. If $f$ is L-lipschitz, then the following holds with the probability $\exp\left(-2n\epsilon^{2}\right)$:

$$ $0\leq\mathbb{E}_{z\in S}f(\hat{\omega}^{*},z)-\mathbb{E}_{z\in S}f(\omega^{*}, z)\leq 2L\epsilon$ (2) $$

###### Proof.

According to definition of $\omega^{*}$ and $\hat{\omega}^{*}$, we have

$$ $\mathbb{E}_{z\in\hat{S}}f(\hat{\omega}^{*},z)\leq\mathbb{E}_{z\in\hat{S}}f( \omega^{*},z)\quad$ (3) $$

and

$$ $\mathbb{E}_{z\in S}f(\omega^{*},z)\leq\mathbb{E}_{z\in S}f(\hat{\omega}^{*},z).$ (4) $$

Therefore, we can obtain

$$ $\mathbb{E}_{z\in S}f(\hat{\omega}^{*},z)-\mathbb{E}_{z\in S}f(\omega^{*},z)\geq 0$ (5) $$

Furthermore, the following holds with the probability $1-\exp\left(-2n\epsilon^{2}\right)$ based on Lemma 1:

$$ $\begin{split}&\mathbb{E}_{z\in S}f(\hat{\omega}^{*},z)-\mathbb{E}_{z\in S}f( \omega^{*},z)\\ &\leq[\mathbb{E}_{z\in S}f(\hat{\omega}^{*},z)-\mathbb{E}_{z\in\hat{S}}f(\hat{ \omega}^{*},z)]+[\mathbb{E}_{z\in\hat{S}}f(\omega^{*},z)-\mathbb{E}_{z\in S}f( \omega^{*},z)]\\ &\leq\int_{z\in S}|f(\hat{\omega}^{*},z_{1})-f(\hat{\omega}^{*},z_{1}+\epsilon )|dp(z)+\int_{z\in S}|f(\omega^{*},z_{1})-f(\omega^{*},z_{1}+\epsilon)|dp(z)\\ &\leq L(z_{1}+\epsilon-z_{1})+L(z_{1}+\epsilon-z_{1})\\ &=2L\epsilon,\end{split}$ (6) $$

where $p(z)$ is the probability density function. The last inequality is due to the L-lipschitz of $f$.
∎

Lemma 2 demonstrates the utility loss when transferring the estimated distribution. Consider $K$ clients participating in federated learning with their data $\mathcal{D}_{k}=\{(x_{k,i},y_{k,i})\}_{i=1}^{n_{k}}$. Denote $\mathcal{D}=\mathcal{D}_{1}\cup\cdots\cup\mathcal{D}_{K}$ and $\mathcal{D}$ follows the distribution $S$.
Then we prove Theorem 1 according to Lemma as follows:

###### Theorem 1 .

A data distribution across clients being disentangled is a sufficient condition for the existence of a privacy-preserving federated algorithm that requires only a single communication round and achieves the $\epsilon$ utility loss with the probability $1-\exp(-\frac{\min\{n_{1},\cdots,n_{K}\}\epsilon^{2}}{2L^{2}})$.

###### Proof.

If the data distribution across clients is disentangled, it means each client $k$ can use their own $n_{k}$ data to estimate their data distribution. Therefore, according to Lemma 2, when achieving utility loss $\epsilon$ for distribution $S_{k}$ (i.e., $\mathbb{E}_{z\in S_{k}}f(\hat{\omega}^{*},z)-\mathbb{E}_{z\in S_{k}}f(\omega^{
*},z)\leq\epsilon$, the probability is $\exp(-\frac{n_{k}\}\epsilon^{2}}{2L^{2}(a-b)^{2}})$.

Then each clients can transfer estimated $\hat{S}_{k}$ to the server, thus,

$$ $\begin{split}&\mathbb{E}_{z\in S}f(\hat{\omega}^{*},z)-\mathbb{E}_{z\in S}f( \omega^{*},z)\\ &=\frac{1}{K}\sum_{k=1}^{K}\mathbb{E}_{z\in S_{k}}f(\hat{\omega}^{*},z)- \mathbb{E}_{z\in S_{k}}f(\omega^{*},z)\\ &\leq\frac{1}{K}K\epsilon=\epsilon.\end{split}$ (7) $$

And the probability to achieve $\epsilon$ utility loss is $1-\exp(-\frac{\min\{n_{1},\cdots,n_{K}\}\epsilon^{2}}{2L^{2}})$.

∎

###### Lemma 1 .

###### Lemma 2 .

###### Proof.

###### Theorem 1 .

###### Proof.

### C.2 Analysis on Near-disentangled case

###### Lemma 3 .

Consider $K$ unit vectors $\vec{a}_{k}=(a_{1,k},\cdots,a_{m,k})$ such that $\sum_{k=1}^{K}a_{i,k}=1$ for any $i$, if $<\vec{a}_{i},\vec{a}_{j}>\leq\xi$ for any $i\neq j$ and $\xi<\frac{1}{(K-1)^{2}}$, then $\max\{a_{i,1},\cdots,a_{i,K}\}\geq 1-(K-1)\sqrt{\xi}$

###### Proof.

Since $<\vec{a}_{i},\vec{a}_{j}>\leq\xi$ for any $i\neq j$, $\min\{a_{i,k_{1}},a_{i,k_{2}}\}\leq\sqrt{\xi}$ for any $i,k_{1}\neq k_{2}$. Therefore, we have

$$ $\begin{split}&\max\{a_{i,1},\cdots,a_{i,K}\}\geq 1-\sum_{k\neq 1}\min\{a_{i,1} ,a_{i,k}\}\\ &\geq 1-(K-1)\min\{a_{i,1},a_{i,k}\}\\ &\geq 1-(K-1)\sqrt{\xi}\end{split}$ (8) $$

∎

For the near-disentangle case (that small $\xi$), we have the following theorem:

###### Theorem 2 .

When the distributions of across $K$ clients satisfy near-disentangled condition, specifically, $\xi<\frac{1}{(K-1)^{2}}$, then there exists a privacy-preserving federated algorithm that requires only one communication rounds and achieves the $\epsilon$ expected loss error with the probability $\exp(-\frac{(1-(K-1)\sqrt{\xi})n\epsilon^{2}}{2mL^{2}})$.

###### Proof.

In simplify the analysis, we assume $\vec{\pi}_{k}$ to be unit vector and the number of data for each base distribution to be the same ($\frac{n}{m}$), i.e., $\sum_{k=1}^{K}\pi_{i,k}=1$. According to the Lemma 3, we have

$$ $\max\{a_{i,1},\cdots,a_{i,K}\}\geq 1-(K-1)\sqrt{\xi}$ (9) $$

which means for each base distribution, there exists one client has at least the $\frac{(1-(K-1)\sqrt{\xi})n}{m}$ data. Thus, the probability when achieving the utility loss less than $\epsilon$ for any base distribution $P_{i},1\leq i\leq m$ is

$$ $\exp(-\frac{(1-(K-1)\sqrt{\xi})n\epsilon^{2}}{2mL^{2}})$ $$

according to the lemma 2.

Therefore, the probability when achieving the utility loss on $S$ less than $\epsilon$ is

$$ $\min_{k\in[K]}\{2\exp(-\frac{(1-(K-1)\sqrt{\xi})n\epsilon^{2}}{2mL^{2}})\}= \exp(-\frac{(1-(K-1)\sqrt{\xi})n\epsilon^{2}}{2mL^{2}}).$ $$

∎

###### Lemma 3 .

###### Proof.

###### Theorem 2 .

###### Proof.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: Please see the weakness above and also the questions below:

1. How to do you plot figure 5? do you add noise somewhere to guarantee DP?
2. It seems the final training is performed over the so-called base distributions, in case of classification, is the learned model useful for the original client dataset?
soundness: 2
strengths: 1. Both communication cost and heterogeneity are important challenges in federated learning;
2. Authors verify the efficacy of the proposed algorithm through numerical experiments;
3. The proposed algorithm saves communication cost compared to baseline federated learning algorithms with little utility loss.
summary: In this work, authors claim that a major challenge of federated learning is due to the entanglement of data distribution across clients. As a result, authors proposed a communication-efficient algorithm named FedDistr which decouples client data distributions into the different base distributions, and then the server merges the base distributions, finally a model trained over synthetic data generated following the base distributions. Authors perform numerical experiments to compare with several baseline federated learning algorithms and demonstrate FedDistr obtains better utility, communication, privacy trade-off.
weaknesses: 1. The paper is not well-motivated. In federated learning, heterogeneity is a more common term compared to the so-call disentanglement. Furthermore, in the first paragraph, authors claims 'There is a consensus that this inefficiency stems from the entanglement of data distribution across clients', can you provide references to this claim of 'consensus'. In fact, in the later part of the introduction, authors show that it is the disentanglement where classical FL algorithms like FedAvg does not perform well.
2. The proposed algorithm is not clearly introduced. Both Algorithm 1 and Section 4 only covers the parts until 'Activate Distribution Alignment', it seems there are steps after this in Figure 3 where clients generate data locally and train locally? Please explain more.
3. The claim of privacy preserving lacks rigorous discussion. How do you guarantee that the proposed algorithm does not leak user privacy. Roughly, the propose algorithm transfer cluster center of local datasets (in latent space) to the server, will this leak sensitive user information?

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. There are some grammatical errors.
   - “FL is hindered by entanglement of data distributions across different clients” —> “FL is hindered by **the** entanglement of data distributions across different clients”
   - “both disentangled base distribution” —>  “both disentangled base distributions”
   - “…for achieving ideal efficiency of federated learning…” –> “…for achieving **the** ideal efficiency of federated learning…” 
   - ”extract data distributions via stable diffusion model, and then upload these decoupled distributions to the server” —>  ”extract data distributions via **a** stable diffusion model and then upload these decoupled distributions to the server”
   - “The server actively identifies the orthogonal or parallel between the base distributions uploaded by clients and aggregate the orthogonal distribution once.” —> “The server actively identifies the orthogonal or parallel between the base distributions uploaded by clients and **aggregates** the orthogonal distribution once.”

2. In the reference “Numerous studies (… Li et al..),” the year is missing; it’s recommended to update the references.

3. It is better to adjust the formatting of the references. For example, ...
   - Use**\citet**: “Latent Diffusion Model (LDM) Ho & Salimans (2022)” —>“Latent Diffusion Model (LDM, Ho & Salimans, 2022) ” 
   - Use**\citep**“: The autoencoder Van Den Oord et al. (2017); Agustsson et al. (2017) “ —>  “The autoencoder (Van Den Oord et al., 2017; Agustsson et al., 2017) “
   - Use**\citep**“: LDM Ho & Salimans (2022); Liang et al. (2024)” –> “LDM (Ho & Salimans, 2022; Liang et al., 2024)” 

4. Some equations lack punctuation at the end, such as Eqs. (8), (9), (10) … 

5. Some minor spelling issues ：
   - “consisting solely of ’cat’ images, … ’dog’ images” —> “consisting solely of **\``cat‘’ **images, … **\``dog’’** images”. 
   - “Theorem 2. When the distributions of across K clients satisfy” —> “Theorem 2. When the distributions across K clients satisfy”
   - “We conduct experiments on two datasets: x‘CIFAR100 has the 20…” —> “We conduct experiments on two datasets: CIFAR100 has the 20…” 
   - “For DomainNet and CIFAR10, we regard…” —> “For DomainNet and CIFAR100, we …” 
   - “the era of large language models (LLMs) (), the time” —> “the era of large language models (LLMs) , the time” 

6. It is recommended to present Algorithm 1 in a single-line format to avoid confusing indentations and spacing.

7. In Section 4.3, each client $ k$ uploads the distribution parameters to the server, but the paper states that “the server does not have access to the sequence of the distribution parameters.” This seems contradictory.

8. The idea of transferring information contained in data rather than the data itself has been reflected in some previous works, such as methods based on distillation and those using data Mixup. It is recommended to include a review and discussion of similar methods in the paper, and ideally, to add these baselines in the experiments.

9. If clients and the server have information about the base distributions, can they infer the distribution information of other clients? Would this potentially lead to privacy leakage? It would be beneficial to include a discussion in the paper regarding the privacy and security implications of the proposed method.
soundness: 3
strengths: 1. The paper is easy to follow.
2. The experiment results look good.
summary: This paper proposes a novel FedDistr algorithm, which employs stable diffusion models to transfer data distributions instead of model parameters between clients and the server. The theoretical and experimental results demonstrate the effectiveness of the proposed method. Overall, the paper is easy to understand, but I still have some concerns regarding the fundamental settings in federated learning and the risks of privacy leakage.
weaknesses: 1. The proposed method may have the potential to leak privacy.
2. Some related work on methods that transmit knowledge instead of models between clients and the server may need to be reviewed and discussed.
3. There are some minor spelling and grammatical errors.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - Marfoq et al. [1] provide an impossibility result that shows that federated learning under a mixture of distributions is impossible without further assumptions on the distributions. It seems that this result is applicable to your case. How can this be reconciled with your paper?
- Shamir and Srebro [3] show that one-shot distributed stochastic learning via averaging of parameters can be arbitrarily bad. How does this result relate to your work?
- It can be shown that one-shot federated learning is provably efficient for convex learning problems when using robust aggregation [4]. How does this relate to your approach?
- I have troubles understanding why Lemma 2 should be correct. It seems that without further assumptions on $S$ and $\widehat{S}$ we cannot bound the expected loss. Assume that under $S$ only $z$ has probability 1 and all others have probability 0 and that under $\widehat{S}$ some $z'$ has probability 1 and all others have probability 0. Now we can assume that $f(w,z)=0$ for all $w$ and for any number $C>2L\epsilon$, we can assume $f(w,z')=C$. Then the bound in Lemma 2 should not work. Did I understand an assumption wrongly? 
- The proof of Thm 1 heavily relies on Lemma 2. It basically says: if local distributions have nothing to do with each other, then we can estimate the global distribution from local distributions. Now with Lemma 2 this would imply that one can learn a good model on that. This not only seems to contradict the impossibility result in Marfoq et al., it also seems to contradict classical transfer learning results - if local tasks have nothing to do with each other, how can we learn from each other? 
- The paper does not state in the main text that for the theoretical results to hold, one has to find the global risk minimizer (as stated in the assumptions of Lemma 2). Since this is in practice infeasible for non-convex problems, such as deep learning, the paper should address the limitations of their theory.
- How is the model trained? Is it trained on the estimate of the global distribution $\widehat{S}$ by generating examples and training on them? How can such a model be good if local distributions are disentangled? 
- SCAFFOLD generally performs well on heterogeneous FL tasks; however, your results show it underperforms. Could you elaborate on potential reasons for this discrepancy and clarify whether experimental settings or parameters might have affected SCAFFOLD's performance?


References:

[1] Marfoq, Othmane, et al. "Federated multi-task learning under a mixture of distributions." Advances in Neural Information Processing Systems 34 (2021): 15434-15447.

[2] Wu, Yue, et al. "Personalized federated learning under mixture of distributions." International Conference on Machine Learning. PMLR, 2023.

[3] Shamir, Ohad, and Nathan Srebro. "Distributed stochastic optimization and learning." 2014 52nd Annual Allerton Conference on Communication, Control, and Computing (Allerton). IEEE, 2014.

[4] Kamp, Michael, et al. "Effective parallelisation for machine learning." Advances in Neural Information Processing Systems 30 (2017).
soundness: 1
strengths: - Novel application of diffusion models to disentangle and aggregate data distributions in federated learning.
- Demonstrates some improvements in communication efficiency and utility for specific datasets.
summary: The paper tackles the problem of heterogeneous federated learning by estimating the data distribution in a federated way. Specifically, it assumes that the global data distribution can be represented as a mixture of a set of base distributions, with local distributions varying from the global one in their mixture coefficients. The proposed approach uses diffusion models to disentangle data into base distributions and communicates these distributions using only a single round of communication.
weaknesses: - The approach is similar to previous approaches [1, 2] but does not discuss them.
- The theoretical results are questionable and assumptions are not clearly stated.
- The results seem to contradict previous theoretical results [1].
- The description of the approach omits important details, in particular, how the local models are trained and if local models can differ.

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 2
questions: I have a singular, but deep question that I need answered, because it is not at all clarified in the paper: Could you please clarify how a singular, global, classification model is obtained? In the referenced work of Liang et al., a clearly important reference as the authors have a very similar structure to the one presented, the authors learn to generate datasets and labels via the LDM, but further train a classifier in a FedAvg fashion. It seems to me that this work fixes the multiple-communication rounds used to learn the distribution parameters used to generate samples/labels at the edge, but speaks nothing about the classifier model itself.

The answer to this question is at the heart of the paper itself and needs to be answered. If a classifier is trained, how is it done? How does it perform on a test set that includes labels outside of the local training distribution? Does it require aggregation, as is typically done in federated learning?
soundness: 2
strengths: The primary, and possibly only, strength of the paper is the ability to achieve high accuracy scores with only one communication round in a heterogeneous federated learning scenario; this, followed by theoretical support proved in the appendix, the authors proposed a unique solution to a commonly known problem in federated learning. The idea of disentangling the local data is original and significant.
summary: The paper proposes an one-shot method approach to federated learning by introducing the concept of “disentangling” distributions. The authors determine two distributions to be disentangled if given the same mixed distributions, the stochastic weight vector has a cosine similarity of 0. The authors theoretically support their approach via Theorems 1 and 2. The proposed algorithm, FedDistr, trains a Latent Diffusion Model which learns to generate distributions which are aligned in a central server. The alignment process uses the distribution parameters to determine orthogonal and parallel distributions for different clients; this is accomplished by solving an optimization problem on a weighted bipartite graph. The consolidated distribution parameters are broadcasted to all clients.

The work provides a good, singularly novel, idea, that of entangled/disentangled distributions and the creation of parallel/orthogonal sets of parameters. The body of the work relies heavily on the work of Liang et al. (2024), "Diffusion-Driven Data Replay: A Novel Approach to Combat Forgetting in Federated Class Continual Learning." Unlike in Liang et al., which thoroughly explains each step of their algorithm, the authors sweep factors underneath the rug. For example, Liang et al. state that the Latent Diffusion Model (LDM) is used to generate images at the edge which are used to train a classifier. The classifier, trained locally, is aggregated through the use of the traditional FedAvg algorithm. The authors of this paper do not clarify if they do or do not do this. It must be assumed by the reader that they do not, but one who reads the related work I have referenced will wonder if that is truly the case. In fact, I believe a re-organization and further clarifications would strengthen the paper significantly.  The authors did not reference Liang et al. and break down their work in the Related Works section; this is an obvious mistake as the two are without a doubt built on the same core with a minor, although interesting and significant, improvement. I believe it would be beneficial to the paper if the authors added a flow chart or table comparing/contrasting the work of Liang et al. and their work.

If the authors had more clearly explained how a final classifier is made, if one was even trained, I would be willing to increase the score. It would be prudent to add a subsection on how classification works, if a model is trained or not, and how one achieves a singular classification model used by all agents. I believe the paper has good merit but with further clarifications and much clear writing and flow, the paper would be much stronger than it currently is. A key problem with the writing is the authors are relying on too much prior knowledge rather than opting in for clarity and verbosity, both of which would make it a much better read with much clearer communication of ideas. Furthermore, I believe the experiments section needs to be clearer with regard to what was done. The additional, but few, comments in the Appendix mention the authors used subsets of the classes of the datasets; what does this mean and how is the data distributed to the agents? If 20 classes are used and 5 clients exist in your network, are you forcing each agent to have a different subset of the 20 classes to force heterogeneity? This further speaks to the overall lack of clarity in the writing and explanation provided by the authors. I would ask the authors to better explain their experimental setup and the associated variables. I would also ask the following: How many classes of each dataset were used in total? How were these classes distributed among the clients? Was heterogeneity artificially induced, and if so, how?

Overall, independently of the lack of clarity, I believe the paper is relevant to readers in the field of Federated Learning, as it proposes an interesting one-shot approach to heterogeneous federated learning. If possible, I would prefer to see the paper rewritten for clarity but the results and ideas are important enough to warrant acceptance.
weaknesses: The weaknesses of the paper are all indirectly related to their contribution. The writing lacks extreme amounts of clarity all around. The authors are not clearly explaining concepts and do not mention how extremely similar the core of their work is to Liang et al. (2024). They cited the paper in a small section of the paper, but it is much more significant than they lead the readers to believe and because of this it needs to be introduced in the related works section with a proper break down. The authors further fail to clarify if the ResNet classifier is trained and fused or if the one-communication step is only to learn to generate images and labels through the Latent Diffusion Model. These are important questions. Furthermore, there are several things to cover that I will double-up and put in the "Questions" section because I believe the authors need to answer them. Here they are:
1) In the experiments section, the authors did not explain the exact precisely enough what the federated architecture looks like. While stating that K=5,10,20, does this mean there are a total of K clients in the network and all K are used when training? It is typical to show results with a subset of the agents as the contributing agents in the learning step. For example, in the work that introduced FedAvg, the authors design a network with 100 agents and use subsets as the contributing agents when fusing models, such as 10 of the 100 are used to average their models together.
2) The authors also don’t clearly state what the distribution of the data looks like at the edge. While their proposed algorithm aim to disentangle the distributions, one wonders what the underlying initial distribution of data does to the result of the model. Perhaps the authors’ goal was to limit to cases of disentangled and nearly-disentangled, but it would be better to make this clear. Furthermore, of the subclasses of the datasets selected, how are those samples distributed? This is also not clearly defined.
3) The paragraph between eq 2 and 3 contains the unsubstantiated claim, other than their own claim on the matter via the presented “Theorem 1”: “Both theoretical analysis and empirical studies show that data heterogeneity is a blessing rather than a curse, as long as data distributions among different clients can be completely disentangled.” This claim needs to be supported by citations such that readers can confirm such a bold claim, which is commonly known to be a difficult hurdle for FL-algorithms to overcome.
4) In the paragraph between equation 10 and 11, the authors state the following: “To disentangle the data distribution for client k, the data is segregated based on these latent feature embeddings.” What exactly does this mean and how is this accomplished?
5) The second paragraph of section 4.3, “Active Distribution Alignment,” needs clarification on what it meant by “the server does not have access to the sequence of the distribution parameters.” Further, when the authors state “orthogonal and parallel data base distributions,” what does this mean? 
6) The proofs in the appendix need more clarifications to be more clear and need grammatical cleaning (e.g., the proof of Theorem 1 has a "}" in the "exp( . )" above equation 7).