Distributed In-Context Learning under Non-IID Among Clients

TL;DR

This paper addresses the challenge of in-context learning with distributed, non-IID client data by proposing a data-driven client contribution allocation method, improving performance in multi-client scenarios.

Contribution

It introduces a novel approach for allocating data usage budgets among clients based on query preferences in non-IID distributed settings for in-context learning.

Findings

Outperforms baseline methods in diverse datasets

Effective client contribution allocation improves ICL performance

Addresses non-IID data challenges in distributed environments

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.