Understanding and reasoning over long contexts is a crucial capability for language models (LMs). Although recent models support increasingly long context windows, their accuracy often deteriorates as input length grows. In practice, models often struggle to keep attention aligned with the most relevant context throughout decoding. In this work, we propose DYSCO, a novel decoding algorithm for improving long-context reasoning. DYSCO leverages retrieval heads--a subset of attention heads specialized for longcontext retrieval--to identify task-relevant tokens at each decoding step and explicitly up-weight them. By doing so, DYSCO dynamically adjusts attention during generation to better utilize relevant context. The method is training-free and can be applied directly to any off-the-shelf LMs. Across multiple instruction-tuned and reasoning models, DYSCO consistently improves performance on challenging long-context reasoning benchmarks, yielding relative gains of up to 25% on MRCR and LongBenchV2 at 128K context length with modest additional compute. Further analysis highlights the importance of both dynamic attention rescaling and retrievalhead guided selection for the effectiveness of the method, while providing interpretability insights into decoding-time attention behavior. Our code is available at https://github.com/princeton-pli/DySCO.
Figures & tables
Figure 1: Top: An illustrative Path Traversal task (simplified). Solving the task requires dynamically locating relevant context during decoding. Bottom: Accuracy as a function of context length for Qwen3-8B with and without DySCO . Despite the total context being only 16K tokens, both models exhibit severe performance degradation as context length increases.
Figure 2: Overview of DySCO algorithm. At each decoding step, DySCO consists of three stages: (1) Aggregation : We run a partial forward pass over the input sequence to obtain attentions of retrieval heads, such as QRHead , and use them to assign relevance scores to context tokens; (2) Selection : We use the relevance scores to select the important tokens; (3) Rescaling : We up-weight the important tokens by intervening attention logits of all attention heads and run a full forward pass to sample the next token.
Figure 3: Left : Performance of Qwen3-8B on Path Traversal as context length increases. Middle : Fractions that the gold edge appear among the top-5% edges ranked by attention score (sum of attention over all tokens in the span) from QRHead versus random heads. Right : Attention mass assigned to gold edges by QRHead and random heads. Despite severe performance degradation and a reduction in attention mass on gold edges, QRHead consistently allocates substantially higher attention to the gold edges.
Path Traversal
4K
8K
16K
32K
Qwen3-32B
80
67
32
21
+ UniAttnS
82
68
40
19
+ DySCO
79
71
46
33
Qwen3-8B
57
45
19
1
+ UniAttnS
59
46
25
2
+ DySCO
65
51
33
2
Table 1: Performance across different context lengths with Uniform Attention Scaling ( UniAttnS ) and DySCO .
Figure 4: Performance on MRCR , LongBenchV2 , and Clipper . DySCO substantially outperforms vanilla decoding, and UniAttnS . YaRN is applied to Qwen models at 128K context length, but not to Llama-3.1-8B-Instruct, which natively supports 128K.
MRCR
Hotpot
InfMC
InfQA
Avg
128K
128K
128K
128K
128K
Qwen3-4B YaRN
16.7
37
50
27.1
32.7
+ UniAttnS
18.4
34
39
24.1
28.9
+ DySCO
19.0
39
54
31.3
35.8
Qwen3-8B YaRN
16.7
61
61
39.6
44.6
+ UniAttnS
17.5
58
63
39.7
44.6
Table 2: Performance on long-context recall tasks with 128K input length. We apply YaRN for Qwen models.
MRCR
LongV2
Path
InfMC
64K
64K
16K
128K
Think
Yes
No
Yes
No
No
Vanilla
24.8
24.8
43.4
19
61
DySCO
27.3
26.3
51.2
33
66
w/ RandomHead
26.6
24.2
47.0
20
63
w/ StaticScaling
22.5
23.1
47.5
23
64
Table 3: Ablations of DySCO on head selection and dynamic rescaling with Qwen3-8B. DySCO outperforms its ablations with random attention heads and with static scaling.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
EAGER (Implementation)
Flash (Theoretical)
Vanilla
6.1
17.8
DySCO
3.6
10.5
Appendix
Table 4: Throughput (tokens/s) comparison between vanilla decoding and DySCO , under standard attention (EAGER) and an estimated Flash Attention implementation. Results are obtained on Qwen3-8B with 100K input tokens and 4K output tokens on an H100 GPU.
Figure 5: Comparison between DySCO , RAG (Stella), LongLLMLingua, and vanilla decoding. For RAG and LongLLMLingua, we report the results after reducing the context to different length (4K, 8K, and 16K tokens).
Figure 6: Performance of Qwen3-8B on MRCR and LongBenchV2 at 64K and 128K under four decoding settings: direct answer vs. Think, with vanilla decoding or DySCO . DySCO yields larger gains when combined with CoT.
p
K
β
MRCR
LongBenchV2
Vanilla
–
–
–
16.1
28.6
Default
0.975
8192
2.5
19.6
34.7
Vary p
0.95
8192
2.5
18.4
31.7
Varying K
0.975
2048
2.5
19.6
36.8
Appendix
Table 5: Evaluation of hyperparameter robustness of DySCO with Qwen3-8B. DySCO is robust to modest variations in the rescale strength β , selection constraint K , and selection threshold p .