Revisiting Ensembling in One-Shot Federated Learning

We thank the reviewer for their positive feedback and constructive comments. We address the reviewer's questions below.

Q. Does the communication cost comparison in Figure 2 assume quantization? If so, does the communication cost for the iterative FL methods also use quantization?

$\rightarrow$ Yes, FENS is presented with quantization in Figure 2. While the iterative FL baseline (FedAdam) in Figure 2 does not use quantization, we have included the relevant baseline with quantization, called FedAvg STC in Figure 4. The quantization in FedAvg STC results in an accuracy drop compared to FedAdam.

Q. The performances reported in Appendix D are impressive, but Table 15 (Fed-ISIC2019) shows a limitation/non-robustness of one-shot methods on a real-world dataset. Why do you think the gap here is large?

$\rightarrow$ This gap can be explained through the quality of local models in case of the FedISIC2019 dataset together with our experimental results from Section 3.3.1. Our experiments in Section 3.3.1 show that local data size plays a key role in determining when FENS can match the performance of iterative FL. The FedISIC2019 dataset, in particular, exhibits high variance in the amount of local data held by the clients, and consequently their local model performance. As reported in Table 3 (Appendix A), the biggest client has 12k samples while the smallest client has only 225 samples. Thus, we speculate that the FedISIC2019 dataset falls in the low local training fraction regime of the curves in Fig. 5, exhibiting a larger gap compared to iterative FL.

Collecting the logits of all clients for every forward pass of the ensemble is not elegant, especially given the memory and computational costs. The authors do acknowledge this limitation and offer some directions to mitigate this, but they still do not offer a practical solution that scales well.

$\rightarrow$ We currently use FP32 to INT8 quantization to mitigate this issue (Appendix B.3). This reduces the costs by $4\times$. An alternative and complementary approach is to consider an ensemble such that the total size is comparable to the single model trained in iterative FL. Preliminary results indicate that under a comparable memory footprint, FENS still remains competitive with the original FedAdam. We will include the complete ablation in our revised version.

I think running experiments where some clients are sampled every round of training the ensemble is straightforward to implement, so it would be interesting to see that. How would you concatenate the logits?

$\rightarrow$ This is indeed a highly relevant suggestion. We explored sampling clients each round by setting logits from non-sampled models to zero in the concatenated vector. Additionally, we also tried learning mask tokens that fill for non-sampled logits. Preliminary results show an accuracy drop under high heterogeneity, likely because each model’s contribution is significant in such scenarios. To improve performance, a more sophisticated logit concatenation method may be required.

On suggested comparison to personalized FL baselines and Split Learning.

$\rightarrow$ We thank the reviewer for their insightful suggestions. We agree that Split Learning might help make the computational budget more feasible. In FENS, we focus on simultaneously achieving OFL-like communication efficiency and FL-like accuracy. In this sense, the objectives of Split Learning and personalized FL slightly differ for a direct comparison. Personalized FL focuses on individual client accuracy, while FENS targets global data accuracy. Similarly, Split Learning reduces computational burden, whereas FENS prioritizes communication efficiency. We believe that both these paradigms can be orthogonally explored within FENS and are keen on exploring these aspects in future work.

We thank the reviewer for their positive feedback and comments. We address the reviewer's questions below.

Q. How to determine the best number of iterations for the aggregator model?

$\rightarrow$ This is determined as in standard FL schemes by stopping when the validation loss has converged or validation accuracy is not improving further.

Q. How is the MoE method implemented over the local model ensemble? Why did it incur such a high cost?

$\rightarrow$ MoE aggregation involves both the input and the logits in the following form: $f(x)=\sum_{i \in [M]}G(x)\_i\pi_i(x)$. Here, $G: \mathcal{X} \rightarrow [0,1]^M$ is the gating network that generates scalar weights for every expert $[\pi_1,\ldots,\pi_M]$ based on the input $x \in \mathcal{X}$. In FENS with MoE aggregation, only the gating network is trained via the federation while $[\pi_1,\ldots,\pi_M]$ correspond to the (frozen) locally trained client models. For the gating network, we employ a CNN with two convolutional blocks, followed by two fully connected layers and a final classification head. Since the gating model has several layers, it incurs a bigger communication overhead of iterative transfers as compared to just an MLP based shallow aggregator (called NN aggregator) in Fig. 7. We also tried using smaller models, however these resulted in accuracy deterioration. We have included the MoE training details in Appendix B.4.

Lack of theoretical analysis to ground the design of the proposed method.

$\rightarrow$ We acknowledge that deriving a formal theoretical explanation for our empirical findings is an interesting and valuable research direction. However, as briefly discussed in Section 2.2, our algorithm is motivated by the study of stacked generalization [1] in the centralized ensemble literature. It has been demonstrated that such stacking of models leads to higher-level models correcting for the biases of lower-level models in the stack, thereby improving overall generalization performance. While stacked generalization has been primarily studied in centralized setting, through FENS, we demonstrate that this scheme can be efficiently realized under the communication constraints of FL through our novel two-phase training procedure.

Dispatching the whole ensemble to clients raises privacy concerns.

$\rightarrow$ We agree that privacy is an important concern. We would like to note however that FENS still does not require the clients to transfer their local data but only their fully trained models. Exploring how to ensure that these model exchanges do not leak information about the data is of independent interest, and could have applications in various FL contexts. While we leave an in-depth exploration for future work, we believe that existing privacy schemes could be directly incorporated into FENS. For instance, local model training can be done using differentially private SGD [2], thereby ensuring differential privacy of local data despite the sharing of the local models. Similarly, the aggregator training can be conducted via a differentially private FL algorithm [3]. We will include this discussion in our revised version.

[1] David H. Wolpert. Stacked generalization. Neural Networks, 5(2):241–259, 1992.

[2] Abadi, Martin, et al. "Deep learning with differential privacy." Proceedings of the 2016 ACM SIGSAC conference on computer and communications security. 2016.

[3] Noble, Maxence, et al. "Differentially Private Federated Learning on Heterogeneous Data". Proceedings of the 25th International Conference on Artificial Intelligence and Statistics (AISTATS). 2022.

We thank the reviewer for their positive assessment of our work and the suggestions for improvements. We address the reviewer's questions below.

Q. Why does the MoE approach to aggregation require such a large communication overhead? Can't the size of the dense layers used for routing be tuned?

$\rightarrow$ Yes, we considered tuning the size of the routing network. We had experimented with models of various sizes, and observed that smaller models cause a deterioration in accuracy. Our current setup employs a CNN with two convolutional blocks, followed by two fully connected layers and a final classification head (Appendix B.4). This configuration achieves a reasonable balance between accuracy and communication overhead. Since the routing network needs to process images as input, it must be a reasonably sized CNN. Consequently, this results in a larger communication overhead compared to using a shallow MLP-based aggregator (referred to as the NN aggregator in Fig. 7), which processes concatenated logits.

Q. What happens if you apply a method like FedAdam to a model whose total size is comparable to the total FENS model size?

$\rightarrow$ This is an interesting ablation. We evaluated this suggestion by downsizing client local models such that the total size of FENS approximately matches the size of the single model in FedAdam. Preliminary results indicate that under a comparable memory footprint, FENS still remains competitive with the original FedAdam. On the other hand, using the downsized model for FedAdam induces an accuracy drop. We will include the complete ablation in the revised version of our paper.

Q. How do the results in Figures 2 and 4 change if we instead consider total client computation instead?

$\rightarrow$ There is indeed a trade-off here and we expect the client computation cost of FENS to be the highest in Fig. 2. This can be seen as the cost of FENS in providing OFL-like communication efficiency and FL-like accuracy which are its primary metrics of interest.

On the limitations of synthetically partitioned and vision datasets.

$\rightarrow$ We appreciate the reviewer’s recognition of our use of the FLamby benchmark. We are keen on extending our empirical analysis with a language dataset and will incorporate this in the revised version of the paper.

On trimming Sections 1 and 2 and including other ablations.

$\rightarrow$ Thank you for the suggestion. We will shrink the first two sections to enable the inclusion of Appendix C and above ablation in the main body of our revised paper.

On expanding the discussion on privacy.

$\rightarrow$ Privacy is indeed an important concern and we believe that existing privacy schemes could be directly incorporated into FENS. One approach could be to perform client local training using differentially private SGD [1]. Thus each client can provide a differentially private local model for the ensemble. Similarly, the aggregator training could be conducted via a differentially private FL algorithm [2]. As suggested by the reviewer, we will include a discussion of potential directions in our revised paper.

[1] Abadi, Martin, et al. "Deep learning with differential privacy." Proceedings of the 2016 ACM SIGSAC conference on computer and communications security. 2016.

[2] Noble, Maxence, et al. "Differentially Private Federated Learning on Heterogeneous Data". Proceedings of the 25th International Conference on Artificial Intelligence and Statistics (AISTATS). 2022.

We thank the reviewer for the feedback and insightful comments. We noticed that there was a confusion regarding the usage of public dataset in our algorithm. We also noticed that most of the identified weaknesses seem to arise from this confusion. Below, we clarify our algorithm design and address the reviewer's questions.

On the usage of public dataset at the server.

$\rightarrow$ Our method does not rely on any public dataset. In FENS, the data remains decentralized throughout. In the second phase of the algorithm, FENS trains the shallow aggregator model via standard FL (lines 131-139). We are happy to provide further details or address any additional questions.

On the lack of strong baselines such as Fed-ET [1] and FedGKT [2].

$\rightarrow$ In light of the above clarification, we do not consider Fed-ET [1] since it is an iterative FL method that relies on the usage of public data. Similarly, we do not compare to FedGKT [2] since it is an iterative algorithm that focuses on reducing the computational burden of edge clients. Hence, the objectives and design of the suggested baselines differ from that of FENS. In particular, FENS focuses on simultaneously achieving OFL-like communication efficiency and FL-like accuracy. Accordingly, we have considered 5 one-shot baselines and 6 iterative FL baselines that are more relevant to FENS. We will further clarify our choices in the paper.

Q. I am curious about how "Revisiting" is discussed in the paper. I do not see any revisiting part in the writing.

$\rightarrow$ Ensembling has been studied in OFL prior to our work, and we revisit this concept by applying it in a different and more powerful way. Unlike existing methods, FENS employs a shallow aggregator model atop locally trained client models, thereby generalizing traditional aggregation functions. By introducing a novel two-phase training procedure, which keeps data decentralized, FENS significantly improves accuracy while maintaining low communication costs. We will further clarify the "revisiting" aspect in our paper.

Q. I suggest the author to list a pseudo-algorithm in the Section-2 to let the audience better understand the proposed method.

$\rightarrow$ Thank you for the suggestion. We will do so in the final version of the paper.

[1] Cho, Yae Jee et al. “Heterogeneous Ensemble Knowledge Transfer for Training Large Models in Federated Learning.” International Joint Conference on Artificial Intelligence (2022).

[2] He, Chaoyang et al. “Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge.” arXiv: Learning (2020): n. pag.