Dynamic Uncertainty Ranking: Enhancing In-Context Learning for Long-Tail Knowledge in LLMs

2024-10-31

arXiv.org

AI Summary

To address prediction instability in large language models (LLMs) for long-tail domain question answering—caused by uncertainty in retrieved samples—this paper proposes a reinforcement learning–based dynamic uncertainty ranking method. The approach models retrieval sample uncertainty in real time and introduces a learnable, adaptive threshold that triggers re-ranking upon detecting negative prediction shifts, thereby balancing stability and efficiency. Unlike static retrieval-augmented generation (RAG), this work pioneers an uncertainty-aware dynamic sample re-ranking framework that tightly integrates uncertainty modeling, reinforcement learning, and in-context learning (ICL). Evaluated on multi-domain QA benchmarks, the method outperforms the strongest baseline by 2.76% overall and improves accuracy on long-tail questions by 5.96%, significantly mitigating zero-shot failure.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Uncertainty Representations
• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval

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

Large language models (LLMs) can learn vast amounts of knowledge from diverse domains during pre-training. However, long-tail knowledge from specialized domains is often scarce and underrepresented, rarely appearing in the models' memorization. Prior work has shown that in-context learning (ICL) with retriever augmentation can help LLMs better capture long-tail knowledge, reducing their reliance on pre-trained data. Despite these advances, we observe that LLM predictions for long-tail questions remain uncertain to variations in retrieved samples. To take advantage of the uncertainty in ICL for guiding LLM predictions toward correct answers on long-tail samples, we propose a reinforcement learning-based dynamic uncertainty ranking method for ICL that accounts for the varying impact of each retrieved sample on LLM predictions. Our approach prioritizes more informative and stable samples while demoting misleading ones, updating rankings based on the feedback from the LLM w.r.t. each retrieved sample. To enhance training efficiency and reduce query costs, we introduce a learnable dynamic ranking threshold, adjusted when the model encounters negative prediction shifts. Experimental results on various question-answering datasets from different domains show that our method outperforms the best baseline by $2.76%$, with a notable $5.96%$ boost in accuracy on long-tail questions that elude zero-shot inference.

Problem

• Enhance long-tail knowledge retrieval in LLMs
• Reduce LLM prediction uncertainty in ICL
• Improve accuracy on specialized domain questions

Innovation

• Reinforcement learning dynamic ranking
• Learnable dynamic ranking threshold
• Enhanced retrieval-augmented in-context learning