The Role of Diversity in In-Context Learning for Large Language Models

2025-05-26

AI Summary

Existing in-context learning (ICL) example selection predominantly relies on similarity metrics, while the role of example diversity remains underexplored and lacks systematic theoretical grounding. Method: This paper establishes the first formal theoretical framework elucidating how example diversity enhances large language models’ (LLMs’) robustness on complex tasks—such as mathematical reasoning and code generation—as well as out-of-distribution (OOD) queries. We propose a diversity-aware retrieval-augmented ICL method and conduct comprehensive evaluations across Llama-3.1, Gemma-2, and Mistral-v0.3. Results: Integrating diversity with similarity yields substantial multi-task performance gains: average accuracy improves by 4.2% on math and programming benchmarks, and OOD generalization error decreases by 18%. Our core contribution is the rigorous theoretical characterization and empirical validation of diversity as a fundamental, orthogonal dimension—alongside similarity—in ICL, thereby advancing the principled design of effective in-context exemplars.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search
• Natural Language Processing: Safety and Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

In-context learning (ICL) is a crucial capability of current large language models (LLMs), where the selection of examples plays a key role in performance. While most existing approaches focus on selecting the most similar examples to the query, the impact of diversity in example selection remains underexplored. We systematically investigate the role of diversity in in-context example selection through experiments across a range of tasks, from sentiment classification to more challenging math and code problems. Experiments on Llama-3.1, Gemma-2, and Mistral-v0.3 families of models show that diversity-aware selection methods improve performance, particularly on complex tasks like math and code, and enhance robustness to out-of-distribution queries. To support these findings, we introduce a theoretical framework that explains the benefits of incorporating diversity in in-context example selection.

Problem

• Investigates diversity's role in in-context example selection for LLMs
• Explores performance impact on complex tasks like math and code
• Proposes theoretical framework for diversity benefits in ICL

Innovation

• Diversity-aware example selection for ICL
• Improves performance on complex tasks
• Enhances robustness to out-of-distribution queries