Implementation with RoPE about lm-infinite

[Glaciohound]

Hello! Thanks for your interests in our work!

Basically, the difference in rotating angles determines the "attention distance" between two tokens. For example, in original RoPE encoding, every token at position $i$ is just rotated with angle $i\theta$. Then the dot product between query at $i$: $q_i$ and key at position $j$: $k_j$, i.e., $q_i^\top k_j$, depends only on $i-j$. (For why this is the case, a more detailed explanation and derivation can be found in the original paper. Basically the authors carefully constructed the mechanism just to achieve this effect!)

Going back to our work, we want to let "global branch" tokens be exactly $L_{pretrain}$ distance when being attended, regardless of their "true" distance. So a token at position 100 being attended by another token at position 8k will be attended "as if" the distance is $L_{pretrain}$ instead of $8k-100=7900$. To achieve this effect efficiently, in global branch only, we rotate all $k$ vector by angle $0$, rotate all $q$ vectors by angle $L_{pretrain}\theta$, then for the same reason mentioned above, their attention will always be as if their distances are $L_{pretrain}$.

Does this make sense to you? Feel free to follow up if I missed anything!

[sdc17]

Truly helpful explanations, thanks!