Long-Context Reasoning Through Proxy-Based Chain-of-Thought Tuning

Miao Li; Irina Saparina; Alexander Gurung; Mirella Lapata

cs.CLApr 6, 2026

Long-Context Reasoning Through Proxy-Based Chain-of-Thought Tuning

Authors: Miao Li, Irina Saparina, Alexander Gurung, Mirella Lapata

Organizations: School of Informatics, University of Edinburgh

Abstract

Recent large language models support inputs of up to 10 million tokens, yet they perform poorly on long-context tasks that require complex reasoning. Such tasks can be solved using only a subset of the input -- a proxy context -- rather than the full sequence. Despite sharing the same underlying reasoning process, models exhibit a significant performance disparity between proxy and full contexts. To improve long-context reasoning, we propose ProxyCoT, a novel training framework that transfers reasoning capabilities from short proxy contexts to full long contexts. Specifically, we first obtain high-quality chain-of-thought reasoning traces on proxy contexts through reinforcement learning or distillation from a larger teacher model, and then ground the generated traces in full long contexts with supervised fine-tuning. Experiments across different datasets demonstrate that ProxyCoT consistently outperforms strong baselines with reduced computational overhead. Furthermore, models trained with ProxyCoT generalize their long-context reasoning capabilities to out-of-domain tasks.

Figures & tables

Figure 1: The disparity of model performance in the zero-shot setting on SciTrek Li et al. (2025a) when showing full long contexts vs short proxy contexts in terms of exact match. The full context includes 128K tokens, while there are only around 650 tokens on average in a proxy context. Models work better on proxy contexts, while the proxy context requires the same reasoning process as their corresponding full long context.

Panel of Figure 2 — Figure 2: General two-stage pipeline of ProxyCoT (left), and two instantiations (right): ProxyCoT-ZS and ProxyCoT-RL. Given a target model $M_{\text{init}}$ , ProxyCoT-ZS employs a large off-the-shelf model $M^{\mathcal{ZS}}$ as the teacher to generate CoTs from proxy contexts, and then fine-tunes $M_{\text{init}}$ as the student to generate CoTs on corresponding long contexts. ProxyCoT-RL first optimizes $M_{\text{init}}$ using RLVR to obtain CoTs on proxy contexts, and subsequently fine-tunes the RL-optimized model $M^{\mathcal{RL}}$ as the student to ground these CoTs in corresponding long contexts.

Specification	Qwen3-4B	Gemma3-4B
Number of Parameters	4B	4B
Native Context Size	256K	32K
Supporting Context Size	256K	128K
Attention Architecture	Dense	Sparse
Release Date	July 2025	March 2025

Table 1: Details of experimental models: Qwen3-4B-Instruct-2507 and Gemma3-4B-IT. Gemma3-4B-IT supports the 128K-token context based on RoPE scaling.

SciTrek
$q$	What is the smallest number of authors for any article in the collection?
$C$	{7 full-text scientific articles}
$C^{p}$	Article title: Existence and uniqueness for Legendre curves There are 6 words in the title (separated by spaces). There are 2 authors: Tomonori Fukunaga, Masatomo Takahashi There are 9 references in the reference section. The other provided articles are not cited by this article. { we cut down the text because of limited space. } Article title: Effect of higher-order interactions on synchronization of neuron models with electromagnetic induction There are 12 words in the title (separated by spaces). There are 4 authors: Mohanasubha Ramasamy, Subhasri Devarajan, Suresh Kumarasamy, Karthikeyan Rajagopal There are 37 references in the reference section. The other provided articles are not cited by this article.
HotpotQA
$q$	Luis Gianneo was teacher of which chief exponent of Argentine folk music?
$C$	{67 full-text Wikipedia articles}

Table 2: Example questions (

q

), abbreviated long contexts(

C

), and their proxy contexts (

C^{p}

) for SciTrek (top) and HotpotQA (bottom).

Dataset	Full Context	Proxy	Question	Answer
SciTrek	83,018.3	659.4	18.9	16.7
HotpotQA	77,764.4	301.1	24.5	4.5

Table 3: Dataset statistics for SciTrek and HotpotQA: average number of tokens in full contexts, proxy contexts, questions, and ground-truth answers, based on the tokenizer of Qwen3-4B-Instruct.

Model	Training Strategy	Proxy $\uparrow$	Full $\uparrow$
Qwen3-235B-Instruct	Zero-shot	72.7	45.2
Qwen3-235B-Thinking	Zero-shot	85.6	48.8
Qwen3-4B-Instruct	Zero-shot	67.2	30.8
	SFT on $C$	39.0	19.5
	RLVR on $C$	66.1	32.9
	SFT on $C$ , CoT ∗	45.8	31.6

Table 4: Performance of Qwen3-4B-Instruct and Gemma3-4B-IT on SciTrek, in terms of exact match (%). Models are evaluated with Proxy and Full contexts as input.

C

: full long contexts, and CoT ∗ : chain-of-thought reasoning traces generated by Qwen3-235B-A22B-Thinking on full long contexts.

Model	Training Strategy	Proxy $\uparrow$	Full $\uparrow$
Qwen3-235B-Instruct	Zero-shot	92.1	60.8
Qwen3-235B-Thinking	Zero-shot	93.2	50.7
Qwen3-4B-Instruct	Zero-shot	91.3	44.5
	SFT on $C$	92.6	48.8
	RLVR on $C$	88.6	48.1
	SFT on $C$ , CoT ∗	84.5	40.2

Table 5: Performance of Qwen3-4B-Instruct on HotpotQA evaluated with GPT5-mini as judge in terms of accuracy (%). Models are evaluated with Proxy and Full contexts as input.

C

: full long contexts, and CoT ∗ : chain-of-thought reasoning traces from Qwen3-235B-A22B-Thinking on full long contexts.

Training Strategy	CoT Tokens $\downarrow$	EM $\uparrow$
Zero-shot	1,744	30.8
RLVR on $C$	937	32.9
SFT on $C$ , CoT ∗	6,683	31.6
ProxyCoT-ZS	5,520	38.8
ProxyCoT-RL	617	46.5

Table 6: Models trained with different training strategies use different numbers of chain-of-thought (CoT) tokens when showing full long contexts. Results are based on Qwen3-4B-Instruct on SciTrek. CoT ∗ : chain-of-thought reasoning traces are generated by Qwen3-235B-A22B-Thinking on full long contexts.

Model	Training	Financial $\uparrow$	Academic $\uparrow$
Qwen3-4B-Instruct	Zero-shot	37.76	24.91
Qwen3-4B-Instruct	ProxyCoT-RL	40.83	42.51
Gemma3-4B-IT	Zero-shot	25.85	3.55
Gemma3-4B-IT	ProxyCoT-RL	32.05	24.32

Table 7: Results on Loong (domains of Financial Reports and Academic Papers) when the model is zero-shot or trained on SciTrek with ProxyCoT-RL, in terms of answer quality (scored on a 1–100 scale by GPT-5 mini as the judge).

Model	(RLVR)	(SFT)	Proxy $\uparrow$	Full $\uparrow$
	Stage 1	Stage 2
Qwen3-4B-Instruct	✓	✓	88.5	46.5
	✓	✗	91.5	29.0
	✗	✓	77.5	46.3
Gemma3-4B-IT	✓	✓	69.8	43.7
	✓	✗	88.5	8.0

Table 8: Ablations on ProxyCoT-RL using Qwen3-4B-Instruct and Gemma3-4B-IT with SciTrek. We report results in terms of exact match (%) on proxy and full contexts. There are two stages in training: reinforcement learning with verifiable rewards (RLVR) in Stage 1, and supervised fine-tuning (SFT) in Stage 2.

	Proxy Type	#Tokens	Perf. $\uparrow$
	Random sentences from full context	1,102	3.4
	Titles, authors and references of the context articles	20,344	24.6
SciTrek	Descriptions of the structured metadata (provided by SciTrek)	659	91.5
	Random sentences from full context	385	15.8
	Sentences semantically related to the question	364	35.1
HotpotQA	Ground-truth evidence sentences (provided by HotpotQA)	301	92.2

Table 9: Performance of Qwen3-4B-Instruct after RLVR with different types of proxy contexts on SciTrek and HotpotQA in terms of exact match (%) and accuracy (%, evaluated with GPT5-mini as judge), respectively.

Dataset	Oracle: Noise	#Tokens	Performance $\uparrow$
SciTrek	1:5	5,471	85.3
	1:2	2,588	89.9
	1:1	1,622	91.3
	1:0	659	91.5
HotpotQA	1:5	1,853	83.7
	1:2	922	88.4

Table 10: Performance of Qwen3-4B-Instruct trained with ProxyCoT-RL on noisy proxy contexts: human-annotated proxies augmented with randomly sampled sentences as noise. Results on SciTrek are reported as exact match (%); on HotpotQA as accuracy (%), judged by GPT5-mini. The Oracle:Noise column gives the ratio of oracle evidence sentences to injected noise sentences; 1:0 denotes no injected noise.

Long-Context Reasoning in Language Models

See all 41 papers

PERK: Long-Context Reasoning as Test-Time Learning
Jul 8, 2025Zeming Chen, Angelika Romanou, Gail Weiss +1Memory-Augmented Language Models Long-Context Modeling
EPFL
ReContext: Recursive Evidence Replay as LLM Harness for Long-Context Reasoning
Jul 2, 2026Yanjun Zhao, Ruizhong Qiu, Tianxin Wei +6Long-Context Inference Long-Context Reasoning in Language Models
University of Illinois Urbana-Champaign
LongR: Unleashing Long-Context Reasoning via Reinforcement Learning with Dense Utility Rewards
Feb 5, 2026Qize Yu, Zijun Chen, Chen Wu +6Long-Context Reasoning in Language Models RL for Language Model Reasoning
Shanghai Jiao Tong University · Tsinghua University · Beijing University of Posts and Telecommunications, Beijing, China +1

Jul 8, 2025J/K move · Enter open · S save

PERK: Long-Context Reasoning as Test-Time Learning

Zeming Chen, Angelika Romanou, Gail Weiss, Antoine Bosselut

EPFL

Memory-Augmented Language Models Long-Context Modeling Test-Time Training Long-Context Inference Long-Context Reasoning in Language Models

Open paper AskPDF arXiv

Long-context reasoning requires accurately identifying relevant information in extensive, noisy input contexts. In this work, we propose PERK (Parameter Efficient Reasoning over Knowledge), a scalable approach for learning to encode long contexts using gradient updates at test time. Specifically, PERK employs two nested optimization loops in a meta-training phase. The inner loop rapidly encodes contexts into a low-rank adapter (LoRA) that serves as a parameter-efficient memory module for the base model. Concurrently, the outer loop learns to use the updated adapter to accurately recall and reason over relevant information from the encoded long context. Our evaluations on several long-context reasoning tasks show that PERK significantly outperforms the standard long-context finetuning, achieving average absolute performance gains of up to 20% for Qwen-2.5 (0.5B & 7B) on synthetic and real-world long-context reasoning. PERK also maintains its advantages across model scales and families. Compared to specialized long-context LLMs, PERK matches or surpasses their performance. Finally, our analyses show PERK is more robust to reasoning complexity, length extrapolation, and the positions of relevant information in contexts. https://perk-long-context.web.app

Abstract

Long-Context Reasoning in Language Models

Explore similar work

PERK: Long-Context Reasoning as Test-Time Learning

ReContext: Recursive Evidence Replay as LLM Harness for Long-Context Reasoning

LongR: Unleashing Long-Context Reasoning via Reinforcement Learning with Dense Utility Rewards