Yu Zhao, Zekun Zhang, Fan Jiang, Bo Zeng, Xintong Wang, Chonghan Liu, Shimin Shan, Yichao Du 等 9 位作者
Long chain-of-thought (CoT) reasoning models such as DeepSeek-R1 keep lengthening their inference chains under test-time scaling, yet standard self-attention incurs quadratic cost in context length, making long-CoT deployment costly in production. We therefore introduce LISA (Linear-Indexed Sparse Attention), a plug-and-play replacement for self-attention that requires no full pretraining. LISA combines a linear-attention branch for long-range memory with a Lightning Indexer that dynamically selects the most relevant tokens for a sparse self-attention branch. The two branches are fused through a learned gate, so the model retains broad context while focusing computation on informative tokens. We train LISA in two stages. In Stage 1, a linear-attention branch is integrated with sliding-window attention and distilled from a frozen teacher's full-attention distribution. In Stage 2, the Lightning Indexer replaces the static window and is trained with a per-head KL-divergence objective to mimic teacher attention patterns. The resulting model reduces eager-mode self-attention computation from quadratic in sequence length to linear in the number of generated tokens and the selected context budget. Experiments on DeepSeek-distilled Qwen models show that LISA cuts measured eager-mode mean sample latency by 23.9% at 16K context without degrading long-context reasoning quality. Across nine mathematical and code reasoning benchmarks, the gated LISA-1.5B model improves the equal-weight benchmark average by 4.2 percentage points relative to the full-attention baseline.
Yu Zhao, Fan Jiang, Tianle Liu, Bo Zeng, Xintong Wang, Jiaxuan Teng, Yu Liu, Yichao Du
Recent advancements in Large Reasoning Models (LRMs), exemplified by DeepSeek-R1, have underscored the potential of scaling inference-time compute through Group Relative Policy Optimization (GRPO). However, GRPO frequently suffers from gradient signal attenuation when encountering problems that are either too trivial or overly complex. In these scenarios, the disappearance of inter-group advantages makes the gradient signal susceptible to noise, thereby jeopardizing convergence stability. While variants like DAPO attempt to rectify gradient vanishing, they do not alleviate the substantial computational overhead incurred by exhaustive rollouts on low-utility samples. In this paper, we propose Difficulty-Estimated Policy Optimization (DEPO), a novel framework designed to optimize the efficiency and robustness of reasoning alignment. DEPO integrates an online Difficulty Estimator that dynamically assesses and filters training data before the rollout phase. This mechanism ensures that computational resources are prioritized for samples with high learning potential. Empirical results show that DEPO achieves approximately 1.86x faster rollouts than DAPO in the profiled configuration, while improving benchmark performance over GRPO at comparable training cost. Our approach significantly lowers the computational barrier for training high-performance reasoning models, offering a more sustainable path for reasoning scaling.