Dual Dimensionality for Local and Global Attention

2026-06-16

AI Summary

This work addresses the limitation of conventional Transformers, which employ uniform-dimensional key/value (KV) representations for all historical tokens, thereby neglecting the differing contributions of local and distant context to prediction. The authors propose Distance-Adaptive Representation (DAR), a method that retains full-dimensional KV representations within a local window while using reduced-dimensional representations for tokens beyond this window. This approach provides the first systematic validation that local and global tokens can be effectively represented with different dimensions, challenging the prevailing assumption that KV dimensions must remain consistent across all tokens. Experiments on decoder-only architectures demonstrate that DAR matches the performance of full-dimensional baselines across pretrained models ranging from 70M to 410M parameters and a 1B fine-tuned model, significantly outperforms global dimensionality reduction schemes, and offers potential reductions in KV cache memory during inference.

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.

Problem

• KV cache
• attention mechanism
• dimensionality allocation
• local and global attention
• Transformer

Innovation

• Distance-Adaptive Representation
• KV cache compression
• decoder-only Transformer
• local-global attention
• adaptive dimensionality