On the token distance modeling ability of higher R o PE attention dimension (2024)

Abstract: Length extrapolation algorithms based on Rotary position embedding (RoPE) have shown promising results in extending the context length of language models. However, understanding how position embedding can capture longer-range contextual information remains elusive. Based on the intuition that different dimensions correspond to different frequency of changes in RoPE encoding, we conducted a dimension-level analysis to investigate the correlation between a hidden dimension of an attention head and its contribution to capturing long-distance dependencies. Using our correlation metric, we identified a particular type of attention heads, which we named Positional Heads, from various length-extrapolated models. These heads exhibit a strong focus on long-range information interaction and play a pivotal role in long input processing, as evidence by our ablation. We further demonstrate the correlation between the efficiency of length extrapolation and the extension of the high-dimensional attention allocation of these heads. The identification of Positional Heads provides insights for future research in long-text comprehension.

## introduction the transformer model has revolutionized natural language processing tasks, but it demonstrates limitations in modeling long sequences. meanwhile, models like mamba @xcite that excel in capturing long-range dependencies struggle to meet the practical requirements of natural language modeling @xcite . consequently, there has been a recent surge of work focused on extending the context length in language models based on the transformer architecture @xcite @xcite . particularly, some of these efforts that leverage and enhance the capabilities of rope (rotary positional embedding) @xcite @xcite , have shown promising results in extrapolating the model's capacity to handle longer contexts @xcite . open-source large language models commonly employ rotary positional embedding (rope) to model sequence positional information @xcite @xcite . rope exhibits two desirable properties. firstly, its exponential positional encoding introduces long-range attention decay, allowing the model to focus more on neighboring semantic information. secondly, by utilizing trigonometric functions to differentiate frequencies, rope effectively captures different distances between tokens, enabling higher attention scores for tokens with longer semantic dependencies, facilitating semantic aggregation. when compared to length extrapolation methods based on sparse attention @xcite or prompt compression @xcite , modifications to rope for length extrapolation do not result in the loss of fine-grained contextual information at a global level. therefore, it possesses distinct advantages in tasks such as long text comprehension @xcite , where the preservation of comprehensive contextual information is essential for practical applications. a prevailing viewpoint suggests that language models based on rope encounter out-ofdistribution (ood) issues when faced with contexts longer than the pre-training text length, specifically affecting the sampling of the trigonometric function component for token distances @xcite . as a result, related studies have adjusted the attention resolution in the context of long texts and fine-tuned the model to adapt to longer token distances. we hypothesize that the effectiveness of such methods stems from rope's ability to decouple information from different distances by representing them through different dimensions with varying rotational frequencies. however, this pattern has not been thoroughly observed and analyzed in the inference process. our paper presents a novel approach by examining the impact of each dimension of attention heads that use rope on modeling text distances. we empirically validate the claim in yarn that lowerfrequency dimensions are responsible for modeling longer text dependencies. furthermore, we discover that not all attention heads exhibit this characteristic, emphasizing the importance of heads that possess this relationship in modeling long texts. our study explores rope's potential for long text modeling from a frequency perspective, shedding light on the relationship between dimensions and text modeling capabilities. our primary findings are as follows: ## defining dimension contribution in rope in rotary position embedding (rope), each dimension of the vectors q m and k n contributes to the attention score via their dot product. to thoroughly investigate the role of different dimensions in rope for semantic modeling, we utilize an algorithm that analyzes the contribution of each dimension to the attention scores. to capture the contribution of each dimension, we employ the hadamard product, i.e., elementwise multiplication, denoted by the symbol ⊙: in rope, every two dimensions correspond to trigonometric functions with the same frequency θ i . we sum the values of these corresponding dimensions to form new vectors: the value of g i reflects the contribution of each dimension in rope to the attention score. a higher value indicates a greater contribution of that dimension to the attention score, where: here, θ i represents the positional encoding frequency for the i-th dimension. the dot product of q m and k n can be expressed as: therefore, we use the value g i to measure the contribution of θ i to the attention score. this methodological framework enables a comprehensive analysis of how each dimension in rotary position embedding (rope) contributes to the attention scores. through this approach, we can delve into the role of different dimensions in rope for semantic modeling. ## related work handling longer contexts in transformer models has seen significant improvements through various methods, including enhanced training techniques, innovative frameworks, memory mechanisms, and adjustments to positional encoding, including enhanced training techniques @xcite , external summary designs @xcite , memory mechanisms @xcite , and adjustments to positional encoding. among these methods, modifying positional encoding stands out due to its simplicity and efficacy. some methods manipulate the token position numbering itself, as seen in pi @xcite and selfextend @xcite . others make adjustments within the encoding layers at the level of rotational positional encoding, exemplified by works like yarn @xcite and clex @xcite . additionally, novel positional encodings, such as cope @xcite , have been proposed to generalize rotational positional encoding and further enhance long-text capabilities. certain studies have delved into the impact of positional encoding in depth. some research indicates that the initial token's position is crucial in longtext contexts @xcite , while other work @xcite highlights that the base of rotational positional encoding can limit a model's capacity to handle long texts. @xcite proposes a comprehensive framework to describe length extrapolation. however, the role of different dimensions within rotational positional encoding for information interaction remains underexplored. moreover, the precise mechanisms by which positional encoding affects information interaction are not yet fully understood. highly relevant to this work are studies focusing on the interpretability of attention heads @xcite . these studies specifically investigate the role of attention heads in information retrieval processes. the function of self-attention mechanisms extends beyond mere replication of highly relevant information; we emphasize the capability of self-attention mechanisms 0-4k 4k-8k 8k+ input length 25 30 35 40 45 50 55 task score yarn_repobench-p 0-4k 4k-8k 8k+ input length 30 40 50 60 70 task score yarn_lcc no masking top 5% top 10% random 5% random 10% figure 7: masking out top-scored heads v.s. random heads. for the code tasks in longbench, the removal of heads with top scores evidently leads to a decline in performance. to integrate information from different positions. this capability is crucial for practical long-text comprehension tasks. ## conclusion we investigated the properties of attention heads with rotary position embeddings (rope) in commonly used transformer architectures. using long text comprehension tasks as a starting point, we explored the modeling of token-to-token distance within the model by deconstructing the contributions of different dimensions within the attention heads to the attention scores. we found that due to the computational nature of rotary position embeddings, higher dimensions of the attention heads, which correspond to lower rotational frequencies, are more effective at distinguishing distances between tokens. furthermore, attention heads that, through training, allocate attention scores across different dimensions according to token distances and exhibit a certain degree of correlation, demonstrate superior capabilities in modeling text distances. these heads are crucial for integrating information from varying distances in long text comprehension tasks. we provide an analytical perspective on the currently popular rotary position embeddings, illustrating the attention patterns of models trained with rope. future research can leverage the properties of these attention heads to address challenging tasks such as long text comprehension.