Disentangling data distribution for Federated Learning

Q1. In the experiments section, the authors did not explain the exact precisely enough what the federated architecture looks like.

We are sorry for this unclear claim. Our experimental setting is all $K$ clients participate for each communication round. We would include more clients sampling strategy in the future version.

Q2. The authors also don’t clearly state what the distribution of the data looks like at the edge.

The data is distributed according to the $\epsilon$. The small $\epsilon$ indicates the large heterogeneity of the data distribution among clients.

Q3. 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?

We apologize for the unclear clarification. This statement means that each client applies a clustering algorithm to the latent feature embeddings to disentangle their own complex distributions. The client then uploads some of these disentangled base distributions to the server.

Q4. 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”.

The server receives the distribution parameters uploaded by different clients. Then it aligns these distributions to identify the orthogonal ($\mathcal{I}_o$) and parallel sets ($\mathcal{I}_p$).

Q5. The authors state “orthogonal and parallel data base distributions”, what does this mean?

The orthogonal set refers to the set of prompts that are largely different, while the parallel set refers to the set of prompt embeddings where the prompts are similar.

Q6. Could you please clarify how a singular, global, classification model is obtained?

We apologize for the confusion. The detailed process is as follows.

After the server actively aligns the uploaded base distributions, each client receives the distributed base distributions $v=[v_p,v_o]$ from server. For each base distribution embedding $v_o^i \in v_o$, clients use the latent diffusion model to generate data $\widehat{x} _ i=\{\widehat{x} _ {i,j} | j\in [n_i] \}$ following the base distribution. $\widehat{x} _ {i,j} = D(f(z _ {T}, T, [p, v_o^i])), j \in [n_i] \text{, for } i \in [|v_0|],$ where decoder $D$ is for the inverse mapping of encoder $E$, $f$ is the denoising process, $T$ is the number of diffusion steps, and $[p, v_i]$ is the conditional prompt corresponding to base distribution $i$.

Using the synthetic data, the clients locally train their downstream task model. $\mathcal{F} _ {k}^{*} = \text{argmin} _ {\mathcal{F} _ {k}} \mathbf{E} _ {i\in [|v_0|]} \mathcal{L} _ {CE} (\mathcal{F} _ {k}(\widehat{x} _ {i,j}), \widehat{y} _ i) \text{, for } k \in [K],$ where $\mathcal{F} _ k$ is the local downstream model of client $k$, $\mathcal{L} _ {CE}$ is the cross entropy loss, $\{\widehat{x} _ {i,j} | j\in [n_i] \}$ is the generated data corresponding to base distribution $i$.

Q7. 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?

We apologize for the confusion. First, the classifier is locally trained on synthetic data generated from the base distributions without any aggregation. Second, we evaluate the trained model using a test dataset that includes all the labels to assess whether achieving a global optimum [1,2].

[1]Li T, Sahu A K, Zaheer M, et al. Federated optimization in heterogeneous networks[J]. Proceedings of Machine learning and systems, 2020, 2: 429-450.

[2]Li Q, He B, Song D. Model-contrastive federated learning[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021: 10713-10722.

Q1. The approach is similar to previous approaches but does not discuss them.

Regarding [1,2,3]: Our paper focuses on improving communication efficiency while other papers [1,2,3] concentrated on enhancing the model utility (global and personalized performance).

Regarding [4]: While our paper primarily aims to reduce the communication rounds, [4] focuses on minimizing computational complexity.

[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).

Q2. The theoretical results are questionable and assumptions are not clearly stated.

I apologize for the earlier lack of clarity. To clarify, $S$ and $\hat{S}$ represent truncated Gaussian distributions, and the mean difference between them is bounded by $\epsilon$. Based on this, we derive Eq.(6).

The intuition behind Theorem 1 is that when local distributions are independent, they can be concatenated into a single global distribution. This global distribution can generate data consistent with any of the local distributions. As a result, the generated data can effectively support the training of the downstream model.

Q3. The results seem to contradict previous theoretical results [1]. The conclusion is not in contradiction with [1]. The key difference lies in the aggregation approach: while [1] averages all model parameters, FedDistr adopts a more selective strategy by disentangling the parameters. Specifically, FedDistr actively separates the disentangled and entangled components, aggregating the entangled parts multiple times while concatenating the disentangled parts only once.

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

Q4. The description of the approach omits important details, in particular, how the local models are trained and if local models can differ.

Please refer to Q3 of Reviewer mZWU for further details.

Q5. SCAFFOLD generally performs well on heterogeneous FL tasks; however, your results show it underperforms.

This is because this paper mainly focus on the disentangled and near-disentangled scenario, i.e., the data distribution of clients are extreme Non-IID. Therefore, SCAFFOLD doesn't perform well.

Q1. The proposed method may have the potential to leak privacy.

In this paper, clients transfer the base data distribution with DP between clients and server, which protect clients' data privacy. For example, we add random noise following normal distribution to the uploaded base distribution embeddings, aiming satisfy differential privacy. In Algorithm 1, each client upload $\{\widetilde{v} _ {k,i}\} _ {i=1}^{m_k}$ to the server. $\widetilde{v} _ {k,i} = v _ {k,i} / max\{ \lbrace 1, \frac{\|v _ {k,i}\| _ 2}{C} \rbrace \} + \mathcal{N}(0, \sigma^2 C^2 I), i \in [m_k], \text{ for } k \in [K].$where $C$ is the norm upper bound, $\sigma$ is the standard derivation of added noise.

Moreover, we recognize the potential leakage of group statistics, such as data means. Techniques such as sparsification[1] and quantization[2] of the uploaded distribution parameters may help mitigate this risk. We plan to explore the protection of these aspects in future work.

Finally, we emphasize that the small size of the transferred parameters and the single communication round make this approach compatible with homomorphic encryption (HE) methods. Using HE would allow computations on encrypted embeddings directly, further eliminating the risk of privacy leakage without significant additional computational overhead. Incorporating HE is an avenue we plan to investigate for scenarios requiring stronger privacy guarantees.

[1]Aji A F, Heafield K. Sparse communication for distributed gradient descent[J]. arXiv preprint arXiv:1704.05021, 2017.

[2]Shlezinger N, Chen M, Eldar Y C, et al. UVeQFed: Universal vector quantization for federated learning[J]. IEEE Transactions on Signal Processing, 2020, 69: 500-514.

Q2. Some related work on methods that transmit knowledge instead of models between clients and the server may need to be reviewed and discussed.

Thanks for your constructive suggestion. Prior works have explored the idea of transmitting knowledge about private data, such as methods based on knowledge distillation and data Mixup. However, these approaches still require lots of communication rounds.

Specifically, knowledge distillation methods frequently require multiple communication rounds and may introduce significant communication overhead due to the transmission of auxiliary models or generators. For example, FedDF[3] distills the ensemble of client teacher models into a server-side student model, necessitating multiple exchanges of logits. Similarly, [4] proposes a data-free knowledge distillation approach where the server learns a generator to aggregate user information, requiring repeated transmission of a lightweight generator. Methods based on data Mixup, such as [5], also involve multiple communication rounds, as they rely on transmitting a batch of averaged data per round.

In contrast, our proposed FedDistr significantly improve the communication efficiency, which only requires one round of transmitting embeddings. In FedDistr, clients more accurately extract the data distribution by using the latent diffusion model and upload data distributions represented by LDM conditional embeddings. The server actively aligns the conditional embeddings and distribute the embedding to the clients.

We will include a discussion of these related works and add the methods as baseline in the revised manuscript.

[3]Lin T, Kong L, Stich S U, et al. Ensemble distillation for robust model fusion in federated learning[J]. Advances in neural information processing systems, 2020, 33: 2351-2363.

[4]Zhu Z, Hong J, Zhou J. Data-free knowledge distillation for heterogeneous federated learning[C]//International conference on machine learning. PMLR, 2021: 12878-12889.

[5]Yoon T, Shin S, Hwang S J, et al. Fedmix: Approximation of mixup under mean augmented federated learning[J]. arXiv preprint arXiv:2107.00233, 2021.

Q3. In Section 4.3, each client 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.

We apologize for the confusion. The server receives the distribution parameters uploaded by different clients. Then it aligns these distributions to identify the orthogonal ($\mathcal{I} _ o$) and parallel sets ($\mathcal{I} _ p$).

Q4. Errors and Typos

Thank you for your thorough review and for pointing out these grammar, formatting, and presentation issues. I have carefully addressed all the points in the revised manuscript.

Q1. In federated learning, heterogeneity is a more common term compared to the so-call disentanglement.

We use $\epsilon$-entangled to represent the extent of heterogeneity among different clients, as defined in Definition 2 of the main text. Specifically, when $\epsilon = 0$, it corresponds to the disentangled case, where the data distributions of different clients are entirely distinct. Furthermore, Theorem 2 establishes the relationship between utility loss and the entangled extent $\epsilon$.

Q2. Can you provide references to this claim of 'consensus'?

In distributed systems, tightly disentangled tasks exhibit greater efficiency[1], as complex tasks can be decomposed into simpler subtasks that can be assigned to multiple clients for parallel execution. For instance, in a medical imaging diagnosis task, if different clients (e.g., hospitals) naturally have a disentangled focus on specific sub-populations (e.g., one hospital deals more with elderly patients, while another focuses on pediatric cases), this inherent disentanglement enables each client to contribute highly specialized and complementary knowledge to the global model. Such complementary contributions reduce redundancy and conflict during aggregation, allowing the server to achieve efficient convergence with fewer communication rounds compared to scenarios with entangled and homogeneous data distributions.

Existing algorithms such as FedAvg fail to actively leverage such disentanglement. Their reliance on simple aggregation mechanisms like weighted addition overlooks the underlying structure of the data distributions, leading to suboptimal performance, especially in heterogeneous settings. Theoretical results (Theorem 2) demonstrate that when the server actively leverages disentanglement, it is possible to achieve global model utility with only one round of communication while maintaining a tolerable utility loss. This highlights the potential for significant efficiency improvements under disentangled data distributions.

[1]Coulouris G F, Dollimore J, Kindberg T. Distributed systems: concepts and design[M]. pearson education, 2005.

Q3. The proposed algorithm is not clearly introduced. It seems there are steps after this in Figure 3 where clients generate data locally and train locally.

We apologize for the confusion. Briefly, after the server actively aligns the uploaded base distributions, each client receives the distributed base distributions $v=[v_p,v_o]$ from server. For each base distribution embedding $v_o^i \in v_o$, clients use the latent diffusion model to generate data $\widehat{x} _ i= \lbrace \widehat{x} _ {i,j} | j\in [n_i] \rbrace$ following the base distribution. $\widehat{x} _ {i,j} = D(f(z _ {T}, T, [p, v_o^i])), j \in [n_i] \text{, for } i \in [|v_0|],$ where decoder $D$ is for the inverse mapping of encoder $E$, $f$ is the denoising process, $T$ is the number of diffusion steps, and $[p, v_i]$ is the conditional prompt corresponding to base distribution $i$.

Using the synthetic data, the clients locally train their downstream task model. $\mathcal{F} _ {k}^{*} = \text{argmin} _ {\mathcal{F} _ {k}} \mathbf{E} _ {i\in [|v_0|]} \mathcal{L} _ {CE} (\mathcal{F} _ {k}(\widehat{x}_{i,j}), \widehat{y}_i) \text{, for } k \in [K],$ where $\mathcal{F} _ k$ is the local downstream model of client $k$, $\mathcal{L} _ {CE}$ is the cross entropy loss, $\{\widehat{x} _ {i,j} | j \in [n_i] \}$ is the generated data corresponding to base distribution $i$.

Q4. The claim of privacy preserving lacks rigorous discussion.

Please refer to Q1 of Reviewer n8vX for further details.

Q5. 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?

The learned model is indeed useful for the original client dataset. This is because our work focuses on achieving global optimal [2,3], which ensures the model optimality across all clients collectively.

[2]Li T, Sahu A K, Zaheer M, et al. Federated optimization in heterogeneous networks[J]. Proceedings of Machine learning and systems, 2020, 2: 429-450.

[3]Li Q, He B, Song D. Model-contrastive federated learning[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021: 10713-10722.