Provably Near-Optimal Federated Ensemble Distillation with Negligible Overhead

Comments on theoretical assumptions

We believe the reviewer's suggestion regarding $L$-smoothness may arise from a different interpretation of the convexity assumption. The reviewer seems to have interpreted the convexity assumption with respect to the model parameter $\theta$, whereas our convexity assumption is in fact with respect to the model output $\hat{y}$.

While deep neural networks like ResNet-18 do not exhibit convexity with respect to $\theta$, our theoretical results do not rely on such a convexity with respect to $\theta$. Instead, we assume the convexity of the loss function with respect to $\hat{y}$, which is standard in the literature. We acknowledge that our explanation may have caused some confusion, and we will revise the statement in our paper more clearly.

Comments on GAN-based assumptions

We agree that, to theoretically guarantee the achievement of optimal single-model performance, the generator and discriminator must be properly trained. However, our experimental results in Appendix F.7 demonstrate that the proposed method is robust to the quality of the generator and discriminator. In particular, in Appendix F.7.1, we showed that FedGO even with a completely untrained generator outperforms FedDF in terms of both server test accuracy and ensemble test accuracy. Also, in Appendix F.7.2, we can see that FedGO still significantly outperforms baselines even when client discriminators were trained at only one-sixth of the main setting.

Next, regarding privacy leakage due to the provision of the discriminator, we have conducted a comprehensive privacy leakage analysis in Appendix G. Table 15 demonstrates that by incorporating local differential privacy (LDP), FedGO can guarantee a client-side privacy level comparable to FedAVG. Furthermore, as a measure to prevent excessive client distribution leakage, we adopted a simple four-layer CNN for the discriminator, as detailed in Appendix E.2 and F.7.3, and implemented the output activation using a double composite sigmoid function, restricting the discriminator’s output range to [sigmoid(0), sigmoid(1)].

Comments on client-side data limitations

Our experimental results demonstrate that, contrary to the reviewer's concern, FedGO is effective even with low-volume client datasets. We present experimental results not only for 20 clients but also for 100 clients in Appendix F, where each client has an average of only 250 data samples. Figure 7 shows that FedGO achieves significant performance improvements over existing baselines even in such a client data-deficient situation.

Comments on FedGO and data heterogeneity

Federated ensemble distillation algorithms leverage additional unlabeled datasets at the server to perform pseudo-labeling and knowledge distillation, thereby enhancing server model performance. However, in client data heterogeneous situations where client distributions are highly diverse, inference quality for server unlabeled data $x$ varies significantly across clients. Using a fixed weighting function (e.g., uniform weighting) for pseudo-labeling can degrade the quality of pseudo-labels in proportion to the average discrepancy between the client average distribution $p$ and each individual client distribution $p_k$ (as summarized in Table 1 of our paper, specifically (1.1)).

Thus, research has focused on developing weighting methods that assign higher weights to clients with higher inference quality per data sample $x$. DaFKD is one such method. However, in large-client scenarios, its generalization bound become vacuous, and there has been little research on weighting methods that are theoretically guaranteed and robust to client data heterogeneity.

In this paper, we demonstrate in Theorem 3.4 that assigning weights in a specific manner ensures that pseudo-label performance remains independent of client data heterogeneity while providing the tightest existing generalization bound (as discussed in Table 1, specifically (1.3), and Definition 3.1). Furthermore, we show that the generalization bound for the server model trained on these pseudo-labels is expressed in terms independent of client data heterogeneity, proving that our ensemble distillation scheme is theoretically robust to client data heterogeneity.

We modeled this weighting function through client discriminators, as presented in Theorem 3.6, and implemented it via FedGO. Our experimental results across various settings confirm that FedGO is significantly more robust to client data heterogeneity compared to baseline algorithms.

We hope that this answers your question. We will emphasize this aspect further in the paper to enhance clarity.

Comments on data-free setting

We believe we have thoroughly examined both scenarios, with and without a server dataset. For data-free setting, we have conducted experiments using both off-the-shelf generator and generator trained via federated learning and analyzed the results in Appendix F.5. Also, we have analyzed the communication, privacy, and computational complexity of data-free FedGO in Appendix G.

Comments on security risks

In accordance with the reviewer’s comment, we have conducted additional experiments where 5 and 10 out of 20 clients were Byzantine, outputting only the maximum value for the discriminator. The results showed that while the accuracy on the CIFAR-10 classification task under $\alpha=0.05$ was initially 72.35$\pm$9.01, it dropped to 69.75$\pm$5.05 with 5 Byzantine clients and 66.38$\pm$4.97 with 10 Byzantine clients. Even in this extreme scenario where half of the participants were Byzantine, our method significantly outperformed all the baselines that did not utilize a discriminator. We will report this result in the final paper.

Comments on theoretical assumptions

As shown in the experimental results in Appendix F.1 and Appendix F.5, FedGO outperforms existing baselines even when clients have limited data. Specifically, the experiments were conducted with 100 clients, each having an average of only 250 data samples—an amount that is insufficient for training an optimal discriminator.

Comments on generator assumptions

For the case (G3) where a generator and a discriminator are trained using an FL approach, we have already addressed the reviewer’s concern by conducting experiments with a small number of training samples per client (250 images per client when there are 100 clients) in Appendix F.5 and by showing that additional client-side computational overhead is negligible compared to FedDF, which does not train a generator and a discriminator.

For the case (G2) where an off-the-shelf generator is used, we have also conducted experiments where the distribution of the generator differs from that of the client data. These results are presented in Table 5 of Section 4.2. Specifically, when the off-the-shelf generator was trained on ImageNet and the client datasets were CIFAR-10, CIFAR-100, or ImageNet100, the performance remained comparable to the case where the generator was trained on data matching the client distribution. This indicates that even when the off-the-shelf generator is trained on a different dataset from clients’ data, clients can still train their discriminators effectively.

We hope our response addresses your concerns. If we have misunderstood your question, please feel free to clarify.

Comments on experimental design The simplest weighting method for ensemble distillation would be the uniform weighting that FedDF incorporates. The paper demonstrates that the improvement from FedAVG to FedDF stems from the benefit of ensemble distillation itself, while the improvement from FedDF to FedGO is attributed to GAN-based weighting. Additionally, extensive experiments were conducted by fixing the ensemble distillation process while varying the weighting method, effectively quantifying the contribution of weighting in Figure 2 of our main paper.

Comments on deployment challenges

• Latency: We provided a comparison of the MFLOP counts between the baseline and FedGO algorithms in Table 16 of Appendix G. The comparison of MFLOP counts can serve as a proxy for latency comparison.
• Scalability: We have already provided experimental results in Appendix F.1 and F.5 with a large-scale setup of 100 clients under various settings.
• Potential biases in GAN-based weighting: In accordance with the reviewer’s comment, we have additionally conducted experiments with malicious clients and demonstrated the effectiveness of our weighting method in such a challenging scenario. We have reported the experimental results in the response to the second comment.
• Generalization to non-image tasks: In accordance with the reviewer’s suggestion, we have additionally conducted experiments with a tabular healthcare dataset, confirming performance improvements over FedAVG and FedDF as shown in the table below. In this experiment, we used a total of four clients, all of whom participated in every communication round. Regarding NLP, FedDF and some GAN-based approaches have demonstrated promising results, indicating the strong generalization potential of our weighting method. This appears to be an interesting direction for future research, and we appreciate the suggestion.

\begin{array}{|l|cc|cc|cc|} \hline & \alpha = 0.1 & \alpha = 0.05 \\ \hline \text{Central training} & 36.21\pm0.15 \\ \hline \text{FedAVG} & 34.20\pm0.56 & 33.82\pm0.86 \\ \hline \text{FedDF} & 34.66\pm0.22 & 34.21\pm0.46 \\ \hline \text{FedGO} & 34.81\pm 0.36 & 34.64\pm0.32 \\ \hline \end{array}

Comments on pre-trained generator robustness

Our theoretical analysis already takes into account such discrepancy: Theorem 3.6 says that our weighting method produces the optimal weight $w_k^*$ for the data point on supp$(p)\cap$ supp$(p_g)$, where $p$ is the average client data distribution and $p_g$ is the generator’s data distribution. Thus, as we mentioned in the last paragraph of Section 3.1, it is theoretically guaranteed that our proposed method works properly as long as the generator is capable of producing sufficiently diverse samples. However, empirically, we showed even stronger results in Appendix F.7.1: FedGO with a completely untrained generator still performs better than FedDF.

Comments on communication overhead

We confirm that the additional communication cost of our method is indeed negligible, because (1) the number of parameters for the generator and the discriminator is much smaller than that of the classifier architecture and (2) a small number of communication rounds is sufficient for training generator and discriminator, i.e., only 5 rounds for our experiment. More specifically, in Appendix G, we showed that FedGO incurs only an additional 0.47% in communication overhead even in the most challenging scenario where the generator needs to be federated, i.e., case (G3) + (D2). Note that for the scenario where the number of parameters for the generator and discriminator should be large due to the input dimension, the number of parameters for the classifier will also be large, thereby ensuring that the relative overhead remains negligible.

Comments on server-side resource consumption

We reported the computational complexity in terms of the number of MFLOPs (rather than runtime) in Appendix G.2. The cost of the data generation process is already included in the MFLOPs. In cases where there is no unlabeled dataset on the server, we generated 25,000 images using the generator before main FL stage. Once generated, this dataset is reused for distillation in every communication round. The image generation process takes only 3 seconds and is performed just once per experiment, making its time cost negligible.

Comments on extreme data heterogeneity

We first would like to emphasize that our weighting method is proposed to address data heterogeneity, and our theoretical analysis already takes into account the discrepancy between the client data distribution and the generator distribution. To validate the effectiveness of our method in highly heterogeneous setting, we considered Dirichlet parameter $\alpha = 0.05$, for which it is common for only two or three clients out of 20 clients to possess all the images of a particular class. The data distributions for $\alpha = 0.1$ and $\alpha = 0.05$ are visualized in Figures 5 and 6 in the Appendix E.2. Additionally, to account for more complex data scenarios, we conducted experiments using the ImageNet100 dataset. Since ImageNet is sufficiently complex, we believe these experiments offer meaningful insights into FedGO’s performance under realistic and challenging conditions.

Comments on theoretical guarantees under heterogeneity

The proposed method is designed to address data heterogeneity, and our theorems and their proofs already take into account possible imbalance or heterogeneity in the client data. For “model” heterogeneity, we assumed homogeneous model structures in Theorem 3.2, Corollary 3.3, and throughout all of our experiments. As discussed in Appendix H (Limitations), defining an optimal model ensemble becomes challenging when dealing with multiple hypothesis classes (synonym for multiple client model structures).

Comments on experimental completeness - We appreciate this valuable suggestion. To address this concern, we have implemented and tested the following baseline algorithms targeting data heterogeneity (published in 2024) under the main experimental setting in Section 4.1:

1. FedUV (CVPR 2024)
2. FedTGP (AAAI 2024)

As shown in the table below, FedGO outperforms both baselines across all settings. We will incorporate these results into the final version of our paper.

\begin{array}{|l|cc|cc|cc|} \hline & \text{CIFAR-10} && \text{CIFAR-100} && \text{ImageNet100} & \\ & \alpha = 0.1 & \alpha = 0.05& \alpha = 0.1 & \alpha = 0.05& \alpha = 0.1 & \alpha = 0.05 \\ \hline \text{FedUV} & 62.58 \pm 4.83 & 53.80 \pm 5.68 & 38.84 \pm 0.79 & 36.17 \pm 1.24 & 30.09 \pm 1.09 & 27.32 \pm 0.65 \\ \hline \text{FedTGP} & 61.16 \pm 6.98 & 61.51 \pm 7.78 & 39.58 \pm 0.10 & 36.56 \pm 0.11 & 29.21 \pm 1.13 & 26.34 \pm 1.02 \\ \hline \text{FedGO} & \mathbf{79.62} \pm 4.36 & \mathbf{72.35} \pm 9.01 & \mathbf{44.66} \pm 1.27 & \mathbf{41.04} \pm 0.99 & \mathbf{34.20} \pm 0.71 & \mathbf{31.70} \pm 1.55 \\ \hline \end{array}

Comment 1: Important Limitation: The theoretical analysis — including the derivation of the optimal weighting functions and generalization bounds — is restricted to binary classification tasks. This limitation is underemphasized in the main text, yet all experiments are conducted on multi-class classification problems. While the empirical results are compelling, it remains unclear how well the theoretical results translate to the multi-class case. Clarifying or extending the theoretical framework to multi-class settings would strengthen the paper's claims considerably.

Thank you for reviewing our paper. In the following, we have clarified the assumption of binary classification tasks in deriving the optimal weighting functions and generalization bounds.

The optimality of our proposed weighting function, based on Theorem 3.4 and Theorem 3.6, is not restricted to binary classification tasks. Theorem 3.4 only requires the convexity of loss function and Theorem 3.6 does not require either the convexity of the loss function or binary classification tasks.

Our generalization bound in Theorem C.1 indeed assumes binary classification tasks. However, we would like to highlight that obtaining a tight generalization bound is challenging even in binary classification tasks and remains an active area of research. For example, FedDF, Fed-ET, and DaFKD derived generalization bounds under binary classification assumptions, but their bounds either tend to degrade in data-heterogeneous settings or become vacuous in large-client scenarios. However, our bound is tighter than those bounds, making it a significant contribution even within the binary classification framework. For the multi-class case, there are some existing results on obtaining generalization bounds in a single-model (non-federated learning) setup, but they remain limited and loose.