Federated Instruction Tuning of LLMs with Domain Coverage Augmentation

Authors: Zezhou Wang, Yaxin Du, Zhuzhong Qian, Siheng Chen
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Federated Learning, Large Language Model, Instruction Tuning

Abstract:
Federated Domain-specific Instruction Tuning (FedDIT) utilizes limited cross-client private data together with server-side public data for instruction augmentation, ultimately boosting model performance within specific domains. To date, the factors affecting FedDIT remain unclear, and existing instruction augmentation methods primarily focus on the centralized setting without considering distributed environments. Our experiments reveal that the cross-client domain coverage, rather than data heterogeneity, drives model performance in FedDIT. In response, we propose FedDCA, which optimizes domain coverage through greedy client center selection and retrieval-based augmentation. For client-side computational efficiency and system scalability, FedDCA$^*$, the variant of FedDCA, utilizes heterogeneous encoders with server-side feature alignment. Extensive experiments across four distinct domains (code, medical, financial, and mathematical) substantiate the effectiveness of both methods. Additionally, we investigate privacy preservation against memory extraction attacks utilizing various amounts of public data. Results show that there is no significant correlation between the volume of public data and the privacy-preserving capability. However, as the fine-tuning rounds increase, the risk of privacy leakage reduces or converges.

Decision: reject
Reviewer scores: 6, 5, 3, 3

Reviews and discussion:
<thread>
creation date:
1730036083881

thread_type:
official_review

<thread_object>
creator: reviewer_8vfk

<content>
summary:
This paper proposes a federated instruction tuning approach. The authors assume the existence of a large public dataset that covers all domain data, while each domain only has a small subset of data. The paper introduces multiple techniques, including greedy client center selection and domain data retrieval, to achieve data selection from the public dataset and improve the performance of federated instruction tuning. Furthermore, a heterogeneous encoder approach is introduced to further enhance efficiency. The experiments demonstrate that this data selection approach improves performance compared to client-data-only fine-tuning baselines.

strengths:
1. This approach highlights domain coverage as the main bottleneck for federated instruction tuning.

weaknesses:
1. The finding that motivates this research—"there is no linear correlation between the degree of non-independent and identically distributed (IID) and LLM performance in the context of FedDIT"—requires further evidence and rephrasing to avoid being misleading. First, "linear correlation" is a very strong term and may be inappropriate here; the authors likely mean "monotonicity" instead of "linearity." Second, the current experiments are insufficient to draw definitive conclusions. The results are not conducted over multiple rounds, nor do they report the mean and standard deviation of performance, which makes them highly dependent on the seed, especially in the Dirichlet instance split. Third, the number of clients (10, with 2 clients involved in each round) is relatively small and may not be significantly affected by \$α\$. Fourth, averaging performance across multiple datasets or tasks is problematic. For instance, in Table 1, Method 2 is clearly better than Method 1 on multiple datasets, but the average shows the opposite.

**Table 1. Performance of Method 1 and Method 2 on Binary Classification Tasks**

| Method  | Method 1  | Method 2 |
| ------- | --------- | -------- |
| Data 1  | 98%       | **99%**  |
| Data 2  | 98%       | **99%**  |
| Data 3  | 98%       | **99%**  |
| Data 4  | 98%       | **99%**  |
| Data 5  | **52%**   | 47%      |
| Average | **88.8%** | 88.6%    |

2. The assumption that there exists a public dataset requires further clarification regarding its prevalence in practice. Clarification is needed on whether there are real-world applications that align with this assumption. It is also necessary to determine how much of the improvement is due to the additional data versus the algorithmic design. Furthermore, it should be clarified whether the public multi-domain dataset needs to be within the domain of the clients, or if it can be out-of-domain data.

3. The method of retrieving similar public instructions to domain instructions requires the public dataset to already cover all domains (i.e., contain full domain knowledge). However, it is unclear which baseline uses this large, powerful public dataset for fine-tuning. It would be interesting to see the results of:

- Fine-tuning with the same public dataset.
- Fine-tuning with the public dataset followed by further federated fine-tuning with private clients' local dataset.

Including these baselines would significantly enhance the convincing nature of the experiments.

4. The experimental setting is unrealistic. In line 996, the authors randomly select 1,000 samples from the in-domain instruction set and split them into 10 shards as each client's local dataset, leaving the remaining 20x to 200x samples for the public dataset. The necessity of federated learning is questionable if such a large, independent and identically distributed (IID), noise-free public dataset exists.

**Minor Comments**

1. The format of Table 1 is difficult to read. It is unclear which entries represent methods and which represent metrics, as well as how to interpret the different metrics.

questions:
1. Does random sampling baseline make a sample of each client's domain data? In Table 1, random sampling generally outperforms full-data fine-tuning. However, the introduction claims that FedIT requires a sufficient amount of instruction data. If this is true, why does down-sampling lead to an increase in average performance? More evidence is needed to support the motivation.

2. In "Results demonstrate that domain coverage significantly influences model performance in the corresponding domain," is this conclusion drawn from the experiments in this manuscript or from Wang et al. (2023)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730366650792

thread_type:
official_review

<thread_object>
creator: reviewer_tyyt

<content>
summary:
This work proposes a federated domain-specific instruction tuning method to enable federated clients to utilize data from the server to augment their local data. The objective of this work is meaningful, as FL does indeed face the issue of insufficient data on the client-side. However, the assumption that there exists a large dataset on the server encompassing all domains seems unrealistic. If such a dataset is available on the server, why would there still be a need for FL to tune LLMs? Additionally, compared to traditional FL, this approach poses a risk of exposing additional information from the client-side data (such as domain information) to the server. In summary, I believe this work requires further refinement.

strengths:
1. This manuscript is well-written and well-organized, making this work easy to understand.
2. Extensive experiment evaluations are provided.

weaknesses:
1. The main distinction between the proposed method and previous approaches is its ability to leverage server-side data, with the effectiveness of the method relying on the server holding a dataset that encompasses multiple domains. In fact, many LLM tuning methods based on FL aim to explore the possibility of tuning LLMs with client-side data when centrally collected data has been exhausted. If there is already a dataset on the server that encompasses multiple domains, why would we still need FL to tune the model? Wouldn't it be better to fine-tune the LLM directly on the central server?
2. In the experiments, it is mentioned that "each client has 100 local instructions and obtains 5000 augmented public instructions from the server." (line 885). In this case, the majority of the data is contributed by the server, and the significance of introducing FL seems unclear.
3. This work builds upon the original FL paradigm by providing additional client-side information, specifically client-side cluster centers. Although the original data remains on the client side, it is inappropriate to directly claim that "it does not violate the client’s privacy." Such a kind of information may reveal the distribution of client-side data, thereby exposing client data to privacy leakage threats such as membership inference attacks, which do not require direct access to the original data. As FedDCA describes as a distinguishing feature from other existing methods, it sends data from the server's public dataset that are related to a specific client to that client as augmentation. Doesn’t this imply that the server can determine the data domains of that client to some extent?
4. One of the claimed contribution, "no linear correlation between the degree of heterogeneity and model performance" has already been revealed by one of the cited works [1], which weakens the contribution of this work to a certain extent.
5. To the best of my knowledge, the appropriate citation of FedIT should be [2], which has been proposed in 2023, rather than (Ye et al., 2024b;a) in Line 764.
6. The codes are not available, making the reproducibility of this work difficult to assess.

[1] FederatedScope-LLM: A Comprehensive Package for Fine-tuning Large Language Models in Federated Learning

[2] Towards Building the Federated GPT: Federated Instruction Tuning.

questions:
1. How are the four baselines based on FedIT implemented? Do they involve directly tuning and transmitting all parameters, or only PEFT adapters?
2. How is the weight of the projector in FedDIT* obtained? It seems that many details have not been provided in Appendix A.4.
3. Please refer to Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730637748443

thread_type:
official_review

<thread_object>
creator: reviewer_29fa

<content>
summary:
The paper proposes a new method for federate instruction fine tuning for LLMs.
The main idea involves around using cluster center info coming from clients to select publicly available data stored on the server to augment the local data to achieve better coverage.

strengths:
+ The paper is motivated well.

+ The paper is easy to read.

+ Provides a basic clustering method that could improve the federated fine tuning process.

weaknesses:
- To me, the main question I have is that whether FedAvg and variants developed for dealing with Non-IID data combined with simple random sampling of the public data could achieve similar results.   The main reason is that Federated learning by itself is not sending data to the clients and just aggregates the client updates. Therefore, it is not surprising to me that sending more data to clients help even if this is done carefully using clustering as suggested in the paper. Therefore, how about sending some random data to clients and then use standard FL techniques especially the ones developed for non-iid settings.    Hence, I highly recommend the authors to explore this setting. One initial experiment could be to send the random data (e.g., the same amount sent using the clustering approach) to each client and then use FedAvg to see the performance. 

- The privacy experiments are good but not itself it is not adequate ( I understand that this is not the focus of this paper and I did not use this comment in any negative way in my ranking), I think a setting where noise is added during aggregation (e.g., something of differential privacy sense) can be evaluated. As a starting point, Gaussian noise could be added to achieve \epsilon-\delta differential privacy. Again this could be considered in the future work.

questions:
Please see my first comment in the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730695511914

thread_type:
official_review

<thread_object>
creator: reviewer_4a8x

<content>
summary:
This paper introduces FedDCA, a novel approach to federated learning for domain-specific LLM instruction tuning that focusses on coverage, as opposed to heterogeneity (non-iid ness). Instead of worrying about how different each client's data is from others (data heterogeneity), the authors show that what really matters is how well the combined data from all clients covers the target domain. This insight leads to their proposed method which uses client selection and data augmentation strategies to maximize domain coverage.

The authors demonstrate their approach across four domains (code, medical, financial, and mathematical), showing significant improvements over existing methods. They also introduce a more computationally efficient variant called FedDCA* that uses different-sized encoders for clients and servers. While the paper lacks some theoretical guarantees and has limited privacy analysis, it presents a promising new direction for federated learning with LLMs by focusing on domain coverage rather than trying to manage data heterogeneity.

The authors do not provide any additional methods to guarantee privacy (no DP), though it is common in FL papers, also they mount extraction attacks to measure memorization.

strengths:
1. novel idea and surprising insight, regarding coverage vs. heterogeneity.

2. Extensive and thorough experimentation

weaknesses:
1. There is one main concern I have, maybe because I have not been following FL literature recently, but it seems like the baselines are all from 2020! I wonder why, is the setup not realistic? The paper cites Ye et al. 2024 (openfedllm) which is recent, but they only compare against in table 3, why not table 2? is it because the FedIT method there is the backbone to this method?


2. Minor: seems like clients are limited to 10, which is a bit low. Also I wonder how DP would come into play, but this is not necessary to study.

questions:
Please respond to the weaknesses above. 

Also

How do you isolate whether the improvements come from your domain coverage method versus simply having more training data? Did you control for total training data size when comparing methods?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732085558032

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We have noticed that some reviewers have concerns about the setting of FedDCA and whether to use public datasets for data augmentation. Thus, we provide a unified response here, hoping to help you better understand and agree on the necessity of public datasets.

---

A public dataset is an abstraction of open-source datasets available on the internet. In the context of this paper, the presence of a public dataset on the server is only one possible scenario. The core innovation of this work lies in maximizing domain coverage, and the proposed algorithm is independent of the presence of a public dataset on the server. Even if the server does not have a public dataset, clients can upload their cluster centers to the server, which selects a set of client centers and sends them back to the clients. Each client can then retrieve data based on the received client centers, thereby achieving data augmentation that maximizes domain coverage. This approach was chosen to simplify the setting, much like the first step in solving an equation is simplification, allowing us to focus more on the algorithm itself.

I believe the reason why federated domain instruction fine-tuning requires data augmentation is to prevent potential performance degradation from using only local data and to enhance the model's generalization capability. I don't think the value of local data will be overshadowed by public data. 

Regarding whether the public dataset must include in-domain data, first of all, the concept of "domain" is flexible and hierarchical. Currently, there is no precise definition of a domain. For example, in Explore-Instruct [1], "Brainstorming" and "Rewriting" are considered two domains, but they could also be regarded as two tasks under the general domain.

Secondly, even when using the clearer domain classification adopted in this paper (e.g., code, medical, financial, math), if the client aims to solve tasks based on existing knowledge, the public dataset will inevitably contain knowledge relevant to those domains. This could come from the original corpus (which can be converted into instruction-response pairs using GPT) or from pre-constructed instruction datasets. So the public dataset contains in-domain but out-of-task is a more reasonable and intuitive scenario.

Consider a scenario where clients hold data for solving task $t_1$, which belongs to domain $\mathcal{D}$. The public dataset includes domain $\mathcal{D}$, containing a set of tasks denoted as $T=\{t_2,\dots,t_N\}$, which does not include $t_1$. The motivation for the client to participate in federated learning is to gain additional benefits compared to purely local training; therefore, the test set should also consist of data from domain $\mathcal{D}$ and task $t_1$. In this scenario, if we do not use local data and merely train the LLM by randomly sampling from the public data, intuitively, its performance is inferior to using both local data and a portion of the sampled public data.

To further clarify, we have conducted additional experiments based on this out-of-task scenario. Considering this setting in the financial domain, to construct the out-of-task setting, given that the training set FinGPT and the test sets FPB, FiQA, and TFNS are all related to sentiment analysis tasks, we choose to continue using FPB, FiQA, and TFNS as the test sets. This means that the clients aim to train a model adept at performing sentiment analysis through federated learning. Meanwhile, the FinGPT's instructions in public data are replaced with data from the $\texttt{Sujet-Finance-Instruct-177k}$ dataset where `task_type="qa"`. The local data of the clients is still randomly sampled from FinGPT. This approach yields out-of-task public data. The following experimental results show the performance on the test set after 30 rounds of training, using Accuracy (%) as the metric.

[**Table R1**. Performance on the test set after 30 rounds of training. Random is short for Random Sampling.]
| Method | FPB | FiQA | TFNS |
|:------------:|:------------:|:------------:|:------------:|
| Zero Shot | 55.94 | 18.54 | 59.21 |
| Base Data | 58.25 | 14.18 | 66.62 |
| Random | 60.39 | 9.45 | 65.45 |
| **FedDCA** | **60.89** | **18.91** | **67.37** |

As shown in Table R1, it can be observed that even when the public data is out-of-task, FedDCA still achieves performance improvements compared to other baselines. Additionally, using the Random Sampling data augmentation strategy resulted in performance degradation on the FiQA dataset. This further underscores the necessity of selecting an appropriate data augmentation strategy.

[1] Wan, Fanqi, et al. "Explore-instruct: Enhancing domain-specific instruction coverage through active exploration." arXiv preprint arXiv:2310.09168 (2023).

---

We hope that our responses can fully address your concerns and will be grateful for any feedback.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732459284214

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Once again, we sincerely thank the reviewers for their valuable feedback and constructive suggestions. We have uploaded a revised version of our paper, with the changes highlighted in red for your convenience. The key modifications are as follows:

- In Section 1, we add a paragraph to briefly introduce FedDIT's setting and lead readers to Appendix A.2 for a more detailed explanation.
- We revised the description of the relationship between data heterogeneity and LLM performance, changing it from nonlinear to non-monotonic for greater precision and clarity.
- In Section 5, we add a table to highlight the key differences between FedDCA and other baselines.
- We modify the citation for FedIT, including both OpenFedLLM and Shepherd.
- In Appendix A.2, we further elaborate FedDIT's setting, including why this setting, the definition of domain and the distribution of public data.
- In Appendix A.3.1, we update the results of Table 5, which reports the average performance and the standard deviation in each domain.
- In Appendix A.6, we add the discussion of domain inference attack.
- In Appendix A.10, we perform the experiments on the held-out setting, where the public dataset does not contain the task that clients aim to solve.
- In Appendix A.10, we provide the experimental results when each client only samples 100 data from the server.
- In Appendix A.10, we provide the experimental results when the LLM is firstly fine-tuned on public data then performs FedDIT based on different settings.
- In Appendix A.10, we show the scalability of FedDCA.
- We adjust some of the figures and tables to make them more self-explanatory and to create a more compact layout for the article.

We deeply appreciate the reviewers' time and effort in evaluating this paper and hope that this revised version effectively addresses your concerns.
</content>
</thread_object>

</thread>


Metareview:
metareview:
## Summary: 
The paper presents a novel approach called Federated Domain Coverage Augmentation (FedDCA) for federated learning in the context of domain-specific instruction tuning for Large Language Models (LLMs). In the proposed approach, each client performs clustering of local instructions and sends the cluster centers to the server, while the server selects N (N is the number of clients) centers in a greedy manner to solve a maximum coverage problem. They assume that there is a public instruction tuning dataset and use dense retrieval to find the public instructions closest to the selected center for each client. The selected subset of the public instruction dataset is then sent to the client as an augmentation to its local private dataset for local training. The experiments conducted across four diverse domains (code, medical, financial, and mathematical) show improvement over existing federated instruction tuning methods. 

## Strengths:
1. The paper raises that domain coverage is important to FedIT and proposes an idea to address it
1. Extensive comparison to baseline approaches.

## Weaknesses: 
1. The assumption that there exists a public and diverse instruction-tuning dataset is not fully convincing to reviewers. If the data are of web-scale as the author claims in the response or can be easily accessed by each client, then the motivation of applying federated learning is questionable. If the public dataset is small, then it might be hard to guarantee its diversity and full coverage of all the domains. Moreover, whether federated learning is necessary and realistic in the paper-proposed setting is not well explained. Although the authors provide additional experiments and clarifications, after reading them, the reviewers and the meta-reviewer are still not fully convinced.
1. Reviewers are concerned about the privacy of the proposed method. The statement regarding privacy is not rigorous and only provides some vague intuitions without theoretical guarantees such as differential privacy. 
1. The greedy selection algorithm is a heuristic and why it targets and can approximately solve a maximum coverage problem (the coverage is not well defined as in Eq. 3) is not clear. A rigorous optimization problem formulation and a straightforward connection to the greedy algorithm need to be presented
1. Some other major claims in the paper are not rigorous and might be misleading, e.g., the linear correlation between the degree of heterogeneity and model performance. 
1. The original draft lacks experiments on different non-IID settings, >10-client setting, public dataset of comparable size as the local datasets, variance of the performance under different random seeds. Although the author provided partial results in the rebuttal, they did not cover the main scope of the experiments presented in the original draft. 
1. The math notations in the paper are difficult to understand and need further corrections and simplifications. 
1. FedDCA does not always show advantages, as reflected by the additional experiments during the discussion.  
1. It is necessary to compare FedDCA with existing clustered federated learning approaches (with or without using the public dataset). 

## Decision:
The authors provided detailed responses to the reviewers' concerns with additional experimental results and further clarifications. Authors and all the reviewers actively participate in the discussion. Although some concerns have been successfully addressed by the rebuttal, reviewers found that several major concerns above are still not resolved and may undermine the proposed setting and approach. The mete-reviewer carefully read all the discussion and responses, as well as the original paper. The meta-reviewer agrees with most reviewers that the paper is not ready for publication. 

Based on the review comments, the author rebuttal, and the author-reviewer discussion, and the final review ratings, I cannot recommend acceptance of this paper to ICLR. The authors are encouraged to address the remaining concerns of the reviewers, clarify/justify the motivations of the problem settings better, formulate the problem and derive the algorithms in a more rigorous manner, and provide a fine-grained ablation study to demonstrate the advantages of the proposed methods and submit it to the next conference.