Dual Dimensionality for Local and Global Attention

Abstract: Decoder-only Transformers compute attention over the KV cache of preceding tokens. Keys (and Values) are typically represented with the same dimensionality, regardless of its distance from the prediction target. In natural language, however, the next word is most strongly influenced by the immediately preceding tokens. We hypothesize that local and distant tokens impose asymmetric demands on representational capacity: local tokens are more critical for predicting immediate outputs and thus require richer representations, whereas distant tokens primarily serve as long-range memory, for which lower-dimensional representations may suffice. We formalize this idea as Distance-Adaptive Representation (DAR), implemented in a controlled setting that preserves full-dimensional representations within a local context window while assigning reduced-dimensional representations (e.g. 1/4 of the original dimensionality) to tokens beyond that window. Across multiple pretraining scales (70M to 410M parameters), as well as continued supervised fine-tuning on a 1B-scale model, this approach closely matches the performance of full-dimensional baselines. In contrast, uniformly reducing dimensionality across all token positions leads to worse performance. These results challenge the common assumption that key and value dimensionality should be uniform across token positions. Our findings suggest a new direction for designing attention architectures that adaptively allocate representational capacity across sequences, enabling further reductions in KV cache during inference.

• The paper demonstrates that Distance-Adaptive Representation (DAR) allocates full-dimensionality to local tokens while compressing distant tokens.
• Empirical studies show that reducing dimensionality for distant tokens achieves near-baseline performance and significant memory savings.
• DAR challenges uniform KV representations in Transformers, enabling scalable, resource-efficient attention mechanisms for long-context tasks.

Dual Dimensionality for Local and Global Attention

Problem Formulation and Hypothesis

The paper "Dual Dimensionality for Local and Global Attention" (2606.18587) investigates the architectural assumption in decoder-only Transformers that key/value (KV) representations for attention are uniformly high-dimensional across all context positions. The authors hypothesize that local tokens, which are proximate to the next-token prediction target, require richer, high-dimensional representations, while distant tokens may suffice with reduced dimensionality. This leads to the formalization of Distance-Adaptive Representation (DAR), a mechanism to differentially allocate representational capacity based on token distance in the sequence, challenging the conventional design of uniform KV dimensionality.

Distance-Adaptive Representation: Design and Implementation

DAR operates with a two-regime partition of sequence tokens under a sliding window:

• Local Regime: Tokens within a window $w$ of the prediction target retain full model dimensionality ($d$).
• Global Regime: Tokens outside the window are projected into a lower-dimensional space ($d_{\text{down}} < d$) via a lightweight linear projection.

Attention is computed over both representations. For distant tokens, their compressed representations are up-projected back to model dimensions to align with standard attention machinery, albeit their information content is bounded by $d_{\text{down}}$. This hybrid attention requires careful bookkeeping of per-token representations during training and inference, but enables substantial memory reductions during inference if only the compressed state is cached for distant tokens.

Empirical Evaluation: Language Modeling and Downstream Tasks

Pretraining Studies

The authors validate DAR across multiple scales (70M, 160M, 410M parameters), pretraining on 10B tokens from the Pile. Perplexity evaluations on Fine Web-Edu, WikiText-103, and C4 demonstrate:

• Substantial reduction of distant token dimensionality (up to $d/8$) yields minimal performance degradation compared to the full-dimensional baseline.
• Uniform reduction across all positions significantly worsens performance relative to DAR, underscoring the necessity of high-dimensional local representations.

For instance, at 70M scale, DAR with $d_{\text{down}}=128$ achieves a perplexity ratio of $99.61\%$ vs. baseline, whereas uniform reduction at $d_{\text{down}}=128$ sees $105.49\%$. Performance remains comparable across ablations of window size $w$ for $d$0, lending robustness to the dual regime design.

Scaling and Ablation

At higher scales (160M, 410M), DAR with $d$1 continues to match or outperform the full-dimensional baselines. Window size sweeps corroborate that only a modest number of proximate tokens require full dimensionality for optimal prediction.

Supervised Fine-Tuning and Task-Level Evaluation

DAR is further validated through continued supervised fine-tuning of OLMo-2-1B-SFT on instruction-tuning data, with downstream evaluation across MMLU, HellaSwag, CommonsenseQA, GSM8K, MBPP, and Multi-News. Moderate bottleneck configurations ($d$2 or $d$3) preserve downstream performance, with severe compression ($d$4) leading to nontrivial degradation. Relative scores are consistently high ($d$5 with $d$6, and performance remains stable for $d$7), confirming the core hypothesis for practical usage scenarios.

Comparison with Prior Work

Previous work on KV cache reduction primarily utilizes uniform compression (e.g., MLA, CSA, SKVQ) or sparsification via local windowing or heavy-hitter retention (e.g., Sliding Window Attention, H2O). These methods do not adapt representational capacity with token distance. DAR represents a distinct direction: an intrinsic, distance-aware allocation of capacity, decoupled from ad hoc resource constraints, and empirically justified across model scales and tasks.

Theoretical and Practical Implications

By demonstrating that representational capacity can be adaptively allocated, DAR provides direct evidence against the necessity of uniform high-dimensionality in attention KV caches. The findings suggest that long-range memory tokens can be stored with reduced dimensionality, unlocking substantial memory and compute savings for both training and inference. Practically, this enables scalable inference for long-context models and motivates architectural redesigns that further exploit context-based dimensional adaptation.

Future directions include:

• Full integration with hardware-efficient attention kernels and inference systems (e.g., FlashAttention, vLLM).
• Extension to order-of-magnitude larger models and diverse architectures.
• Exploring granular, non-window-based adaptation schemes and interaction-specific dynamic allocation.

Conclusion

The paper establishes DAR as a principled mechanism for allocating attention representation dimensionality based on token distance. Controlled experiments show that reducing dimensionality for distant tokens achieves near-baseline language modeling and downstream task performance, while uniform reduction significantly degrades efficacy. DAR challenges prevailing architectural assumptions, proposes direct evidence for asymmetric representational demands in attention, and sets the stage for context-adaptive, resource-efficient Transformer models.