Distributed In-Context Learning under Non-IID Among Clients

## Abstract

Abstract Advancements in large language models (LLMs) have shown their effectiveness in multiple complicated natural language reasoning tasks.
A key challenge remains in adapting these models efficiently to new or unfamiliar tasks.
In-context learning (ICL) provides a promising solution for few-shot adaptation by retrieving a set of data points relevant to a query, called in-context examples (ICE), from a
training dataset and providing them during the inference as context.
Most existing studies utilize a centralized
training dataset, yet many real-world datasets may be distributed among multiple clients, and remote data retrieval can be associated with costs.
Especially when the client data are non-identical independent distributions (non-IID), retrieving from clients a proper set of ICEs needed for a test query presents critical challenges.
In this paper, we first show that in this challenging setting, test queries will have different preferences among clients because of non-IIDness, and equal contribution often leads to suboptimal performance.
We then introduce a novel approach to tackle the distributed non-IID ICL problem when a data usage budget is present.
The principle is that each client’s proper contribution (budget) should be designed according to the preference of each query for that client.
Our approach uses a data-driven manner to allocate a budget for each client, tailored to each test query.
Through extensive empirical studies on diverse datasets, our framework demonstrates superior performance relative to competing baselines.

## 1 Introduction

Recent significant progress in large language models (LLMs) [, , , ] has demonstrated their effectiveness across various natural language processing (NLP) tasks [, ].
Despite their impressive performances, they still require adaptation to the specific downstream tasks for better performance.
However, adaptation poses challenges due to LLMs’ vast number of trainable parameters.

Figure: Figure 1: Problem overview. When datasets are distributed among clients in a non-IID manner, it creates an obstacle in generating a good context (left). However, by assigning appropriate budgets to leverage per-client expertise, better context can be created (right).
Refer to caption: x1.png

In-context learning (ICL) [] is a notable method that distinguishes itself through both its effectiveness and efficiency.
In brief, ICL adapts to the target task by incorporating context information following two primary steps: i) identify samples from the training dataset helpful to solve the target query by creating a prompt describing a context; ii) feed the constructed prompt with the target query and get the answer.
Previous related works on ICL mainly have focused on the construction of a prompt describing the context, which involves several sub-problems, such as the retrieval of in-context examples (ICEs) [] and determining the optimal sequence for the selected ICEs [].

A common assumption in most existing ICL research is that the system has access to a high-quality centralized dataset used for retrieval.
However, in many application scenarios, such as health informatics, centralized datasets may not be feasible, and data could be distributed in different institutions, which calls for the distributed ICL.
In addition, when the data is proprietary and possesses high value towards inferences, access to data entries may also be bound to data pricing strategies [, ].
For instance, the system needs to pay the local institution based on the number of samples sent to the system [] as a means to share profits from inferences.
Under this scenario, aggregating ICEs from local clients to a center server for ICL entails significant financial costs and lacks efficiency.

In this paper, we focus on integrating knowledge from distributed clients to achieve better ICL performance under the per-query ICE budget constraint. Specifically, we formalize the distributed ICL problem where the ICEs are distributed on clients, and the server has an LLM for ICL inference but can only request a limited number of ICEs from all clients for each query, which we refer to as the ICE budget.

We begin by identifying the key challenge in distributed ICL with ICE budget constraints lies in the non-independently and identically distributed (non-IID) training data, as shown in Section 3.1.
For example, in Figure 1, data samples are spread across $C$ clients, each with a unique data distribution.
Specifically, client $1$ primarily contains ($+$) samples, while client $2$ is mainly constituted by ($-$) examples.
Only limited research [] tried to address the challenge of distributed datasets for ICL, while none considers the challenging real-world setting of non-IID clients.
This leaves a critical question unanswered: What happens to distributed ICL when local clients are non-IID?

To further the understanding of the key challenge in the distributed non-IID ICL, we explore the local retrieval process on non-IID clients.
We found that each query has different preferences for different clients based on local knowledge distribution, that is, the number of samples needed from different clients should vary based on local sample distribution.
As the toy example shown in Figure 1, when the server creates context by uniformly assigning budgets to clients, the answer might be incorrect due to the insufficiency of ($+$) information in the context.
To be more detailed, the server assigns the clients who have expertise on ($-$), ($\times$), and ($\div$) operations with the same budget as on ($+$), without any preference.
Nevertheless, if the server assigns more budget to clients with many ($+$) samples, such as client $1$, it can create a more relevant context to answer the query related to ($+$) operation.
This indicates that under non-IID, the server should allocate the budgets over clients based on the preference of each query itself, as well as the distribution of local training samples.

Motivated by this, we propose a novel distributed ICL framework to collaboratively collect scattered information among non-IID clients by properly assigning ICE budgets to each client.
First, the server will gather the optimal budget statistics using an existing proxy dataset on the server side.
Next, the server will use this dataset to train the budget allocator.
During the deployment stage, the server will predict the proper budget for each client using this trained budget allocator given each test query and perform ICL among clients.
Furthermore, in practical scenarios where privacy concerns arise, we augment our framework with the paraphrasing method [] to secure privacy.

Contributions.
A summary of our contributions:

- •
To the best of our knowledge, we are the first to study the challenging real-world setting of ICL with distributed non-IID clients.
We identify the principal challenge as properly assigning the ICE budget for non-IID clients based on the preference of each test query and local knowledge distribution.
- •
We propose a framework to handle the distributed non-IID ICL. This framework trains a budget allocator on the server with the help of a server-side proxy dataset. Then, the server will use this trained allocator to decide how many ICEs to retrieve from each client for the ICL process, enabling collaborative action among clients.
- •
Across a range of dataset benchmarks featuring various non-IID configurations as well as on different LLM architectures, our approach has been validated to enhance ICL performance.
Notably, we examine both non-private, *i.e., *communicate raw samples directly, and private cases using the paraphrasing method to secure privacy. In both scenarios, our approach shows superiority to the previous method and other reasonable baselines.

## 2 Problem Formulation

In this section, we provide a detailed problem formulation. First, we begin with the specifics of in-context learning (ICL), followed by a description of distributed non-IID ICL.

Figure: Figure 2: Overview of the pipeline: First, the budget allocator assigns a budget to each client based on the question. Subsequently, each client retrieves their relevant samples and sends them back to the server. The server infers the answer by feeding the question, which is composed of concatenated context examples and the query.
Refer to caption: x2.png

### 2.1 In-Context Learning

Notation.
We consider a NLP tasks which have training dataset $\mathcal{D}=\{(x_{i},y_{i})\}_{i=1}^{N}$ with $N$ training samples.
Here, $x_{i}$ is the input text, and $y_{i}$ is the corresponding output. In the test phase, a test query $x_{q}$ is given.

Retrieval.
We employ the off-the-shelf pre-trained retriever KATE [](^1^11We do not fine-tune the retriever for each task, which is impractical because we cannot gather the distributed datasets.), which utilizes $k$-NN example selection. This retriever employs a sentence encoder $\mathcal{E}(\cdot)$ to measure the similarity between the in-context example $x_{i}$ in dataset $\mathcal{D}$ and the query $x_{q}$ as follows:

$$ $\displaystyle d(e_{i},e_{q})=\lVert e_{q}-e_{i}\rVert_{2},$ (1) $$

where $e_{q}=\mathcal{E}(x_{q})$ and $e_{i}=\mathcal{E}(x_{i})$.
We select $k$ samples using the following criterion:

$$ $\mathcal{T}(e_{q},k|\mathcal{D})=\underaccent{e_{i}=\mathcal{E}(x_{i})\forall{ (x_{i},y_{i})\in\mathcal{D}}}{\arg\text{Top-}k}(d(e_{i},e_{q})),$ (2) $$

where $\mathcal{T}(e_{q},k|\mathcal{D})$ denotes the selected samples from the dataset $\mathcal{D}$, and used for inference.

ICL Inference.
In the test phase, given a test query with input $x_{i}$, relevant $k$ training samples called in-context examples (ICEs) are selected, *i.e., *$S=\mathcal{T}(e_{q},k|\mathcal{D})$.
Based on the retrieved samples, we feed the constructed context prompt $s(S,x_{q})$ into LLM for inference and obtain results via:

$$ $\displaystyle\begin{split}y_{t}&=\operatorname*{arg\,max}_{y}p_{{}_{\text{LLM} }}(y|s(S,x_{q}),y_{<t})\\ &s(S,x_{q})=(x_{1},y_{1})\odot\ldots\odot(x_{k},y_{k})\odot x_{q},\end{split}$ (3) $$

where the $\odot$ operation denotes concatenation, and $s(S,x_{q})$ is the context constructed using query $x_{q}$ and samples in $S$;
the term $p_{{}_{\text{LLM}}}$ represents the output softmax probability of the LLM, functioning autoregressive, meaning that the output up to time $t$, *i.e., *$y_{<t}$, is input back into the model to generate the $t^{\text{th}}$ output, $y_{t}$.
Previous works [, ] on ICL mainly focus on the selection of $S$ under a centralized setting.
However, we investigate the scenario where $\mathcal{D}$ is split among several clients, each following non-IID distributions.

### 2.2 Distributed non-IID ICL

Distributed ICL Setting.
We consider $C$ clients with a centralized server in our system. Each client $c\in[C]$ has local training dataset $\mathcal{D}_{c}=\{(x_{i}^{c},y_{i}^{c})\}_{i=1}^{N_{c}}$ with $N_{c}$ training samples.
Note that $\mathcal{D}_{c}$ follows different distributions for different clients. We follow the non-IID conditions as defined in [], with details provided in Appendix A.
In summary, we allocate data on a per-class basis, where each client receives a specific number of classes, meaning each client has samples from only specified classes.
Clients and the server have identical off-the-shelf pre-trained retrievers. Consider the computation resource limitation on clients as in many real scenarios [], only the server is equipped with an LLM. Moreover, the server has limited proxy dataset $\mathcal{D}_{\text{proxy}}=\{(x_{j}^{\text{proxy}},y_{j}^{\text{proxy}})\}_{j=
1}^{N_{\text{proxy}}}$, that $N_{\text{proxy}}\ll\sum_{c=1}^{C}{N_{c}}$.
The server has quite a small $\mathcal{D}_{\text{proxy}}$, and it is an auxiliary dataset to extract information for collaboration to make the problem feasible.

Pipeline.
First, the server requests relevant samples from each client by sending $x_{q}$ to all clients with local budgets $k_{c}$.
Remark that each query $x_{q}$ has its own preference of each client, which can be represented as $k_{c}$.
Here, $x_{q}$ can be anonymized by paraphrasing, as done in previous works [](^2^22Although our main experiments utilize the non-paraphrased dataset, we also present the paraphrased results in Section 5.).
Each client then selects the most relevant $k_{c}$ samples from their local training dataset, *i.e., *$S_{c}=\mathcal{T}(e_{q},k_{c}|\mathcal{D}_{c})\subset\mathcal{D}_{c}$, and returns them to the server.
The server receives $S_{c}$ from clients and generates the context $s$ based on the merged examples, $S=\bigcup_{c=1}^{C}S_{c}$.
In the final step, the server infers $y$ using $s(S,x_{q})$. The entire framework also can be described in Figure 2.
In this paper, we are concentrating on assigning $k_{c}$ to each client as described in Figure 2.

Figure: Figure 3: Non-IID experimental results. It shows that centralized performance is comparable to the IID case, whereas non-IID scenarios exhibit a significant declined performance. This highlights the critical importance of addressing non-IIDness to find a solution.
Refer to caption: x3.png

## 3 Observations

In this section, we describe several empirical supports to handle the distributed non-IID ICL. First, we demonstrate that non-IID distributions hinder the merging of scattered information.
We then establish our goal, termed as *oracle budget*, which reflects the server’s preference for each client if the server knows all distributed data.
Finally, we check if predicting the oracle budget of each test query for inference is feasible.

### 3.1 Non-IIDness Leads to Performance Drop

First of all, we evaluate the effect of non-IIDness. Straightforwardly, we distribute the budget $\{k_{c}\}_{c=1}^{C}$ uniformly according to the following criteria: Given $C$ clients are involved in answering this question, and the number of samples for context is $k$. We first explore the naïve equally assigned local budget scheme in both IID and non-IID settings. That is, each client $c\in[C]$ locally retrieves top-$k_{c}$ samples where $k_{c}=\lceil\frac{k}{C}\rceil$ from local dataset $\mathcal{D}_{c}$. Detailed experimental settings are described in Appendix B.

As illustrated in Figure 3, we observe the followings: (1) There is no significant performance degradation between the centralized case ($\blacksquare$) and the IID case ($\blacksquare$).
This is expected, as the merged top-$k_{c}$ samples in the IID case closely resemble the centralized top-$k$ samples. Any minor discrepancies are attributed to differences in sample ordering. (2) However, performance degradation becomes pronounced in non-IIDness case (refer to the comparison between $\blacksquare$, $\blacksquare$ and $\blacksquare$). Hereinafter, we gather insights to address the distributed non-IIDness ICL.

Figure: (a) SST-5
Refer to caption: x4.png

### 3.2 Proper Budget per Query for Each Client

Oracle budget.
The remaining issue is that to make the server operate similar with the centralized manner, it needs to allocate the budget as if it knows complete knowledge of all clients. We call this budget for each client as the *oracle budget* for query embedding $e_{q}$ and define it as follows:

$$ $k_{c}^{\star}(e_{q})=\Big{|}\mathcal{T}(e_{q},k|\mathcal{D}_{c})\cap\mathcal{T }(e_{q},k|\mathcal{D})\Big{|},$ $$

where $\mathcal{T}(\cdot)$ is defined as Eq. (missing) 2 and $|\cdot|$ is set cardinality. Note that the physical meaning of $k_{c}^{\star}(e_{q})$ is the number of shared samples between the top-$k$ relevant to $e_{q}$ in local $\mathcal{D}_{c}$ and global $\mathcal{D}$ datasets.

Check of predictability of oracle budget.
For the next step, it is necessary to check if $e_{q}$ has sufficient patterns of oracle budget to extract and use it in the inference phase.
Our hypothesis is that similar queries may share similar oracle budget patterns and preferences on the same client, and it can lead to similar budget allocations for that client.
Therefore, to verify this hypothesis,
we perform t-SNE analysis [] on the embeddings obtained from the retriever for queries.
Furthermore, we color each sample based on the oracle budget $k_{c}^{\star}(e_{q})$.
As described in Figure 4, similar query embeddings exhibit similar oracle budget patterns.
This indicates that, given a test query, we can infer the budget assignment
for each client.
However, it is challenging to predict fine-grained budget value since there are no rigid classification patterns.
For instance, determining the detailed budget value seems challenging in the case of client $1$ in SST-5.
Therefore, developing an efficient method to infer the exact
budgets based on these broad patterns for each client are required.

### 3.3 Observation Summary

In summary, our findings and the approach for designing an algorithm are as follows: (1) non-IIDness significantly affects the distributed ICL setting, necessitating the development of a coalition method. To handle this problem, it is straightforward to allocate an appropriate number of budgets to each client, *i.e., *making server work so as it knows client all samples. (2) By analyzing the query embeddings, we can determine the importance of each client per query.

## 4 Method

Figure: Figure 5: Overview of the budget allocator: We train a budget allocator on top of the frozen feature extractor $\mathcal{E}$, which inherits from the retriever. During inference, when a test query $x_{q}$ is provided, this module determines the quantized budget levels for each client and allocates them accordingly.
Refer to caption: x8.png

In this section, we outline the proposed algorithm to mitigate non-IIDness in the ICL framework. Specifically, we show how to train the *budget allocator* and conduct inference.

Figure: Algorithm 1 Top-$k$ sampling, $\mathcal{T}(e,k|\mathcal{D})$

Figure: Algorithm 2 Construct dataset

Figure: Algorithm 3 Inference (Client, Server)

### 4.1 Train a Budget Allocator

Based on Section 3, it is feasible to assign budgets of each client by using the embeddings obtained from the retriever encoder $\mathcal{E}$. We first construct the datasets having the targeting budget values and then train the budget allocator.The pseudo-codes are described in Algorithm 1 and 2.

Construct dataset for oracle budget.
First, we explain how to create a dataset to train the budget allocator for each client, as described in Algorithm 2.
Given proxy dataset $\mathcal{D}_{\text{proxy}}$, for all embeddings $e_{j}=\mathcal{E}(x_{j})$ where $(x_{j},y_{j})\in\mathcal{D}_{\text{proxy}}$, we request $k$ samples from each client $c\in[C]$ using Top-$k$ procedure, *i.e., *$S_{c}=\mathcal{T}(e,k|\mathcal{D}_{c})$.
Once the server receives $k$ examples from each clients, *i.e., *$\{S_{c}\}_{c=1}^{C}$, it merges and re-orders them to obtains $S^{\text{top}}$.
Based on $S^{\text{top}}$, we count the number of samples from each client in $S^{\text{top}}$, *i.e., *compute $k_{c}(e_{j})$. After counting $k_{c}(e_{j})$ for all clients, we quantize the budget levels for each client using the quantization hyper-parameter $\delta$. As a result, the output of this procedure is $B_{\text{proxy}}$ for all clients, composed of embeddings $e$ and their respective budgets $k_{c}(e_{j})$.

Train budget allocator.
Based on the constructed dataset $B_{\text{proxy}}$, we train the *budget allcoators*, *i.e., *$\{f_{c}(\cdot)\}_{c=1}^{C}$, for each $f_{c}(\cdot)$ has Multi-layer perceptrons on top of the frozen feature extractor of the off-the-shelf retriever $\mathcal{E}$. The budget allcoators are trained on the cross-entropy loss, as we have already quantized the optimal budgets using the hyper-parameter $\delta$. Note that if $\delta$ is high, the quantization is severe, otherwise the quantization is mild.

### 4.2 Inference Using Budget Allocator

We derive the response to the test query $x_{q}$ utilizing the LLM $\mathcal{M}(\cdot)$ through the described steps (see Algorithm 3 for specifics).
We first extract the embedding $e_{q}=\mathcal{E}(x_{q})$. Then, we compute the allocated budget $\{\hat{k}_{c}=f_{c}(e_{q})\}_{c=1}^{C}$ and send $\hat{k}_{c}$ to each client. Each client sends back top $\hat{k}_{c}+\alpha$ samples, *i.e., *$S_{c}$, to the server. Note that we summarize how the budget allocator outputs $\hat{k}_{c}$ in Figure 5. Here, $\alpha$ denotes the buffering hyper-parameter, which increases the chances for each client to be involved. After collecting $S_{\text{agg}}=\bigcup_{c\in[C]}S_{c}$, we aggregate them and run the usual ICL procedure.

## 5 Experiment

### 5.1 Experiment setup

First, we summarize the baselines, datasets, and the method for constructing non-IID settings. Finally, we depict the implementation details.

Baselines.
We compare our algorithm with various baselines, including social learning [], which does not account for non-IIDness, and other possible ways for handling distributed non-IID ICL,
such as Zero-shot, Proxy-only, Singleton (single client), Uniform-budget, Random-budget, and $\infty$-budget (oracle case). The detailed explanations are described in Appendix C.

**Table 1: Main results: To address the issue of non-IIDness in distributed ICL, we examined seven datasets and seven straightforward baselines. We run three random seeds and illustrate mean and std values. The top performance is highlighted in bold font, excluding the infinite budget scenario due to its impracticality. In summary, the proposed method effectively mitigates the non-iid distributed ICL problem to a reasonable extent.**
| \hlineB4      Algorithm | Dataset |  | Avg |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| SST-5 | Amazon | Yelp | MR | Yahoo | AGNews | Subj |  |  |  |  |
| Zero-shot |  | 29.19 | 24.70 | 31.23 | 73.95 | 25.87 | 67.60 | 50.55 |  | 43.30 |
| Proxy-only |  | 40.64$\pm$ 2.89 | 28.43$\pm$ 0.11 | 31.85$\pm$ 1.28 | 70.40$\pm$ 1.54 | 54.73$\pm$ 0.93 | 84.65$\pm$ 0.42 | 71.09$\pm$ 1.34 |  | 54.54 |
| Singleton |  | 25.14$\pm$ 4.18 | 24.03$\pm$ 0.57 | 29.44$\pm$ 3.91 | 50.00$\pm$ 0.00 | 38.14$\pm$ 2.03 | 50.60$\pm$ 0.66 | 50.00$\pm$ 0.00 |  | 38.19 |
| Social Learning |  | 36.03$\pm$ 0.27 | 28.42$\pm$ 0.19 | 29.25$\pm$ 0.45 | 58.58$\pm$ 0.18 | 46.03$\pm$ 0.49 | 81.10$\pm$ 0.29 | 71.37$\pm$ 0.71 |  | 50.11 |
| Uniform-budget |  | 32.94 | 25.63 | 26.60 | 33.65 | 43.00 | 73.17 | 63.20 |  | 42.60 |
| Random-budget |  | 32.82$\pm$ 0.82 | 25.69$\pm$ 0.55 | 27.72$\pm$ 0.51 | 34.68$\pm$ 0.59 | 42.46$\pm$ 0.53 | 67.34$\pm$ 0.39 | 65.37$\pm$ 0.80 |  | 42.30 |
| $\infty$-budget |  | 43.26 | 32.70 | 34.80 | 77.20 | 62.67 | 89.37 | 91.4 |  | 61.62 |
| Ours |  | 44.08$\pm$ 0.12 | 31.54$\pm$ 0.22 | 35.48$\pm$ 0.28 | 80.44$\pm$ 0.67 | 61.67$\pm$ 0.25 | 88.52$\pm$ 0.30 | 82.36$\pm$ 0.91 |  | 60.58 |
| \hlineB4 |  |  |  |  |  |  |  |  |  |  |

### 5.2 Main results

We have presented the performance of our algorithm and baselines in Section 5.1.
First, we can observe that performance varies significantly depending on the way the budget is allocated, which indicates that the budget allocation scheme really matters in distributed non-IID ICL.
Additionally, even when using only the proxy dataset, there is a performance improvement, and this performance surpasses that of using other clients which have the tilted local datasets (*e.g., *$29.19\%\to 40.64\%$ in SST-5 case).
This indicates that utilizing a biased dataset can degrade the ICL performance. Although social learning algorithm has shown good performance in the previous paper, it does not perform well under the non-IID cases configured in this research.
If we can use an infinite budget, all settings would exhibit high performance.
However, our proposed algorithm demonstrates better performance than the infinite budget upper limit (*e.g., *$34.86\%\to 35.48\%$ in the Yelp case).
This is likely due to a mechanism that prevents unnecessary information from being selected by the retriever with high importance.
Ultimately, the proposed algorithm shows an average performance improvement of $5.05\%$ percentage points across seven datasets compared to the best performance of baselines using the proxy dataset.
This shows that the proposed algorithm can handle the non-IID case well.

### 5.3 Analysis

In this section, we further examine four key aspects: (1) privacy-preserving case analysis, which encompasses paraphrasing both training and testing queries, (2) sensitivity to hyper-parameters, (3) the performance of the trained budget allocator, and (4) the compatibility of the LLMs.

**Table 2: Analysis of the generated query and training samples: We re-create the datasets using small-sized LLMs and conduct the experiments as in Section 5.1 under the exactly same experimental settings.**
| \hlineB4      Algorithm | Dataset |  | Avg |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| SST-5 | Yelp | Subj |  |  |  |  |
| Zero-shot |  | 27.96 | 31.40 | 51.55 |  | 36.97 |
| Proxy-only |  | 39.39$\pm$ 1.33 | 31.78$\pm$ 1.75 | 73.46$\pm$ 1.46 |  | 48.21 |
| Singleton |  | 25.31$\pm$ 3.89 | 30.78$\pm$ 4.88 | 50.08$\pm$ 0.10 |  | 35.39 |
| Social Learning |  | 33.09$\pm$ 0.68 | 28.80$\pm$ 0.33 | 74.82$\pm$ 0.93 |  | 45.47 |
| Uniform-budget |  | 27.06 | 26.60 | 63.30 |  | 38.99 |
| Random-budget |  | 27.29$\pm$ 0.51 | 27.70$\pm$ 0.46 | 63.88$\pm$ 0.81 |  | 39.62 |
| $\infty$-budget |  | 41.63 | 37.23 | 90.75 |  | 56.54 |
| Ours |  | 40.37$\pm$ 0.27 | 36.52$\pm$ 0.89 | 83.82$\pm$ 1.00 |  | 53.57 |
| \hlineB4 |  |  |  |  |  |  |

## 6 Related Work

In-context learning.
ICL [] is one of the fastest paradigms using pre-trained LLMs by feeding several examples to construct the context to solve the given query.
The main criteria of this research field are to find the most informative samples among the training datasets.
For example [] trains BERT [] oriented encoder and use the $k$ nearest neighbors.
One of the reasonable sparse retriever, rule-based approach, is using BM25 [] which measures the term-frequency.
[] proposed an efficient retriever called EPR.
It trains two encoders by inheriting the method of dense passage retriever (DPR) [] under the loss of positive and negative pairs.
To reduce the domain specificity, [] proposed UDR, which is applicable to multiple domain tasks in a universal way and shows reasonable performance from a single retriever.
PromptPG [] utilized a reinforcement learning framework to train the retriever so that it can generate context to improve the answerability of LLMs.
Similarly, LLM-R [] uses a reward model to train the retriever.
Note that extensive research has not targeted to solve the distributed cases. These works have seen the centralized case.

Distributed ICL.
To the best of our knowledge, only a single study [] tries to address ICL in a distributed manner.
However, this paper solely focuses on merging the distributed information without considering the nature of the non-identically distributed information.
Many studies, such as those on federated learning [, , ], address the non-IID distribution of datasets, highlighting the need to handle distributed non-IID ICL.

## 7 Conclusion

In this paper, we tackle the challenge of ICL when datasets are distributed among clients with non-IID distributions.
Initially, we examine if non-IID distributions lead to performance degradation and discover that they cause significant drops in performance.
We propose an algorithm that learns the task of budget assignment and employs it during inference to allocate appropriate budgets for each query.
Using this proposed algorithm, we achieve performance improvements across various benchmarks.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Why not evaluate using inherently distributed non-IID datasets? Is there an available one or not?
2. How are proxy datasets constructed, and do they ensure coverage of all clients? The description suggests the proxy set is sampled directly from the test set, but what real-world scenario does this correspond to, and how would such a proxy dataset be realistically constructed? Can you provide specific examples of a proxy dataset in real-world application?
soundness: 2
strengths: 1. This paper highlights the important yet underexplored problem of distributed non-IID in-context learning (ICL), offering fresh insights into a new problem.
2. The paper presents some valuable experimental results demonstrating the poor performance of distributed non-IID ICL under simple uniform budget allocation, effectively validating the significance of the problem.
3. The method proposed in this paper have improved the performance in a simple while effective way.
summary: This paper tackles the challenge of distributed non-IID in-context learning (ICL) for LLMs, where data is spread across clients with differing distributions. The paper first shows that uniform in-context examples would fail on non-IID situations. Then the authors propose a method to optimize the allocation of a limited in-context examples (ICEs) budget by training a budget allocator. This allocator predicts query-specific budgets for each client, addressing the inefficiencies of uniform allocation under non-IID settings. The method outperforms baseline approaches like random and uniform budgets in experiments, improving ICL performance on distributed datasets.
weaknesses: 1. The paper does not provide concrete examples of distributed non-IID ICL scenarios. So i don't get that given the server can request budgeted samples from each client, why can't these samples be used to simulate a comprehensive, unbiased retrieval pool for inference?

2. The training dataset for the allocator is constructed by retrieving k samples per query from each client, combining these 𝐶×𝑘
samples to simulate a unified dataset. However, this simulation raises questions. If the simulation is reasonable (e.g., with large 𝑘), why not directly perform retrieval from this simulated dataset instead of training an allocator? If the simulation is unreasonable, can this dataset still be valid for training the allocator?

3. In Figure 8, the proxy size appears to have minimal influence, which is surprising. Since the allocator is trained using questions from the proxy dataset, a poor match between the proxy and test set questions should bias the collected data and hinder the training of a good allocator. However, this surprising result lacks a detailed explanation.

4. The conclusion in Section 3-Observations that query embeddings can determine budget assignments is based on observed clustering patterns in oracle budgets corresponding to different queries. This conclusion may be too strong, as other factors, such as the specific distribution of client data, might play a criticle role. For instance, in the experiments, non-IID clients are constructed based on classes, and these class-based distributions likely influence budget assignments significantly.

5. The experiments simulate non-IID clients based on data classes. Could an LLM directly infer which classes are relevant for a query and decide sample allocations accordingly? Since the allocator effectively behaves like a classifier for assigning budgets based on query classes, a straightforward rule-based budget allocation using known client classes might perform comparably.

6. There are multiple spelling mistakes in the paper. For instance, in Section 2.2, the first two sentences describing the pipeline use k_c with seemingly different meanings, leading to confusion.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: The training of the allocator relies on a proxy dataset. In practice, the distribution of test data is usually unknown and obtaining a proxy dataset with the same distribution is unrealistic. What if the distribution of the proxy dataset differs from that of the test set?
soundness: 2
strengths: 1. Distributed ICL under non-IID conditions is an interesting problem and aligns well with real-world scenarios. This paper explore the challenges under this setting, providing some meaningful insights.
2. The proposed method is simple yet effective across several benchmarks, with low training overhead.
summary: This paper addresses the challenge of distributed in-context learning under non-identically distributed (non-IID) data across multiple clients. It trains a budget allocator to dynamically allocates the number of in-context examples (ICE) retrieved from each client based on the relevance of query. Specifically, it trains a multi-layer perceptron (MLP) for each client to approximate the oracle budget, derived from a server-side proxy dataset, enabling efficient and targeted ICE retrieval per query. This paper additionally explores paraphrasing techniques to ensure data privacy in distributed contexts. Empirical results across multiple datasets show that this approach outperforms several baselines.
weaknesses: 1. In real-world scenarios, data distribution differences can manifest in multiple aspects, such as text length, style, etc., but this paper only focus on Non-IIDess at the class level. I strongly recommend the author to take more aspects into considerations.
2. Since the training of allocator does not require the label of examples, the experiments should not be limited to classfication tasks. The effectiveness of allocator on generation tasks remains to be validated. 
3. The partition for non-IIDness in the main experiments is unreasonable. According to Table 7, for binary classification task like Subj and MR, each client only having access to data under one class is too extreme and can easily lead to biased predictions.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: as shown in the weakness.
soundness: 2
strengths: Pros:
1). The paper is well structured 

2). The experiments are thorough.
summary: This paper aims at addressing the issue of non-iid distributed  in-context examples (ICE) for in context learning. The authors 
introduce an approach to tackle the distributed non-IID ICL problem by calculate the budget for different clients on the number of ICEs for different clients. The principle is that each client’s proper contribution (budget) should be designed according to the preference of each query for that client. This is done by the server who
will gather the optimal budget statistics using an existing proxy dataset on the server side. Basically, the idea is straightforward with limited novelty. The paper is well structured and the experiments are thorough. However, as mentioned, the novelty and technical depth is limited.
weaknesses: Cons:

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

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

3). The paper only uses small open-sourced LLMs, such as GPT-Neo-1.3B GPT-Neo-2.7B Llama-2-7B. Is that possible to provide results using larger ones, such 70B LLama-3.1 and other close-sourced ones, such as Claude 3.5 and GPT-4o?

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. In a real world use case, even when sharing a similar task, the clients might have totally different prompts in terms of structure, length, and specific requirements. Will this affect retrieval effectiveness since it heavily depends on the similarity between query and training examples? 
2.  Could the authors provide more details on the proxy dataset used to train the budget allocator? Specifically, how is the proxy dataset selected, and how does it ensure adequate representation of non-IID distributions across diverse client tasks?
3. The paper references federated learning in related works, have you done any comparisons with FL methods that tackle non-IID distribution, especially regarding data and budget efficiency
4.  How sensitive is the budget allocator to differences across task types? Does the allocator need to be re-trained for different tasks?
5. Table 3 shows the result from Llama-2-7B but line 502 says Gemma-2B was used
6. Missing reference for ICL annotation selection under limited budget - "Mavromatis et al. Which examples to annotate for in-context learning? towards effective and efficient selection"
soundness: 3
strengths: This paper identifies a real-world complexity which is the non-iidess of the dta settings during ICL. It introduced a budge allocator that dynamically selects the most relevant ICEs from each client based on their own queries, which leads to improvement over uniform or random ICE distributions.  It also includes a privacy-preserving option using paraphrasing and demonstrates consistent performance gains over such privacy-preserving setting.
summary: This paper proposed a solution to ICL when non-iid data exist across clients. The authors identify that traditional ICL assumes a centralized and uniform dataset that might not hold in real-world use cases (distributed), where each client's data might vary significantly. Thi could lead to suboptimal results by uniformly retrieving ICEs. In this paper the authors propose a framework that trains a budget allocator to determine the optimal number of ICEs to retrieve from each client based on their data's relevance to the query. The allocator uses a server-side proxy dataset to guide budget assignment per query - and this adjusts the contributions from each client and therefore more relevant context for inference. The experiments in this paper cover various LLM architectures over a series of classification tasks and demonstrate that this method improves ICL performance over existing distributed strategies (especially in non-iid scenarios)
weaknesses: 1. I feel the task covered in the study is quite limited to classification tasks. It would benefit from proving the effectivenss across other more complex tasks that are more context-heavy such as RAG and multihop reasoning. 
2. The 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 the claim
3. The proposed solution relies on the assumption that a high quality, representative proxy dataset presents at server-side. Although experiments have shown the method stability data size. it would be nice to see furtehr experiments on how the proxy data quality or distribution affects performance or alternative methods to reduce such dependency.

Meta-review: In this paper, the authors proposed a new setting for ICL, where training data is stored in a distributed manner. 

There are some major concerns raised by the reviewer. 1, The assumption of the proposed method may not be practical. Though the authors added some experiments, the assumption is still not convincing. 2, Though the problem setup looks new, the novelty of the proposed method is technically limited. 3, The study of the proposed new problem setup is only restricted to classification tasks. It is not a convincing reason to argue that previous ICL studies were only focused on classification problems in the authors' rebuttal.

By considering the concerns mentioned above, this paper does not meet the acceptance standard for ICLR.

Decision: Reject