DySCO: Dynamic Attention-Scaling Decoding for Long-Context Language Models

Paper: 2602.22175 Authors: Xi Ye, Wuwei Zhang, Fangcong Yin, Howard Yen, Danqi Chen Categories: cs.CL

Abstract

Long-context understanding remains a critical challenge for language models despite increasing context window sizes. This paper introduces DySCO (Dynamic Attention-Scaling Decoding), a novel training-free decoding algorithm that addresses the attention misalignment problem in long-context reasoning. DySCO identifies and leverages “retrieval heads”—specialized attention heads that excel at locating relevant information across long contexts—to dynamically up-weight task-relevant tokens during each decoding step. This approach enables models to maintain focus on pertinent context throughout generation without requiring additional training. Evaluated on challenging benchmarks like MRCR and LongBenchV2 at 128K context length, DySCO demonstrates consistent improvements across multiple instruction-tuned and reasoning models, achieving relative performance gains of up to 25% with only modest computational overhead.

Key Contributions

• Training-free decoding algorithm: DySCO can be applied directly to any pre-trained language model without fine-tuning or additional training
• Retrieval head identification: Systematic method for identifying attention heads specialized in long-context retrieval across different model architectures
• Dynamic attention rescaling: Novel mechanism that adjusts attention weights at each decoding step based on retrieval-head guidance
• Substantial performance gains: Up to 25% relative improvement on long-context reasoning benchmarks at 128K tokens with modest compute overhead
• Interpretability insights: Analysis revealing how attention behavior changes during decoding and why certain heads are effective for retrieval

Methodology

DySCO operates through a two-stage process: retrieval head identification and dynamic attention scaling during decoding.

Retrieval Head Identification: The method first identifies which attention heads in a model are specialized for long-context retrieval. This is done by analyzing attention patterns on validation examples where the model must locate specific information within long contexts. Heads that consistently assign high attention weights to ground-truth relevant tokens are classified as retrieval heads. This identification is performed once per model and can be reused across different tasks.

Dynamic Attention Scaling: During generation, at each decoding step $t$, DySCO:

1. Uses the identified retrieval heads to compute relevance scores for all input tokens
2. Aggregates these scores to identify the most task-relevant tokens
3. Applies an up-weighting factor to the attention logits for these tokens across all attention heads
4. Proceeds with standard decoding using the rescaled attention distribution

The scaling factor $\alpha$ controls the degree of up-weighting and can be tuned based on the task. The key insight is that retrieval heads provide a reliable signal for which context is relevant, and explicitly amplifying attention to these regions helps the model maintain focus throughout generation.

Experimental Results

DySCO was evaluated on multiple long-context benchmarks spanning different reasoning types:

MRCR (Multi-hop Reasoning over Contexts): On this challenging multi-hop QA benchmark at 128K context length, DySCO achieved relative improvements of 20-25% across models including Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct. The gains were particularly pronounced on examples requiring multiple reasoning steps across distant context segments.

LongBenchV2: Across the diverse tasks in this benchmark suite, DySCO showed consistent improvements of 10-18% relative gain. Performance improvements were observed across question answering, summarization, and code completion tasks, demonstrating the method’s generality.

Scaling with Context Length: Experiments at 32K, 64K, and 128K tokens revealed that DySCO’s benefits increase with context length. At 32K, improvements were modest (5-8%), while at 128K, gains reached 20-25%, suggesting the method is particularly valuable for extremely long contexts where attention misalignment is most severe.

Computational Overhead: DySCO adds approximately 15-20% to inference time, primarily from the additional forward pass needed to compute retrieval head scores. This overhead is modest compared to the performance gains achieved.

Analysis and Insights

Ablation Studies: The authors conducted extensive ablations showing that both components—retrieval head guidance and dynamic rescaling—are essential. Using random heads instead of retrieval heads reduced gains by 60-70%, while static (non-dynamic) rescaling was only half as effective as the full method.

Attention Visualization: Analysis of attention patterns revealed that without DySCO, models often “drift” during generation, gradually shifting attention away from relevant context toward more recent tokens. DySCO counteracts this drift by continuously reinforcing attention to task-relevant regions identified by retrieval heads.

Retrieval Head Characteristics: Across different model families, retrieval heads were found predominantly in middle-to-late layers and exhibited distinct attention patterns—they tend to have more uniform attention distributions over long ranges rather than focusing on local context. This suggests these heads have learned to perform a different computational role than typical attention heads.

Failure Cases: DySCO shows diminishing returns when the base model’s retrieval heads are poorly calibrated or when the task requires reasoning over information that is uniformly distributed across the entire context rather than concentrated in specific segments.

Takeaways

1. Attention misalignment during decoding is a significant bottleneck for long-context reasoning, and explicitly managing attention can yield substantial improvements without model retraining.

2. Language models naturally develop specialized “retrieval heads” that can be identified and leveraged to guide attention toward relevant context during generation.

3. Dynamic, per-step attention rescaling is more effective than static approaches, as relevance shifts throughout the generation process.

4. The method’s effectiveness scales with context length, making it particularly valuable for applications requiring extremely long context windows (100K+ tokens).

5. Training-free decoding algorithms like DySCO offer a practical path to improving long-context capabilities of existing models without the computational cost of retraining or fine-tuning.