Distributed In-Context Learning under Non-IID Among Clients

1. Task covered in study is limited to classification tasks. It would benefit from proving effectivenss across RAG and multihop reasoning.

Thank you for your thoughtful suggestion on inclusion of RAG & multihop reasoning tasks. Our main contribution in this paper is to demonstrate the feasibility of using ICL under non-IID conditions, and we focus on classification tasks as a starting point. While RAG & multihop reasoning are indeed more text-heavy and complex tasks, further experiments and exploration are needed to evaluate applicability of our method in those settings. This represents a promising direction for future work.

We also note that many prior ICL-related studies [1][2][3] focus on classification tasks, which have proven to provide valuable insights to the research community. We believe our work aligns with this tradition while laying foundation for broader applications.

[1] Lyu, Xinxi, et al. "Z-ICL: Zero-Shot In-Context Learning with Pseudo-Demonstrations." ACL. 2023.

[2] Yoo, Kang Min, et al. "Ground-Truth Labels Matter: A Deeper Look into Input-Label Demonstrations." EMNLP. 2022.

[3] Chen, Huiyao, et al. "Retrieval-style in-context learning for few-shot hierarchical text classification." TACL 2024.

2. Paraphrasing based method to secure privacy during data retrieval - seems to have limited evaluation and more robust evaluation against other privacy-preserving techniques needs to be done to support claim

Thank you for your insightful comment on evaluation of paraphrasing-based method. While paraphrasing is not main focus of our work, it serves to demonstrate that our framework can seamlessly integrate with plug-in privacy-preserving techs, which are orthogonal to current scope of our research. We chose paraphrasing as an example because prior studies [1][2][3] have successfully employed it for privacy preservation in LLM research.

[1] ​​Zhang, Z., Zhang, J., Huang, J., Qu, L., Zhang, H., & Xu, Z. (2024). Fedpit: Towards privacy-preserving and few-shot federated instruction tuning. arXiv preprint arXiv:2403.06131.

[2] Krishna, Kalpesh, et al. "Paraphrasing evades detectors of ai-generated text, but retrieval is an effective defense." NeurIPS (2024).

[3] Yadav, V., Tang, Z., & Srinivasan, V. (2024, July). Pag-llm: Paraphrase and aggregate with large language models for minimizing intent classification errors. 47th ACM SIGIR.

3. Even when sharing similar task, clients might have different prompts in structure, length. Will this affect retrieval effectiveness since it heavily depends on similarity between query and training examples?

Thanks for meaningful comment. We conducted additional experiments for Non-IID with feature-skew (same label distribution but different query distribution) to address this suggestions.

We designed NonIID where 1 client contains only Yelp training samples, while another client contains only Amazon training samples. Yelp & Amazon share same label space (5-class classification), while they show different distributions on queries. On server, we use both Amazon&Yelp samples for test set, and perform t-SNE on test embedding with their budget values on each client. By this setting, we want to show our method intuition on ‘budget value’ and ‘sample embeddings’ still holds even it is feature skew with class-balance.

As shown in [client 1 (with only Yelp)] and [client 2 (with only Amazon)], clustering pattern is even more significant than previous class-based nonIID, indicating our claim still holds under text style non-IID. Also, we found Amazon test samples tend to assign all budget to Amazon client, while Yelp samples tend to assign all budget to Yelp client.

To conclude, our method can also be applied to feature skew (style shifting) non-IID.

4. Paper references federated learning in related works, comparison with FL methods in this setting?

Thank you for your insightful comment. Our current framework does not involve local training or global model aggregation, which differentiates it from FL setting. As such, widely-used FL methods like FedAvg and FedProx are not directly applicable to our approach, which is why they were not included in the experiments.

5. Typo on Llama-2-7B & line 502 Gemma-2B, and missing reference for ICL annotation under limited budget - "Mavromatis et al. Which examples to annotate for in-context learning? towards effective and efficient selection".

Thank you for pointing out these, we have corrected these in revised version (line 503 & related works).

Dear reviewer, thank you for your thoughtful feedback and for providing a positive evaluation of our work. We greatly value your insights, which have been instrumental in improving our submission. As the discussion period is coming to a close, we wanted to kindly check if our rebuttal addressed your concerns satisfactorily. If there are any remaining questions or points you would like us to clarify, we would be happy to address them promptly.

Thank you again for your time and effort in reviewing our work. We truly appreciate your contribution to this process.

1. novelty and technical depth is limited. core idea is server will gather optimal budget statistics using an existing proxy dataset on server side. This however is pretty straightforward.

To the best of our knowledge, there is no existing work that specifically investigates whether this solution is feasible, making our contribution novel and significant. Furthermore, distributed non-IID setting we propose holds meaningful implications for research community, particularly due to its applicability in real-world scenarios, such as collaborations between medical institutions.

We believe problem setting itself is critical, as it reflects practical challenges and opportunities that can inspire future research. Many influential works in field, such as those on in-context learning (ICL) and chain-of-thought (CoT) reasoning, have demonstrated that even straightforward methodologies can lead to profound contributions when they address meaningful and impactful problems.

Similarly, we emphasize importance of introducing a meaningful and relevant problem setting that aligns with real-world needs, which, in our view, is as valuable to community as sophistication of proposed methodology.

2. I think work is highly related to distributed RAG work. authors are suggested to include discussion of difference of existing distributed RAG works and compare with these approaches if possible.

Thank you for your valuable suggestion regarding inclusion of distributed RAG-related works. We sincerely appreciate your insightful feedback and agree that discussing differences between our approach and existing distributed RAG studies will help provide additional clarity and context to our contribution.

In developing this work, we carefully considered related studies in distributed RAG. However, challenges addressed by existing distributed RAG works differ from those tackled in our paper. For instance, [1] focuses on creation of datasets for distributed RAG frameworks and explores LLM-based labeling techniques for engineering pipelines. Their research scope and methodology are distinct from ours and are not directly applicable to our specific problem setting. Similarly, [2] addresses resource consumption and real-time response challenges in distributed RAG, emphasizing local retrieval efficiency and answer accuracy. However, it does not account for non-IID property in distributed settings. Additionally, [2] permits LLM deployment on partial local institutions, which is fundamentally different from our setting.

Real-world distributed non-IID RAG scenarios present a more complex framework, involving numerous challenges that must be addressed for effective deployment. For example:

• How can we effectively decompose a user query into subqueries while considering local knowledge distribution? What is best way to assign these subqueries to different clients with varying local expertise?

• How should we merge knowledge retrieved from multiple clients with overlapping expertise, and should we assign confidence levels to different clients for the same subqueries?

• How can local retrieval process be accelerated when dealing with large local databases?

These challenges represent broader avenues for exploration in distributed non-IID RAG. While our current work cannot directly compare with existing distributed RAG studies due to different settings, we believe it offers an interesting starting point for addressing such challenges. Specifically, our approach focuses on how to enable cooperation among clients with different knowledge distributions. By assigning preferences to clients based on their local knowledge distributions and employing an MLP to learn these distributions without transmitting complete local knowledge to a central server, we offer an intuitive method that could inspire future advancements in distributed non-IID RAG.

We have included this discussion in the revised manuscript to further emphasize these distinctions and highlight unique aspects of our approach. Once again, thank you for your thoughtful feedback, which has been very helpful in refining our work.

[1] Wang, Shuai, et al. "Feb4rag: Evaluating federated search in context of retrieval augmented generation." Proceedings of 47th International ACM SIGIR. 2024.

[2] Li, Jiaxing, et al. "EACO-RAG: Edge-Assisted and Collaborative RAG with Adaptive Knowledge Update." arXiv preprint arXiv:2410.20299 (2024).

1. Limited nonIID, can happen on text length, styles, etc. Authors only focus on class NonIID. Recommend for adding more nonIID setting.

Thank you for valuable suggestion on other nonIID. We appreciate it and agree that non-IID can manifest in various forms. In this paper, our primary goal is to demonstrate the feasibility of using ICL under nonIID conditions, focusing on class-level non-IID as a starting point. Given the scope of a single paper, it is challenging to cover all possible nonIID scenarios comprehensively.

To address your suggestion and evaluate applicability of our method to other nonIID, we added additional experiment involving feature skew. We consider a scenario where class distributions are balanced across clients, but query embedding distributions differ (i.e., style shifting without class nonIID).

In this setting: one client contains only Yelp training samples, while another client contains only Amazon training samples. Yelp & Amazon share same label space (5-class classification), while they show different distribution on queries. On server, we use both Amazon & Yelp samples as test set, and perform t-SNE on test sample's embedding with budget values on each client. By this setting, we want to show our method intuition on budget value and sample embeddings still holds for feature skew (style shifting).

As shown in [client 1 (with only Yelp)] and [client 2 (with only Amazon)], it is clear that clustering pattern is even more significant than previous class nonIID, indicating our claim still holds even without special class-based distribution. And it is very interesting to know that Amazon test samples will tend to assign all budget value to Amazon client, while Yelp test samples will tend to assign all budget value to Yelp client.

To conclude, we believe these results demonstrate that our method can be successfully applied to feature-skew (style-shifting) non-IID, further broadening its applicability. Thank you for your insightful suggestion, which has helped us expand scope of evaluation.

2. Since training of allocator does not require label of examples, experiments should not be limited to classification tasks. Effectiveness on generation task need to validated.

Thank you for the thoughtful suggestion on inclusion of generation tasks. We appreciate it and agree that exploring such tasks could be an interesting direction for future work. However, primary focus of this paper is to demonstrate the feasibility of using ICL under non-IID conditions, with text classification as a representative task.

Many prior ICL-related works [1][2][3] have similarly concentrated on classification tasks, providing significant value to research community. Our work builds on this foundation by examining nonIID for classification in detail, which we believe is a substantial and valuable contribution.

Furthermore, generation tasks in distributed nonIID pose unique challenges and require significant exploration of task-specific nonIID designs & experimental setups. Addressing these would broaden scope of our current work beyond its intended focus. For this reason, we believe it is more appropriate to limit scope of this paper to classification tasks.

We hope this clarifies our rationale, and we appreciate your valuable feedback, which has provided us with useful ideas for extending this work in future research.

[1] Lyu, Xinxi, et al. "Z-ICL: Zero-Shot In-Context Learning with Pseudo-Demonstrations." ACL. 2023.

[2] Yoo, Kang Min, et al. "Ground-Truth Labels Matter: A Deeper Look into Input-Label Demonstrations." EMNLP. 2022.

[3] Chen, Huiyao, et al. "Retrieval-style in-context learning for few-shot hierarchical text classification." TACL.2024.

3. partition for nonIIDness in main experiments is unreasonable. According to Table 7, for binary classification like Subj & MR, each client only having one class data is too extreme and can easily lead to biased predictions.

Thank you for suggestion on experiment design for binary classification. We added experiment on NonIID based on Dirichlet, where each client has samples from both classes with class imbalance. For detailed distribution on each client, please check [MR] & [Subj]. As shown in table, our method still outperforms others under non-extreme nonIID for MR&Subj.

1. Concrete examples of distributed non-IID ICL scenarios. Why can't these samples be used to simulate a comprehensive, unbiased retrieval pool for inference?

• Sorry for the confusion. We have added detailed examples to Appendix G of revised version to address your concerns and assist future readers in understanding this research. Please check it (due to character limitation, we do not post it here).

• “Why not use requested budgeted samples from each client to construct a better retrieval pool for inference?” Yes, it is a meaningful question during our framework design. We provide following reasons to argue that this solution is impractical in real-word settings.

  • During budget allocator training stage, each client only needs to upload “relevant scores” of local top relevant samples, rather than raw samples consisting of input query $x$ and corresponding label $y$, that is, server does not have raw information of each training sample from client. Thus, this information collected for budget allocator training cannot be used as an ICL retrieval pool.

  • During inference stage (with trained budget allocator), can we accumulate collected local samples of previous test queries and use these as a later retrieval pool? Due to privacy concerns, server platform can use local samples to perform inference while not allow caching these samples.

To conclude, a “comprehensive retrieval pool based on budgeted samples” is not applicable in our distributed non-IID scenario due to data pricing and privacy concerns.

2. If simulation of retrieval for proxy is reasonable (e.g., with large 𝑘), why not directly perform retrieval from this simulated data instead of training an allocator? If simulation is unreasonable, can this dataset still be valid for training allocator?

During construction of training data for allocator, we do not request local clients to send raw local training samples (that means, no label information) to server. Local clients only send “relevant scores” of local samples to server, considering cost of both communication and data pricing. Then we use collected ‘relevant scores’ to estimate “oracle budget” for each query in proxy dataset. That means, without these estimated “oracle budget” for queries in proxy set, it is impossible to train budget allocator.

• If simulation is reasonable (with large k): It is impossible for server to do ICL based on this collected training dataset, since retrieval process in budget allocator training stage only contains similarity between samples’ queries while no sample labels collected during this stage.

• If simulation is unreasonable: If it is allowed for transformation of relevance scores on only small k local samples (rather than a large k), then it can happen that estimated “oracle budget” is not accurate enough, which can lead to a sub-optimal budget allocator. Consequently, budget prediction during inference stage would result in sub-optimal budget allocation and lead to suboptimal ICL performance. An extreme case is, given 4 clients and server-side ICE budget for each query is 16, each client is only allowed to send top-1 relevant score to server. Then server can only estimate each local “oracle budget” as [1,1,1,1] (since overall number is even less than 16), which provides no useful information for training on budget allocator.

3. Detailed explanation on Figure 8, why show robustness on different size of proxy set?

• We add experiment results on extremely small proxy set, that is proxy set with only 100 and 50 samples. As shown in [proxy robustness results], with extremely small proxy set, performance does drop a lot (from 80+% to lower than 75%, even lower than 65%). results we present in our paper only shows results of proxy size from 300 to 700. Sorry for confusion.
• Also, as shown in [proxy robustness results], we separately draw performance curve on different quantization resolution $\delta$. As shown in figure, given a fixed quantization resolution value (for $\delta=3$, $\delta=4$), performance slightly increases when proxy size changes from 300 to 700. figure 8 presented in paper is “best result” for each proxy size (red curve). Thus it shows limited change. This may cause confusion, as mentioned in your comment.

We also present complete result figure in revised version. Thank you for pointing out this.

4. Why no inherently distributed non-IID datasets? Is it available? Thank you for suggestion. As far as we know, this is first paper on ICL tasks under distributed non-IID settings. Therefore, we didn't find any available open-source inherently distributed non-IID datasets.

5. Conclusion in section 3-observation on query embedding and budget assignment is too strong. This may be impacted by special non-IID in experiment. Intuition may not work under class-balanced setting.

Thank you for pointing out this. We claim that main reason our method works is not because of special class-based Non-IID setting. To show this, we add extra NonIID setting with feature skew but class-balance. We designed NonIID setting that 1 client contains only Yelp training samples, while another client contains only Amazon training samples. Yelp & Amazon share same label space and similar task, while they show different distribution on queries. On server, we use both Amazon&Yelp samples for test, and perform t-SNE on test embedding and their budget values on each client. By this setting, we want to show our claim on budget value and sample embeddings still holds even under feature skew with class-balance. As shown in [client 1 (with only Yelp)] and [client 2 (with only Amazon)], clustering pattern is even more significant than previous class-based non-IID, indicating our claim still holds even without special class-based distribution. Further, we emphasize that our method design has no relation to special label distribution.

6. Proposed allocator behaves like classifier for assigning budgets based on query classes, a straightforward rule-based budget allocation using known client classes might perform comparably.

Thank you for raising this important point about the behavior of our proposed allocator. While it may seem that our method functions like a straightforward rule-based allocator, we would clarify that our method is not simply trying to learn clients' local class distribution. Instead, our method tries to learn clients' local query embedding distribution. We use experiment where 1 client with only Yelp samples and 1 client with only Amazon samples to explain. In this setting, two clients share same local class distribution, which is balanced over 5 classes. If our method only relies on local class distribution, then each test query should assign equal budgets to two clients. However, as shown in client 1 (with only Yelp) and client 2 (with only Amazon), most queries assign all budget to client with most similar query distribution (from same dataset), while assigning 0 budget to client with different query distribution (from another dataset). This indicates that our method does not simply learn local class distribution, but learns query embedding distribution, which can be a more complex problem.

Dear reviewer aykb, thank you for taking the time to provide detailed suggestions on our submission. We have carefully responsed each of your comments in our rebuttal, and we truly appreciate the opportunity to clarify and expand upon our work based on your valuable insights.

As the discussion period is nearing its end, we wanted to kindly follow up to see if our responses addressed your concerns satisfactorily. If there are any remaining points or additional questions, we would be happy to provide further clarification.

Thank you again for your time and effort in reviewing our work. Your feedback is invaluable, and we greatly appreciate your engagement with our submission.