Rotary Position Embeddings

Apply rotary position embeddings (RoPE) to a matrix of query or key vectors.

Sinusoidal encodings add a position vector to each token. RoPE instead rotates each query and key by an angle that grows with its position. The dot product of two rotated vectors then depends only on the difference of their positions, which gives attention a sense of relative position for free.

Split each vector of width d into d / 2 pairs of adjacent features, (x[2i], x[2i + 1]). Pair i rotates by angle = pos * theta_i, with

theta_i = base ** (-2 * i / d)
out[2i]     = x[2i] * cos(angle) - x[2i + 1] * sin(angle)
out[2i + 1] = x[2i] * sin(angle) + x[2i + 1] * cos(angle)

The first pair turns fastest, and later pairs turn more and more slowly.

Function signature

def apply_rope(x, positions, base=10000.0):

Example

Input

apply_rope(np.array([[1.0, 0.0, 1.0, 0.0]]), positions=np.array([1]))

Output

array([[0.54030231, 0.84147098, 0.99995   , 0.00999983]])

Reasoning

With d = 4 there are two pairs, with theta_0 = 1 and theta_1 = 10000 ** (-0.5) = 0.01. At position 1, the first pair (1, 0) turns by 1 radian to (cos 1, sin 1), and the second turns by only 0.01 radians.

Constraints and notes

• x has shape (T, d) with d even, and positions has shape (T,).
• A rotation keeps the length of every pair, so |out[t]| = |x[t]|.
• Position 0 leaves a vector unchanged.