Technical Summary: Long-Context Reasoning Through Proxy-Based Chain-of-Thought Tuning

Problem Statement

Recent large language models (LLMs) have expanded their context windows to support inputs of up to 10 million tokens. However, these models exhibit significant performance degradation on long-context tasks requiring complex reasoning, such as synthesizing insights from multiple medical reports or analyzing financial documents.

The core issue identified is a substantial performance disparity between full long contexts and proxy contexts. A proxy context is defined as a compact subset of the input containing only the minimal evidence necessary to derive the correct answer. While the underlying reasoning process required to solve a task remains invariant regardless of whether the model sees the full context or the proxy, empirical results show that models perform markedly better on proxy contexts. When conditioned on full long contexts, models often generate plausible high-level reasoning structures but fail to ground specific facts correctly, leading to hallucinations and lower accuracy.

Existing methods to improve reasoning, such as Chain-of-Thought (CoT) distillation and Reinforcement Learning (RL), face critical limitations in long-context settings:

1. CoT Distillation: Relies on high-quality traces from large teacher models. Querying these teachers over full long contexts is computationally expensive, slow, and often yields unreliable traces for complex tasks.
2. Reinforcement Learning: Scaling RL to long contexts is prohibitively expensive due to the need for extensive sampling over long sequences during training.

Methodology: ProxyCoT

The authors propose ProxyCoT, a novel two-stage training framework designed to transfer reasoning capabilities from short proxy contexts to full long contexts. The framework operates on the hypothesis that reasoning traces generated on efficient proxy inputs can be "grounded" into full long contexts via supervised fine-tuning.

The framework consists of two stages:

Stage 1: Acquiring CoTs on Short Proxy Contexts

The goal is to generate high-quality reasoning traces ($t$) conditioned on the proxy context ($C_p$) rather than the full context ($C$). This stage offers two variants for obtaining the traces:

• ProxyCoT-ZS (Zero-Shot): A large, off-the-shelf teacher model ($M_{ZS}$) is queried to generate reasoning traces based on the proxy context ($t \sim p_\phi(t | q, C_p)$). Only traces resulting in correct answers are retained. This avoids the cost of querying teachers on full long contexts.
• ProxyCoT-RL (Reinforcement Learning): If a capable teacher is unavailable, the target model ($M_{init}$) is optimized directly on proxy contexts using Reinforcement Learning with Verifiable Rewards (RLVR). The model learns to maximize the probability of generating traces that lead to correct answers ($t \sim p_\theta(t | q, C_p)$).

Stage 2: Grounding CoTs in Full Long Contexts

In this stage, the reasoning patterns learned in Stage 1 are transferred to the full context. A student model is trained via Supervised Fine-Tuning (SFT) to reproduce the reasoning traces generated in Stage 1, but conditioned on the full long context ($C$) instead of the proxy.

• Objective: Minimize the negative log-likelihood of the trace $t$ given the question $q$ and full context $C$:
  $$L_{SFT} = -E_{(q,C,C_p,t)} [\log p_\theta(t | q, C)]$$
• ProxyCoT-ZS: The student model ($M_{init}$) is fine-tuned to mimic the teacher's traces while seeing the full context.
• ProxyCoT-RL: The RL-optimized model ($M_{RL}$) from Stage 1 is further fine-tuned to apply its learned reasoning patterns to the full context.

Key Contributions

1. Formalization of Proxy Contexts: The paper introduces and formalizes the concept of proxy contexts for long-task reasoning, empirically revealing a significant performance gap between short proxy and full long contexts despite identical reasoning requirements.
2. ProxyCoT Framework: A novel training framework that leverages short proxy contexts to acquire high-quality CoT traces, which are then used to enhance reasoning over full long contexts. This approach circumvents the need for expensive teacher inference on long sequences and reduces RL sampling costs.
3. Empirical Validation: Extensive experiments demonstrate that ProxyCoT consistently outperforms strong baselines (including direct SFT, RL on full contexts, and CoT distillation on full contexts) across different model families (Qwen3-4B, Gemma3-4B) and datasets (SciTrek, HotpotQA). The method also shows strong generalization to out-of-domain tasks.

Experimental Results

The authors evaluated ProxyCoT on SciTrek (long-context QA over scientific articles) and HotpotQA (multi-hop QA over Wikipedia).

• Performance Gains:
  • On SciTrek, ProxyCoT-RL achieved a 46.5% exact match on full contexts with Qwen3-4B, significantly outperforming the zero-shot baseline (30.8%) and SFT on full contexts (19.5%). It even approached the performance of the much larger Qwen3-235B-Thinking model (48.8%).
  • On HotpotQA, ProxyCoT-RL achieved 52.7% accuracy on full contexts, outperforming all baselines.
• Efficiency: ProxyCoT-RL not only improved accuracy but also induced substantially shorter reasoning traces at inference time. On SciTrek, it used an average of 617 CoT tokens compared to 1,744 for zero-shot and 5,520 for ProxyCoT-ZS, while achieving higher accuracy.
• Ablation Studies: Results confirmed that both training stages are necessary. Training only on proxy contexts (Stage 1) yielded high proxy accuracy but poor full-context grounding. Training only on full contexts (Stage 2) without proxy-derived traces resulted in lower performance. The combination was essential for optimal results.
• Out-of-Domain Generalization: Models trained on SciTrek using ProxyCoT-RL showed improved performance on the Loong benchmark (financial reports and academic papers), demonstrating that the method enhances general long-context reasoning capabilities rather than memorizing task-specific patterns.

Significance and Claims

The paper claims that ProxyCoT offers a practical and efficient route to strengthening long-context reasoning without relying on teacher-generated traces over full long contexts. By decoupling the reasoning acquisition (on cheap, short proxies) from the reasoning grounding (on expensive, long contexts), the framework addresses the computational bottlenecks of current long-context training methods.

The authors position ProxyCoT as a solution for scenarios where training and supervising long-context models is challenging due to cost and data constraints. They note that while the approach assumes access to proxy contexts containing sufficient evidence, it provides a robust alternative to direct long-context reinforcement learning or distillation, effectively bridging the performance gap between short and long inputs.