Chain-of-Thought Matters: Improving Long-Context Language Models with Reasoning Path Supervision (2025)

Abstract: Recent advances in Large Language Models (LLMs) have highlighted the challenge of handling long-context tasks, where models need to reason over extensive input contexts to aggregate target information. While Chain-of-Thought (CoT) prompting has shown promise for multi-step reasoning, its effectiveness for long-context scenarios remains underexplored. Through systematic investigation across diverse tasks, we demonstrate that CoT’s benefits generalize across most long-context scenarios and amplify with increasing context length. Motivated by this, we propose a process-supervised framework that teaches models to generate high-quality reasoning paths for enhanced long-context performance. Our framework incorporates a self-sampling mechanism to bootstrap reasoning paths and a novel quality assessment protocol specifically designed for long-context scenarios. This protocol evaluates both answer correctness and process reliability, with the latter decomposed into source faithfulness and intrinsic consistency components for efficient and accurate assessment. Experimental results on various long-context benchmarks demonstrate the effectiveness of our approach, achieving significant improvements over outcome supervision baselines on both in-domain tasks (+13.6/+3.8 points for LLaMA/Qwen on MuSiQue) and cross-domain generalization (+9.3/+8.1 points on average across diverse QA tasks). Our code, data and trained models will be released upon acceptance.

## introduction large language models (llms) have revolutionized language modeling and achieved remarkable success in traditional nlp tasks @xcite . this success has spurred their application to increasingly complex real-world scenarios @xcite @xcite , particularly long-context tasks such as document-level question answering @xcite , summarization @xcite , multi-shot in-context learning @xcite , and repository-level code generation @xcite . these tasks typically require implicit reasoning steps that retrieve and aggregate information dispersed throughout extensive contexts before generating responses, posing significant challenges for contemporary long-context language models (lclms). as a means to elicit reasoning, chain-of-thought (cot) prompting has demonstrated remarkable efficacy in enhancing multi-step reasoning tasks @xcite @xcite . however, its effectiveness in long-context scenarios remains underexplored and only a few works have analyzed the use of cot from limited perspectives, such as context length @xcite , task difficulty @xcite , etc. building on prior work, we conduct first systematic investigation of cot's effectiveness across a diverse set of longcontext tasks of varying lengths and domains on both open-source and proprietary models of different scales. as illustrated in figure 1 (a), cot's benefits generalize across most long-context scenarios, and amplify with increasing context length. even more intriguingly, zero-shot majority voting results (figure 1 ) on the long-context multi-hop reasoning task musique @xcite demonstrate that cot substantially outperforms pure majority voting, and exhibits a steeper improvement trajectory as sampling rounds increase, but still significantly lags behind its oracle performance. this naturally leads to the following research question (rq): how to enhance models' capability in longcontext scenarios to generate high-quality reasoning paths foot_0 for improved performance? to address the question, we further investigate methods to train models in generating high-quality reasoning paths in long-context scenarios, aiming to enhance overall performance. given that realworld annotated data typically only contains outcome labels, we construct training data through a two-stage approach: first bootstrapping reasoning paths from the model itself, then identifying high-quality ones via a novel assessment protocol specifically designed for long-context scenarios. our protocol evaluates both answer correctness and process reliability, with the latter strategically decomposed into source faithfulness and intrinsic consistency for efficient and accurate assessment. these self-sampling and quality assessment mechanisms together form a process-supervised framework that not only guides models toward correct answers but also teaches them appropriate reasoning patterns to reach these conclusions. we train llama3.1-8b @xcite ) and qwen2.5-{7b,32b} @xcite on the musique dataset and comprehensively evaluate their performance across a diverse set of benchmarks, including: (1) musique, for assessing indomain performance; (2) selected qa tasks from longbenchv1 @xcite , for measuring generalization capabilities on both single and multidocument question answering; and (3) selected qa tasks from longbenchv2 @xcite , which feature longer contexts and more diverse domains. compared to the baseline method of outcome supervision, our proposed process-supervised framework demonstrates superior performance on both the in-domain musique dataset (+13.6,+3.8 and +7.5 points for llama/qwen) and other qa tasks (+9.3, +8.1 and +10.4 points on average). these results validate the effectiveness of our approach across long-context scenarios spanning various domains and lengths. our contributions can be summarized as follows: ## related work long-context language modeling context length is one of the most fundamental properties of language models, determining the total amount of information they can process at once. longcontext language modeling seeks to extend this capacity, pushing the boundaries of models' context windows. research for context window extension primarily follows two directions. the first is data-driven strategies, which focus on constructing training data that exhibits long-range dependency patterns @xcite @xcite @xcite @xcite @xcite . the second is architecture-driven approaches, which enhance models' potential for processing long contexts by modifying components such as position encodings @xcite @xcite and attention mechanisms @xcite . building upon these advances, current lclms have demonstrated remarkable performance on simple long-context tasks, such as needle-in-haystack retrieval @xcite @xcite . however, these models still struggle with more complex tasks, particularly long-context reasoning @xcite @xcite . motivated by this challenge, our work investigates how to improve models' longcontext reasoning capabilities by guiding them to generate high-quality chains of thought (cot). process supervision process supervision is a concept proposed to distinguish from outcome supervision @xcite . in process supervision, models are not only guided toward correct answers but also taught appropriate reasoning patterns to reach these conclusions. a significant application of process supervision is in reinforcement learning @xcite @xcite , where trained reward models for process evaluation, known as process reward models (prms), are used to supervise the execution process of policy models. numerous studies have demonstrated that process supervision outperforms outcome supervision in domains requiring multi-step reasoning, such as mathematics @xcite . beyond the reinforcement learning paradigm, many works directly utilize high-quality cot, obtained either through model's self-rejection sampling @xcite @xcite @xcite or sampling from other models, to train models' reasoning process via supervised finetuning @xcite @xcite . our approach falls into the latter category. most closely related to our work is sealong @xcite and real-world scenarios, covering a wide length range from 10k to 128k tokens. for synthetic scenarios, we select the single needle retrieval task (s-niah) from ruler @xcite and the multi-needle reasoning task with 3 supporting facts (mnr3) from babilong @xcite , which specifically focus on assessing models' retrieval and reasoning capabilities within extended contexts. since synthetic datasets allow flexible length control, we sample an equal number of examples at lengths of {16,32,64,128}k for testing. for real-world scenarios, we select three representative tasks from longbench-v2 @xcite : single-document question-answering (sqa), multi-document question-answering (mqa), and long in-context learning (licl), covering various domains including academic, literary, legal, finance, etc. for each task type, we retain test samples with lengths less than 128k. basic statistics of the curated evaluation dataset is presented in ## improving lclms via self-sampled reasoning paths section 3 consolidates the substantial benefits of cot in long-context scenarios. encouraged by these findings, we further explore how to enhance models' ability to generate high-quality cots in long-context scenarios, thereby improving their reasoning capabilities over long contexts. this section details the proposed process-supervised training framework, which comprises a self-sampling phase self-sampling fine-tuning prompt "reasoning": [think step by step] "answer": [your final answer] **context:**[long context] **question:**[related question] quality assessment sample #1 "reasoning": … … "answer": … … sample #2 "reasoning": … … "answer": … … sample #3 … … label answer correctness source faithfulness intrinsic consistency reject quality assessment sample #2 "reasoning": … … "answer": … … context question reject 0.95 sample #1 score : 0.80 sample #2 score : 0.95 sample #3 score : 0.90 top score sample #2 … … accept training data … … training data … … training data … … training set our framework sample #3 sample #1 figure 3 : our process-supervised framework for long-context tasks. we begin by sampling a diverse collection of n reasoning paths from the model. a quality assessment procedure consisting of three criteria is then applied to these samples to select high-quality training samples, which are then used for supervised fine-tuning. ## analysis impact of sampling size first, we analyze how the number of sampled candidate cots per training example affects model performance. as expected, a larger sampling size provides more diverse candidates, potentially enabling the selection of higherquality cots. figure 4 3 . impact of cot source models beyond selfsampling cots from the base model, we investigate the effectiveness of training with cots directly sampled from more capable models (gpt-4o-mini and gpt-4o), as shown in figure 5 . we observe that gpt-4o-generated cots perform best across scenarios: the in-domain musique task, other qa tasks from longbench v1 and v2, followed by gpt-4o-mini and then self-sampled cots. combined with our ablation results in ## conclusion in this study, we propose a process-supervised framework to enhance the long-context reasoning capabilities of llm by guiding them to generate high-quality reasoning paths. by conducting experiments on various long-context benchmarks, we demonstrate that our framework significantly improves model performance on both in-domain and out-domain tasks, highlighting the generalization capabilities of our method. our experiments underscore the critical role of process-supervised reasoning paths in enhancing long-context language models, enabling them to navigate and reason over extensive textual information with greater accuracy and reliability. by reducing the reliance on human annotations and leveraging self-sampled data, our approach offers a scalable solution for enhancing llms' reasoning abilities.