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.

What problem does this paper attempt to address?

### Problems the Paper Aims to Solve This paper aims to address the issue of Non-Independent and Identically Distributed (Non-IID) data encountered during In-Context Learning (ICL) in a distributed environment. Specifically, it focuses on how to effectively retrieve suitable In-Context Examples (ICE) from multiple clients with inconsistent data distributions to enhance the adaptability of Large Language Models (LLMs) to new tasks. ### Detailed Explanation 1. **Background and Challenges**: - Large Language Models (LLMs) perform excellently in various natural language processing tasks, but they need to adapt to new or unfamiliar tasks. - In-Context Learning (ICL) achieves few-shot adaptation by providing relevant context examples during inference. - Existing research mostly assumes a centralized high-quality dataset for retrieval, but in practical applications, data may be distributed across different institutions, and data access may involve costs. - When client data is Non-Independent and Identically Distributed (Non-IID), retrieving suitable ICE from each client becomes highly challenging. 2. **Main Contributions**: - This paper is the first to study the real-world challenges of ICL in a distributed Non-IID client environment. - It proposes a framework to optimize the performance of distributed ICL by reasonably allocating the budget (i.e., the number of ICE retrieved from each client) for each client. - The framework uses a data-driven approach to allocate the budget based on each query's preference for different clients. 3. **Method Overview**: - **Problem Definition**: Formalizes the distributed ICL problem where ICE is distributed across various clients, and the server has an LLM for inference but can only request a limited number of ICE from all clients per query. - **Key Challenge**: Under Non-IID data, each query has different preferences for different clients, so the budget needs to be allocated based on each query's preference and local data distribution. - **Solution**: Proposes a budget allocator that predicts the budget requirement for each query for each client through training and uses these budgets for ICE retrieval during inference. 4. **Experimental Results**: - Extensive experiments were conducted on multiple datasets to validate the effectiveness of the framework under different Non-IID configurations. - Experimental results show that this method outperforms existing methods and other reasonable baseline methods in both non-private and private scenarios. ### Conclusion This paper proposes an effective method to solve the ICL problem in a distributed Non-IID environment by reasonably allocating the budget, significantly improving the model's adaptability to new tasks. This method is not only innovative in theory but also has important practical significance.

Distributed In-Context Learning under Non-IID Among Clients

Implicit In-context Learning

"In-Context Learning" or: How I learned to stop worrying and love "Applied Information Retrieval"

In-Context Learning with Iterative Demonstration Selection

Does In-Context Learning Really Learn? Rethinking How Large Language Models Respond and Solve Tasks via In-Context Learning

Batch-ICL: Effective, Efficient, and Order-Agnostic In-Context Learning

Locally Differentially Private In-Context Learning

ParaICL: Towards Robust Parallel In-Context Learning

Adapting Large Multimodal Models to Distribution Shifts: The Role of In-Context Learning

In-Context Learning Demonstration Selection via Influence Analysis

Many-Shot In-Context Learning

In-Context Learning on a Budget: A Case Study in Named Entity Recognition

A Survey on In-context Learning

Communication-efficient federated continual learning for distributed learning system with Non-IID data

One size doesn't fit all: Predicting the Number of Examples for In-Context Learning

In-Context Learning for Text Classification with Many Labels

Feature-Adaptive and Data-Scalable In-Context Learning

Contextual Client Selection for Efficient Federated Learning over Edge Devices

Which Examples to Annotate for In-Context Learning? Towards Effective and Efficient Selection

Dual Operating Modes of In-Context Learning

Let's Learn Step by Step: Enhancing In-Context Learning Ability with Curriculum Learning