Rotary Position Embedding

hardFREEEmbeddings & Positional EncodingLLM 1: Foundations

Implement Rotary Position Embedding (RoPE).

RoPE rotates pairs of embedding dimensions by position-dependent angles. For a vector x of dimension D at position pos:

Split x into pairs: (x[0], x[1]), (x[2], x[3]), ...
For pair i, compute angle θ_i = pos / 10000^(2i/D)
Apply rotation: x'[2i] = x[2i] * cos(θ_i) - x[2i+1] * sin(θ_i) x'[2i+1] = x[2i] * sin(θ_i) + x[2i+1] * cos(θ_i)

Input: rope(pos, D, x)

pos: position index (integer)
D: dimension (even)
x: NumPy array of shape (D,), the vector

Output: the rotated vector as a NumPy array of shape (D,). The tests print each value rounded to 4 decimal places.

Example

Input

1 4
1.0 0.0 1.0 0.0

Output

0.5403 0.8415 1.0000 0.0100

Reasoning

Inputs: pos $= 1$ is the token's position, D $= 4$ is the embedding size, and x $= (1, 0, 1, 0)$ is the vector. That gives two pairs: $(x_0, x_1) = (1, 0)$ and $(x_2, x_3) = (1, 0)$ .
Angles: pair $i$ turns by $\theta_i = \text{pos} / 10000^{2i/D}$ . Pair 0: $\theta_0 = 1 / 10000^{0} = 1$ rad. Pair 1: $\theta_1 = 1 / 10000^{2/4} = 1/100 = 0.01$ rad.
Rotate pair 0 by 1 rad: $x'_0 = 1\cdot\cos 1 - 0\cdot\sin 1 \approx 0.5403$ and $x'_1 = 1\cdot\sin 1 + 0\cdot\cos 1 \approx 0.8415$ .
Rotate pair 1 by 0.01 rad: $x'_2 = \cos 0.01 \approx 0.99995$ , which rounds to $1.0000$ , and $x'_3 = \sin 0.01 \approx 0.0100$ .
Output: $(0.5403,\ 0.8415,\ 1.0000,\ 0.0100)$ . Rotation keeps each pair's length at 1.
Notice the frequency ladder: early pairs spin fast with position and later pairs barely move. At $\text{pos} = 0$ every angle is 0, so the vector comes back unchanged.

Constraints

0 <= pos <= 100, 2 <= D <= 16 (D is even)
Use base 10000 for frequency
Round to 4 decimal places

Learn this in

Tokenization & EmbeddingsLLM Study Plan Attention MechanismsNLP Study Plan

More from LLM 1: Foundations

Transformer Block Forward Passhard Feed-Forward Networkeasy Layer Normalizationeasy KV-Cache Attentionhard Multi-Head Attentionmedium

Concept primer

Background Knowledge

The problem involves implementing Rotary Position Embedding (RoPE), a technique used in natural language processing and other sequence-based models to incorporate positional information into embeddings. Embeddings are dense vector representations of words, tokens, or other elements in a sequence. Positional encoding is crucial because it helps models understand the order and relationships between elements in a sequence.

In the context of RoPE, the goal is to modify the embedding vectors based on their position in the sequence. This is achieved by applying a rotation transformation to pairs of dimensions in the embedding vector. The rotation angle is determined by the position of the vector and the dimension pair being considered. The use of trigonometric functions like $cos$ and $sin$ to compute the new values of the embedding dimensions is a key aspect of the RoPE method.

Understanding how to manipulate vectors using basic linear algebra operations and how to apply trigonometric functions in a programming context is essential for solving this problem. Additionally, familiarity with the concept of embedding dimensions and how they are used in sequence models will provide a solid foundation for tackling the implementation of RoPE.

Algorithm/Approach

The general approach to solving this problem involves:

Parsing the input to extract the position and dimension of the vector
Splitting the input vector into pairs of dimensions
Computing the rotation angle for each pair based on the position and dimension
Applying the rotation transformation to each pair of dimensions
Combining the results to form the rotated vector

This approach involves a combination of basic vector operations, trigonometric calculations, and looping through the dimensions of the input vector.

Step-by-Step Strategy

To implement the solution:

Read the input position and dimension.
Parse the input vector from the space-separated floats.
Loop through the dimensions of the vector, considering them in pairs.
For each pair, calculate the rotation angle based on the position and the current dimension pair.
Apply the rotation transformation to the current pair of dimensions using the calculated angle.
Store the rotated values for each pair.
Once all pairs have been processed, combine the rotated values to form the final rotated vector.
Round the values of the rotated vector to 4 decimal places before outputting.

Common Pitfalls

Things to watch out for include:

Incorrectly calculating the rotation angle due to misunderstanding the formula or misinterpreting the position and dimension values.
Failing to properly pair the dimensions of the vector, leading to incorrect rotation transformations.
Not rounding the final rotated vector values to the specified precision.

Time & Space Complexity

The time complexity of this solution is expected to be $O(D)$ , where $D$ is the dimension of the input vector, since we are looping through each dimension once. The space complexity is also $O(D)$ , as we need to store the rotated vector, which has the same number of dimensions as the input vector.

solution.py

Python 3.13

Test results