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

TL;DR

This paper introduces ProxyCoT, a training framework that enhances long-context reasoning in large language models by transferring reasoning skills from proxy contexts to full sequences, improving performance and generalization.

Contribution

ProxyCoT is a novel method that leverages chain-of-thought reasoning traces from proxy contexts to improve long-context reasoning in language models.

Findings

ProxyCoT outperforms strong baselines on various datasets.

Models trained with ProxyCoT generalize to out-of-domain tasks.

ProxyCoT reduces computational overhead compared to existing methods.

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.