FuseFL: One-Shot Federated Learning through the Lens of Causality with Progressive Model Fusion

We would like to thank the reviewer for the time in reviewing. We appreciate that you find the proposed algorithm novel and well-motivated, the experiments are extensive and the FuseFL is effective. Please see our detailed feedback for your concerns below.

Q1: Few-shot FL or one-shot FL.

Since FuseFL still communicates multiple rounds, I think it belongs to few-shot FL instead of one-shot FL. Besides, FedAvg (or other FL algorithms) with the same number of communication rounds should be compared as baselines to ensure a fair comparison.

Ans for Q1): Thanks for your important comments. Yes, FuselFL still communicates with a few rounds. In this sense, FuseFL belongs to the few-shot FL. However, the communication cost of FuseFL is as same as other OFL methods, which is the main claim in the introduction (extremely low communication costs). We will highlight this difference in the introduction of the revised paper.

Furthermore, we conduct new experiments (a=0.1) with 2 and 4 communication rounds of baseline methods with their multi-round versions [2,8] to provide a fair comparison from the round perspective. Note that the FedAvg has communication rounds of 10 for all experiments. Results show that for the same communication rounds, FuseFL still achieves the best performance than other baselines.

Q2: Experiment Issues.

For example, in a previous FL benchmark on non-IID data [1], the quantity-based label skew settings are challenging, which are not tested in the current version. Also, in the section of scalability, the authors only test up to 50 clients. I suggest the authors to test on even more clients, e.g. 500 clients, to validate the scalability.

Ans for Q2): Thanks for your valuable comments. The quantity-based label skew is another kind of heterogeneity different from the Dirichlet sampling based Non-IID simulation. Following your suggestions, we conduct new experiments with the quantity-based label skew of #C=2 and 10 clients. The test accuracies are shown as follows, illustrating that the #C=2 is a harder setting than Dirichlet with $a=0.1$. While FuseFL can outperform other baseline methods.

To validate the scalability of FuseFL, we follow the scale of FL in [2,3] ($M=5\sim 50$ clients) in original paper. And most OFL methods simulate with less than 50 clients with ResNet and 100 with basic machine learning models [5,6,7]. Simulating 500 clients within the OFL methods is challenging as 500 local models need to be saved to conduct knowledge distillation or ensembling. Loading 500 ResNets into GPU memory is difficult, or offloading and switching between GPU and CPU memory requires much time. To this end, we tried our best to simulate 100 clients to verify the scalability as follows. Results show that FuseFL outperforms all baseline methods except for the Ensemble, which requires significant storage and computation overheads.

Reference

[1] Li, Qinbin, et al. "Federated learning on non-iid data silos: An experimental study." 2022 IEEE 38th international conference on data engineering (ICDE). IEEE, 2022.

[2] J. Zhang, C. Chen, B. Li, L. Lyu, S. Wu, S. Ding, C. Shen, and C. Wu. Dense: Data-free one-shot federated In NeurIPS 2022.

[3] R. Dai, Y. Zhang, A. Li, T. Liu, X. Yang, and B. Han. Enhancing one-shot federated learning through data and ensemble co-boosting. In The Twelfth International Conference on Learning Representations, 2024.

[4] Q. Li, B. He, and D. Song. Practical one-shot federated learning for cross-silo setting. In IJCAI 2021.

[5] C. E. Heinbaugh, E. Luz-Ricca, and H. Shao. Data-free one-shot federated learning under very high statistical heterogeneity. In The Eleventh International Conference on Learning Representations, 2023.

[6] Jhunjhunwala, D., Wang, S. and Joshi, G., 2024, April. FedFisher: Leveraging Fisher Information for One-Shot Federated Learning. In the International Conference on Artificial Intelligence and Statistics (pp. 1612-1620). PMLR.

[7] Towards Addressing Label Skews in One-Shot Federated Learning. In ICLR 2023.

[8] Ensemble distillation for robust model fusion in federated learning. In NeurIPS 2020.

We would like to thank the reviewer for the time in reviewing. We appreciate that you find the problem interesting, important, and practical, our idea is clear and performance is great. Please see our detailed feedback for your concerns below.

Q1: Experiment Issues.

Important baseline [1] and higher heterogeneity levels.

Ans for Q1): Thanks for your important comments. We have conducted new experiments of CoBoosting [1] and FuseFL with higher heterogeneity levels (a=0.05). The new results are provided following and added in our revision. Results show that FuseFL outperforms CoBoosting and provides higher improvement than the baselines. Note that for CIFAR-10, Ensemble outperforms all baselines but requires extremely large storage and computational costs.

And we provide following discussion of differences between FuseFL with [1,2,3].

• Knowledge distillation based FL: Both CoBoosting [1] and FEDCVAE-KD [2] focus on exploiting knowledge distillation methods to improve the global model performance, while our method focuses on how to aggregate the models together. Thus, the knowledge distillation is orthogonal with our method and may be utilized to enhance FuseFL. For example, one can consider running FuseFL first to obtain a fused global model. Then, this model can be used to conduct knowledge distillation to guide the local model training with the FuseFL once again.
• Average-based FL: FedFisher [3] focuses on how to better average models to obtain a better global model. The fisher information is utilized to identify element-wise averaging weights of parameters in different local models. FedFisher outperforms previous OFL methods [4]. However, FuseFL focuses on a different methodology that concatenates model blocks together instead of averaging. Nevertheless, we consider utilizing the fisher information and average-based methods to enhance FuseFL.

More differences between FuseFL and Average-based [3,4] and Knowledge distillation [1,2] based OFL can be referred in Table 1 and Appendix C in original paper.

Q2: Security issues.

By introducing multiple rounds of communication, FuseFL increases vulnerability to attacks.

Ans for Q2): Thanks for your important comments. Yes, in this work we do not explicitly consider the security issue. However, the vulnerability to attacks of FuseFL will not be higher than previous multi-round FedAvg, which requires many more communication rounds to achieve the same model performance as FuseFL. For example, FedAvg may require more than 100 rounds to achieve 70% test accuracy as FuseFL with $2\sim 4$ rounds, which introduces more communicated information and a higher possibility of attacks. And FuseFL has the same communication size with other OFL methods.

Nevertheless, while the communication size is not increased, FuseFL indeed increases the communication rounds than other OFL methods. The possible higher risk of attack is injecting adversarial modules. The possible solution is to detect and reject such malicious uploading also through the lens of causality.

For the model inversion or membership attack, we can consider adding differential privacy or using noised samples with invariant features while ensuring the model performance.

Q3: Extra computation overheads.

Similarly, FuseFL appears to increase client-side overheads by letting clients run training K times instead of once as in standard OFL. While the computation cost appears to decrease with K, it is nevertheless higher than client-side overheads in standard OFL approaches which mostly add server-side overheads.

Ans for Q3): Thanks for your valuable comments. Sorry for our missed experiment descriptions of FuseFL. We have considered this extra computational problem of repeated local training. To address this problem, we decrease the local training epochs of each client by $K$ times. For example, the FedAvg and other OFL methods require 200 epochs. With $K=4$, for training and merging each block in each round of FuseFL, the local training epochs are decreased to 50. Thus, the number of total training epochs of FuseFL is the same as other OFL methods. The motivation of this design can be referred to the progressive freezing during training DNNS [5]. We have added the descriptions of this design in our revision.

Reference

[1] R. Dai, Y. Zhang, A. Li, T. Liu, X. Yang, and B. Han. Enhancing one-shot federated learning through data and ensemble co-boosting. In The Twelfth International Conference on Learning Representations, 2024.

[2] C. E. Heinbaugh, E. Luz-Ricca, and H. Shao. Data-free one-shot federated learning under very high statistical heterogeneity. In The Eleventh International Conference on Learning Representations, 2023.

[3] Jhunjhunwala, D., Wang, S. and Joshi, G., 2024, April. FedFisher: Leveraging Fisher Information for One-Shot Federated Learning. In the International Conference on Artificial Intelligence and Statistics (pp. 1612-1620). PMLR.

[4] Q. Li, B. He, and D. Song. Practical one-shot federated learning for cross-silo setting. In IJCAI 2021.

[5] M. Raghu, J. Gilmer, J. Yosinski, and J. Sohl-Dickstein. Svcca: singular vector canonical correlation analysis for deep learning dynamics and interpretability. In NeurIPS 2017.

We would like to thank the reviewer for the time in reviewing. We appreciate that you find the problem is interesting, important and practical, our idea is clear and performance is great. Please see our detailed feedback for your concerns below. Some answers can be referred to the global response due to the limited space.

W1 & Q2: The cost of the multiple-block communication and the disconnection.

Ans for W1 & Q2): Thanks for indicating this practical problem. The synchronous waiting problem also exists in the traditional FL methods like FedAvg and others. Given the number of clients $M$, and model size $S$, the traditional FL methods require communication costs $MSR$, in which the $R$ is the communication round, normally ranging from 100 to 1000. However, our method only requires $MS$ communication cost. For the communication rounds in FuseFL, we can communicate more communicating more layers to reduce the number of blocks (like different $K=[2,4,8]$ in Table 2 and 4 in original paper).

Consider the disconnection, FuseFL needs to wait for the clients. This also happens in the traditional FL. One possible solution is to continue the training process of other clients. After reconnecting, the client loads the newest global model and continue training.

W2 & Q1: The explanation part.

Ans for W2 & Q1): Thanks for your important comments. We will revise the paper to make the causality explanation more clear. Specifically, we will rewrite lemma 3.1 to show how the locally learned spurious $R^{spu}_m$ features influence the mutual information between $H^k_{loc}$ and global $Y$. The "block" means a single or some continuous layers in a DNN. For example, a single or two consecutive convolutional layers in a CNN can be seen as a block. Features refer to the outputs of blocks. For simplicity, the $H^k_m$ represents both block $k$ on client $m$ and features output from it as shown in Figure 1 in original paper. Also, please refer to more explanation of Q1 of Reviewer #mHC6 and the global response.

W3 & Q4: Large-scale models like LLMs.

Ans for W3 & Q4): Thanks for your important comments. To deploy transformer-based frameworks like current LLMs with the core idea of FuseFL in one-shot FL, we envision two methods here.

• Concat-and-freeze: similar to training ResNet in FuseFL, we can block-wisely train and collect the transformer blocks together for each round; during local training, the output features of all transformer blocks are concatenated to feed into the subsequent layers. Due to the large resource consumption of pretraining, we do not evaluate this idea here.

• Averaging-and-freeze LoRA: here we consider the finetuning scenarios with LoRA [1]. LoRA blocks can be seen as additional matrix mapping applied on the local Q V attentions and MLP layers. The output is the original feature plus the LoRA output. To use LoRA in FuseFL, we can follow the MoE style [2] or the averaging style [1]. Specifically, we consider averaging LoRAs on different clients together, then averaging and freezing all LoRAs in each transformer block to freeze the obtained aggregated features in each communication round. The following Table shows performance of finetuning Llama2-7B with FuseFL with few-shots FedAvg with OpenFedLLM [3] benchmark (20 clients with Alpaca-GPT4 and MedAlpaca). Results show that FuseFL with much fewer communication rounds can outperform FedAvg of 50 communication rounds.

Q3: Error bars.

Ans for W3 & Q4): Thanks for your important comments. We have added the error bars of main experiments as follows and revised the paper. We report a part of them here due to the limited space.

Reference

[1] LoraHub: Efficient Cross-Task Generalization via Dynamic LoRA Composition. In COLM 2024.

[2] LoRAMoE: Alleviate World Knowledge Forgetting in Large Language Models via MoE-Style Plugin. In ACL 2024.

[3] OpenFedLLM: Training Large Language Models on Decentralized Private Data via Federated Learning. In KDD 2024.

We would like to thank the reviewer for taking the time to review our work. We appreciate that you find the solved problem, our observation, and our core idea are very interesting. Please see our detailed feedback for your concerns below.

Q1: Clarity of theoretical analysis.

Rspu is something that has information about Y even though they are not causally related, , e.g., in Waterbirds dataset, many waterbirds are photographed in front of water background. So, water background (B) naturally has a correlation with bird label (Y) and is not nuisance. But, it is not desirable to use B in the model since it performs poorly on minority groups. Here B is Rspu but not Nuisance. The authors also write I(Rspu;H)=0 means spurious feature is not being used. This would also hold for nuisance but not spurious feature as in the example above. One wants the model to not mechanically use the background B. But, what it actually uses instead of the background can still have correlations with the model outputs, and hence mutual information can still be there. I think what you call Rspu is actually Nuisance at a global level. You actually assume that for the global dataset, Rspu satisfies the definition of nuisance, i.e., being independent with Y, but is dependent at a local level?

Ans for Q1): Thanks for your important comments on the theoretical analysis. Yes, here we call Rspu is actually Nuisance at a global level (global dataset across all clients). Thus, on one local client $m$, the spurious feature $R^{spu}_m$ is independent with the $Y^{global}$ in the global dataset, but dependent on the $Y^{m}$ in the local dataset $\mathcal{D}_m$. I(Hk; Rspu)=0 is a sufficient condition for not overfiting to spurious correlations. Because when I(Hk; Rspu)=0 is not satisfied and Hk takes some spurious features, one can adjust the final classifier to avoid spurious fitting [1,2,3]. Besides, our analysis in Section 3.2 and the Empirical study in Section 3.3 aligns with your understanding.

We will revise the paper and include the section of clear definitions. Specifically, we will include the definitions of the local $X_m$, $Y_m$, global $X^{global}$, $Y^{global}$, and the local spurious features $R^{spu}_m$. Furthermore, we will rewrite Lemma 3.1 to make it clear in the federated learning scenario.

Q2: Experiment issues.

In spurious correlation papers, one is often interested in worst-group/minority group accuracy, i.e., waterbirds on land background. Overall accuracy even decreases sometimes. Could you comment on this in the context of your experiments? How are the local accuracies?

Ans for Q2): Thanks for your important questions on the worst-group accuracy. In FL, the local accuracies refer to the accuracy on the local datasets on clients, which is common in personalized FL [4]. Specifically, local test datasets have the similar data distributions to the training datasets (according to the ratio of the labels). One goal of designing experiments with the backdoored Fl CIFAR-10 datasets is to study the majority/minority accuracy. For a backdoored client, images with local spurious shapes and colors become majority group on it (Figure 5 in original paper), while the normal images become minority groups. Now we report the local/global accuracy here. Backdoored (BD) clients fit on the handcrafted spurious features thus having lower global accuracy than normal clients.

Q3: Overreliance on causality literature.

There is an overreliance on causality literature including citing Pearl. However, I feel the strategy and the novelty are not that reliant on causality.

Ans for Q3): Thanks for your valuable comments. Yes, our strategy does not explicitly adopt the causal discovery to improve one-shot FL. Our motivation is to interpret the mechanism and the root causes of why previous one-shot FL fails through the lens of causality. We appreciate your kind suggestion and will highlight the main contribution in the revision.

Q4: Unclear fusion strategy.

Everything points to one equation (5) which is not described clearly.

Ans for Q4): Thanks for your question. Eq.(5) describes the chain structure of layers in DNNs. For example, a module $H_{12} = H_1 \circ H_2$, we have the $H_{12}(x) = H_1(H_2(x))$.

Reference [1] Class-Balanced Loss Based on Effective Number of Samples. In CVPR 2019.

[2] Learning debiased classifier with biased committee. In NeurIPS 2022.

[3] Towards last-layer retraining for group robustness with fewer annotations. In NeurIPS 2024.

[4] On Bridging Generic and Personalized Federated Learning for Image Classification. In ICLR 2022.

Dear Reviewer #mHC6,

Thanks for providing these related works [1,2]. Removing local biased feature is an interesting idea to enhance the global fairness. We will add these discussions and related works into our revision. Also, we will modify the paper to highlight the information bottleneck/mutual information and the title. Thanks a lot for your suggestions.

Best regards and thanks,

Reference

[1] Yahya H Ezzeldin, Shen Yan, Chaoyang He, Emilio Ferrara, and A Salman Avestimehr. Fairfed: Enabling group fairness in federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 7494–7502, 2023.

[2] F. Hamman and S. Dutta, "Demystifying Local & Global Fairness Trade-offs in Federated Learning Using Partial Information Decomposition,” International Conference on Learning Representations (ICLR 2024).