用强化学习选关键段落,让大模型更懂长文本推理。
Chunks as Arms: Multi-Armed Bandit-Guided Sampling for Long-Context LLM Preference Optimization
- 将上下文分块视为老虎机臂,动态选择最相关段落输入模型。
- 在Llama和Qwen上实现长文本推理提升超4%。
- 适合需要优化长文本理解的AI研发者使用。
长上下文建模对问答、摘要和复杂推理等实际任务至关重要。近期研究尝试通过合成数据微调大语言模型以增强其长上下文能力,但生成数据多样性低且存在事实错误,限制了效果。为此,我们提出LongMab框架,采用多臂老虎机(MAB)策略,在生成过程中动态选择最能激发高质量响应的上下文片段,构建用于直接偏好优化(DPO)的偏好数据对。具体地,将上下文分块视为MAB中的“臂”,根据预期奖励分数选择输入模型生成回复,并基于奖励反馈迭代更新分数。探索与利用并行的推演机制使模型聚焦于最相关的上下文部分,从而生成多样且高质量的响应。在Llama和Qwen上的实验表明,LongMab在长文本推理基准上实现超过4%的性能提升。所有数据与代码将开源至https://github.com/NEUIR/LongMab-PO。
原文摘要 · Abstract (English)
Long-context modeling is critical for a wide range of real-world tasks, including long-context question answering, summarization, and complex reasoning tasks. Recent studies have explored fine-tuning Large Language Models (LLMs) with synthetic data to enhance their long-context capabilities. However, the effectiveness of such approaches is often limited by the low diversity and factual inconsistencies in the generated data. To address these challenges, we propose LongMab, a novel framework that leverages a Multi-Armed Bandit (MAB) rollout strategy to identify the most informative chunks from the given long context for sampling high-quality and diverse responses and constructing preference data pairs for Direct Preference Optimization (DPO) training. Specifically, we treat context chunks as arms of MAB, select chunks based on their expected reward scores to input into LLMs to generate responses, and iteratively update these scores based on reward feedback. Both exploration and exploitation during the rollout process enable the LLM to focus on the most relevant context segments, thereby generating and collecting high-quality and diverse responses. Experimental results on both Llama and Qwen show the effectiveness of LongMab by achieving more than a 4% improvement on long-context reasoning benchmarks. All data and code will be released on https://github.com/NEUIR/LongMab-PO.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。