# Chunks as Arms: Multi-Armed Bandit-Guided Sampling for Long-Context LLM Preference Optimization

arXiv: https://arxiv.org/abs/2508.13993 · HTML: https://mlsanity.com/paper/2508.13993/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in LLM preference optimization methods, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 35 of its 47 references that are in the map.
- First posted on arXiv 19 Aug 2025.

## Placed in

- LLM alignment / RLHF / preference optimisation / [LLM preference optimization methods](https://mlsanity.com/alignment/llm-preference-optimization-methods/)

## Builds on (most cited first, 30 of 35)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [SWE-bench: Can Language Models Resolve Real-World GitHub Issues?](https://mlsanity.com/paper/2310.06770/) (2023)
- [Self-Instruct: Aligning Language Models with Self-Generated Instructions](https://mlsanity.com/paper/2212.10560/) (2022)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [GLM: General Language Model Pretraining with Autoregressive Blank Infilling](https://mlsanity.com/paper/2103.10360/) (2021)
- [MuSiQue: Multihop Questions via Single-hop Question Composition](https://mlsanity.com/paper/2108.00573/) (2021)
- [Large Language Models Can Be Easily Distracted by Irrelevant Context](https://mlsanity.com/paper/2302.00093/) (2023)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [An Empirical Study of Catastrophic Forgetting in Large Language Models During Continual Fine-tuning](https://mlsanity.com/paper/2308.08747/) (2023)
- [The NarrativeQA Reading Comprehension Challenge](https://mlsanity.com/paper/1712.07040/) (2017)
- [LongLLMLingua: Accelerating and Enhancing LLMs in Long Context Scenarios via Prompt Compression](https://mlsanity.com/paper/2310.06839/) (2023)
- [∞Bench: Extending Long Context Evaluation Beyond 100K Tokens](https://mlsanity.com/paper/2402.13718/) (2024)
- [Effective Long-Context Scaling of Foundation Models](https://mlsanity.com/paper/2309.16039/) (2023)
- [Make Your LLM Fully Utilize the Context](https://mlsanity.com/paper/2404.16811/) (2024)
- [LV-Eval: A Balanced Long-Context Benchmark with 5 Length Levels Up to 256K](https://mlsanity.com/paper/2402.05136/) (2024)
- [ChatQA 2: Bridging the Gap to Proprietary LLMs in Long Context and RAG Capabilities](https://mlsanity.com/paper/2407.14482/) (2024)
- [SQuALITY: Building a Long-Document Summarization Dataset the Hard Way](https://mlsanity.com/paper/2205.11465/) (2022)
- [LongReward: Improving Long-context Large Language Models with AI Feedback](https://mlsanity.com/paper/2410.21252/) (2024)
- [RAG-DDR: Optimizing Retrieval-Augmented Generation Using Differentiable Data Rewards](https://mlsanity.com/paper/2410.13509/) (2024)
- [Never Lost in the Middle: Mastering Long-Context Question Answering with Position-Agnostic Decompositional Training](https://mlsanity.com/paper/2311.09198/) (2023)
- [Large Language Models Can Self-Improve in Long-context Reasoning](https://mlsanity.com/paper/2411.08147/) (2024)
- [Making Long-Context Language Models Better Multi-Hop Reasoners](https://mlsanity.com/paper/2408.03246/) (2024)
- [LongPO: Long Context Self-Evolution of Large Language Models through Short-to-Long Preference Optimization](https://mlsanity.com/paper/2502.13922/) (2025)
- [Chain-of-Thought Matters: Improving Long-Context Language Models with Reasoning Path Supervision](https://mlsanity.com/paper/2502.20790/) (2025)
- [LongFaith: Enhancing Long-Context Reasoning in LLMs with Faithful Synthetic Data](https://mlsanity.com/paper/2502.12583/) (2025)
- [Forgetting Curve: A Reliable Method for Evaluating Memorization Capability for Long-context Models](https://mlsanity.com/paper/2410.04727/) (2024)

## Built on by (most cited first, 3 of 3)

- [Autoencoding-Free Context Compression for LLMs via Contextual Semantic Anchors](https://mlsanity.com/paper/2510.08907/) (2025)
- [GOLD PANNING: Strategic Context Shuffling for Needle-in-Haystack Reasoning](https://mlsanity.com/paper/2510.09770/) (2025)
- [A Component-Based Survey of Interactions between Large Language Models and Multi-Armed Bandits](https://mlsanity.com/paper/2601.12945/) (2026)

## Abstract (arXiv)

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.