Distributed In-Context Learning under Non-IID Among Clients

Siqi Liang, Sumyeong Ahn, Jiayu Zhou

Keywords: in-context learning, distributed system, large language model

Abstract:

Advancements in large language models (LLMs) have shown their effectiveness in multiple compli-
cated 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.

Primary Area: foundation or frontier models, including LLMs

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 3853