Dual Dimensionality for Local and Global Attention
This work challenges the uniform key/value dimensionality assumption in attention mechanisms, offering a new direction for efficient Transformer architectures that adaptively allocate representational capacity based on token distance.
The paper proposes Distance-Adaptive Representation (DAR) for decoder-only Transformers, which uses full-dimensional representations for local tokens and reduced-dimensional representations for distant tokens. Across models from 70M to 410M parameters and fine-tuning on a 1B-scale model, DAR matches full-dimensional baselines while enabling KV cache reduction, whereas uniform dimensionality reduction performs worse.
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.