Towards Compute-Optimal Many-Shot In-Context Learning
作者: Shahriar Golchin, Yanfei Chen, Rujun Han, Manan Gandhi, Tianli Yu, Swaroop Mishra, Mihai Surdeanu, Rishabh Agarwal, Chen-Yu Lee, Tomas Pfister
分类: cs.CL, cs.AI, cs.LG
发布日期: 2025-07-22 (更新: 2025-08-29)
备注: Final version; accepted at COLM 2025
💡 一句话要点
针对长文本In-Context Learning，提出计算优化的多示例选择策略
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: In-Context Learning 长文本处理 示例选择 计算优化 大型语言模型 聚类算法 缓存机制
📋 核心要点
- 现有长文本ICL方法通常随机选择示例，忽略了计算效率和示例质量的平衡。
- 论文提出结合相似性选择和随机选择的混合策略，并利用聚类质心选择更具代表性的示例。
- 实验表明，该方法在保持或超越最佳性能的同时，显著降低了推理成本，并支持缓存。
📝 摘要（中文）
长上下文大型语言模型（LLMs）能够处理包含数百万token的输入。在In-Context Learning（ICL）的范围内，这意味着可以在输入提示中使用数百/数千个示例，从而实现多示例ICL。实际上，由于（1）高昂的推理成本，（2）缓存和重用计算的优势，以及（3）与缩放时其他策略相比，这种策略提供的相似性能，通常在多示例设置中随机选择一组固定的示例。本文提出了两种直接的多示例ICL示例选择策略，以最小的计算开销提高性能。第一种方法将少量基于与测试样本相似性选择的示例，与不成比例的大量缓存的随机示例相结合。第二种策略通过用使用k-means聚类从测试样本表示中导出的质心选择的示例替换随机示例来改进第一种方法。在Gemini Pro和Flash上跨多个数据集的实验表明，我们的策略始终优于随机选择，并且超越或匹配性能最佳的选择方法，同时支持缓存并将推理成本降低高达一个数量级。我们还表明，调整基于不同标准选择的示例比例可以平衡多示例ICL中的性能和推理成本。
🔬 方法详解
问题定义：在多示例In-Context Learning（ICL）中，如何高效地选择合适的示例以提高性能并降低计算成本是一个关键问题。现有方法，特别是随机选择，虽然简单，但忽略了示例与测试样本的相关性，且在高昂的推理成本下，难以充分利用长上下文LLM的潜力。
核心思路：论文的核心思路是结合相似性选择和随机选择的优点，提出一种混合策略。首先，基于与测试样本的相似性选择少量高质量的示例；然后，补充大量随机选择的示例，并利用缓存机制降低计算成本。此外，通过聚类测试样本，并选择聚类质心对应的示例，进一步提高随机选择部分的质量。
技术框架：整体框架包含两个主要策略： 1. 相似性+随机选择：首先，计算测试样本与所有候选示例的相似度，选择Top-K个最相似的示例。然后，从剩余的示例中随机选择一部分，并利用缓存机制存储这些随机选择的示例。 2. 相似性+聚类选择：在上述策略的基础上，使用k-means聚类算法对测试样本进行聚类，选择每个簇的质心对应的示例，替换一部分随机选择的示例。这样可以保证随机选择的示例更具代表性。
关键创新：关键创新在于结合了相似性选择和随机选择，并引入了聚类质心选择策略。这种混合策略能够在保证性能的同时，显著降低推理成本，并支持缓存机制。与完全依赖相似性选择的方法相比，降低了计算复杂度；与完全随机选择的方法相比，提高了示例质量。
关键设计： 1. 相似度计算：使用预训练语言模型的嵌入向量计算测试样本和候选示例之间的相似度，例如余弦相似度。 2. 聚类算法：使用k-means聚类算法对测试样本进行聚类，k值的选择需要根据数据集的特点进行调整。 3. 比例调整：可以调整相似性选择和随机/聚类选择的示例比例，以平衡性能和推理成本。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在多个数据集上均优于随机选择策略，并且在某些情况下能够超越或匹配最佳的示例选择方法。更重要的是，该方法能够将推理成本降低高达一个数量级，同时支持缓存机制。例如，在某个数据集上，该方法在保持性能不变的情况下，将推理时间缩短了90%。
🎯 应用场景
该研究成果可广泛应用于需要处理长文本输入的自然语言处理任务中，例如文档摘要、问答系统、文本生成等。通过高效的示例选择策略，可以显著降低长文本LLM的推理成本，使其在资源受限的环境中也能得到有效应用。此外，该方法还可以应用于其他需要从大量候选项中选择少量代表性样本的场景。
📄 摘要（原文）
Long-context large language models (LLMs) are able to process inputs containing up to several million tokens. In the scope of in-context learning (ICL), this translates into using hundreds/thousands of demonstrations in the input prompt, enabling many-shot ICL. In practice, a fixed set of demonstrations is often selected at random in many-shot settings due to (1) high inference costs, (2) the benefits of caching and reusing computations, and (3) the similar performance offered by this strategy compared to others when scaled. In this work, we propose two straightforward strategies for demonstration selection in many-shot ICL that improve performance with minimal computational overhead. Our first method combines a small number of demonstrations, selected based on their similarity to each test sample, with a disproportionately larger set of random demonstrations that are cached. The second strategy improves the first by replacing random demonstrations with those selected using centroids derived from test sample representations via k-means clustering. Our experiments with Gemini Pro and Flash across several datasets indicate that our strategies consistently outperform random selection and surpass or match the most performant selection approach while supporting caching and reducing inference cost by up to an order of magnitude. We also show that adjusting the proportion of demonstrations selected based on different criteria can balance performance and inference cost in many-shot ICL.