(torch.Size([1, 3, 2, 16]), torch.Size([1, 3, 2, 16]))

tensor([1.0000e+00, 3.1623e-01, 1.0000e-01, 3.1623e-02, 1.0000e-02, 3.1623e-03,
        1.0000e-03, 3.1623e-04])

torch.Size([3, 8])

tensor([ 0,  2,  4,  6,  8, 10, 12, 14])

tensor([ 0.,  2.,  4.,  6.,  8., 10., 12., 14.])

The [: (dim // 2)] slice ensures we only take the first half of the dimensions.

Looking at your example:

• dim = 16 (head_dim)
• torch.arange(0, dim, 2) gives [0, 2, 4, 6, 8, 10, 12, 14] - that's 8 elements
• dim // 2 = 8
• So [: (dim // 2)] takes the first 8 elements, which is all of them in this case

The reason for this slice is that RoPE works in pairs - each pair of dimensions gets rotated together. So for a 16-dimensional vector, you have 8 pairs: (0,1), (2,3), (4,5), etc.

The torch.arange(0, dim, 2) generates the even indices [0, 2, 4, 6, 8, 10, 12, 14], and the slice ensures you get exactly dim//2 frequency values - one for each rotation pair. Each frequency will be applied to rotate one pair of dimensions.

tensor([0.0000, 0.1250, 0.2500, 0.3750, 0.5000, 0.6250, 0.7500, 0.8750])

The division by dim creates normalized frequency positions that range from 0 to just under 1.

This normalization ensures that:

1. The frequencies are independent of the actual head dimension size
2. Lower dimensions get lower frequencies (slower rotation)
3. Higher dimensions get higher frequencies (faster rotation)

When you raise theta (10000.0) to these normalized powers, you get a geometric progression of frequencies that decay exponentially. The smallest dimensions rotate very slowly (good for capturing long-range dependencies), while larger dimensions rotate faster (good for shorter-range patterns).

Without the /dim normalization, the frequency pattern would change dramatically based on the head size, making the model behavior inconsistent across different architectures.

tensor([1.0000e+00, 3.1623e+00, 1.0000e+01, 3.1623e+01, 1.0000e+02, 3.1623e+02,
        1.0000e+03, 3.1623e+03])

We raise theta to these normalized powers to create a geometric sequence of wavelengths that get exponentially longer.

Taking the reciprocal 1.0 / (theta ** ...) gives us frequencies that decay exponentially:

• First pair: frequency = 1.0 (fastest rotation)
• Second pair: frequency ≈ 0.316 (slower)
• Third pair: frequency = 0.1 (even slower)
• And so on...

This creates a spectrum where:

• Lower dimension pairs rotate quickly (capture short-range patterns)
• Higher dimension pairs rotate very slowly (capture long-range dependencies)

The exponential decay ensures each dimension pair operates at a distinctly different timescale, allowing the model to encode positional information across multiple ranges simultaneously.

tensor([1.0000e+00, 3.1623e-01, 1.0000e-01, 3.1623e-02, 1.0000e-02, 3.1623e-03,
        1.0000e-03, 3.1623e-04])

We take the inverse to convert wavelengths into frequencies.

The theta ** powers gives us exponentially growing wavelengths (how far apart the rotations are). But for RoPE, we need frequencies (how fast the rotations happen).

Since frequency = 1/wavelength, taking the inverse 1.0 / (theta ** powers) gives us:

• High frequencies for early dimension pairs (fast rotation, short-range patterns)
• Low frequencies for later dimension pairs (slow rotation, long-range dependencies)

This creates the desired frequency spectrum where different dimension pairs rotate at different speeds to encode positional information across multiple scales.

You're absolutely right! I mixed that up. Let me correct that:

• High frequencies (like 1.0) rotate quickly → capture short-range positional relationships
• Low frequencies (like 3.1623e-04 = 0.00031623) rotate slowly → capture long-range positional relationships

Think of it like a clock:

• A second hand (high frequency) completes many rotations quickly - good for distinguishing nearby moments
• An hour hand (low frequency) rotates slowly - good for distinguishing longer time periods

In RoPE, the high frequency dimensions help the model distinguish between adjacent tokens, while the low frequency dimensions help distinguish between tokens that are far apart in the sequence.

Thanks for catching that error!

Exactly right! Looking at the code:

torch.arange(0, dim, 2)[: (dim // 2)].float() / dim

This creates values like [0.0000, 0.1250, 0.2500, 0.3750, ...] that increase with dimension index.

When we do theta ** (increasing_values), we get larger denominators for higher dimensions, resulting in smaller frequencies.

So yes:

• Lower dimension indices → higher frequencies → faster rotation
• Higher dimension indices → lower frequencies → slower rotation

This creates a nice hierarchy where early dimensions capture fine-grained positional differences (adjacent tokens) and later dimensions capture coarse-grained positional differences (distant tokens).

tensor([0, 1, 2])

tensor([[0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
         0.0000e+00, 0.0000e+00],
        [1.0000e+00, 3.1623e-01, 1.0000e-01, 3.1623e-02, 1.0000e-02, 3.1623e-03,
         1.0000e-03, 3.1623e-04],
        [2.0000e+00, 6.3246e-01, 2.0000e-01, 6.3246e-02, 2.0000e-02, 6.3246e-03,
         2.0000e-03, 6.3246e-04]])

This plot shows how each dimension pair's rotation angle changes linearly with position.

The key insight is that freqs contains the actual rotation angles that will be applied:

• Position 0: all angles are 0 (no rotation)
• Position 1: each dimension gets its base frequency as the angle
• Position 2: each dimension gets 2× its base frequency as the angle

The linear scaling means:

• High frequency dimensions (dim_pair_0) accumulate rotation angles quickly
• Low frequency dimensions (dim_pair_7) accumulate rotation angles slowly

This creates the relative positional encoding - tokens at different positions will have different rotation angles, and the difference between any two positions depends on both their distance and the dimension pair's frequency.

Exactly! torch.outer(t, freqs) multiplies each frequency by each position index.

Looking at your output:

• Position 0: All frequencies × 0 = 0 (no rotation)
• Position 1: All frequencies × 1 = original frequencies
• Position 2: All frequencies × 2 = doubled frequencies

So yes, the farther the position, the more the frequencies are scaled up. This means:

• Token at position 0 gets no rotation
• Token at position 1 gets rotated by the base frequencies
• Token at position 2 gets rotated by 2× the base frequencies
• Token at position n gets rotated by n× the base frequencies

This creates the relative positional encoding effect - tokens that are further apart will have larger rotation differences between them, making their dot products in attention naturally encode distance information.

tensor([[ 1.0000+0.0000e+00j,  1.0000+0.0000e+00j,  1.0000+0.0000e+00j,
          1.0000+0.0000e+00j,  1.0000+0.0000e+00j,  1.0000+0.0000e+00j,
          1.0000+0.0000e+00j,  1.0000+0.0000e+00j],
        [ 0.5403+8.4147e-01j,  0.9504+3.1098e-01j,  0.9950+9.9833e-02j,
          0.9995+3.1618e-02j,  0.9999+9.9998e-03j,  1.0000+3.1623e-03j,
          1.0000+1.0000e-03j,  1.0000+3.1623e-04j],
        [-0.4161+9.0930e-01j,  0.8066+5.9113e-01j,  0.9801+1.9867e-01j,
          0.9980+6.3203e-02j,  0.9998+1.9999e-02j,  1.0000+6.3245e-03j,
          1.0000+2.0000e-03j,  1.0000+6.3246e-04j]])

tensor([[ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000],
        [ 0.5403,  0.9504,  0.9950,  0.9995,  0.9999,  1.0000,  1.0000,  1.0000],
        [-0.4161,  0.8066,  0.9801,  0.9980,  0.9998,  1.0000,  1.0000,  1.0000]])

Docstring:
polar(abs, angle, *, out=None) -> Tensor

Constructs a complex tensor whose elements are Cartesian coordinates
corresponding to the polar coordinates with absolute value :attr:`abs` and angle
:attr:`angle`.

.. math::
    \text{out} = \text{abs} \cdot \cos(\text{angle}) + \text{abs} \cdot \sin(\text{angle}) \cdot j

.. note::
    `torch.polar` is similar to
    `std::polar <https://en.cppreference.com/w/cpp/numeric/complex/polar>`_
    and does not compute the polar decomposition
    of a complex tensor like Python's `cmath.polar` and SciPy's `linalg.polar` do.
    The behavior of this function is undefined if `abs` is negative or NaN, or if `angle` is
    infinite.


Args:
    abs (Tensor): The absolute value the complex tensor. Must be float or double.
    angle (Tensor): The angle of the complex tensor. Must be same dtype as
        :attr:`abs`.

Keyword args:
    out (Tensor): If the inputs are ``torch.float32``, must be
        ``torch.complex64``. If the inputs are ``torch.float64``, must be
        ``torch.complex128``.

Example::

    >>> import numpy as np
    >>> abs = torch.tensor([1, 2], dtype=torch.float64)
    >>> angle = torch.tensor([np.pi / 2, 5 * np.pi / 4], dtype=torch.float64)
    >>> z = torch.polar(abs, angle)
    >>> z
    tensor([(0.0000+1.0000j), (-1.4142-1.4142j)], dtype=torch.complex128)
Type:      builtin_function_or_method

torch.polar converts the frequencies into complex numbers representing rotations.

It takes:

• abs = magnitude (all 1s, so unit circle)
• angle = the scaled frequencies we calculated

This creates complex numbers of the form cos(angle) + i*sin(angle), which are rotation operators in the complex plane.

So freqs_cis contains the actual rotation amounts for each position and dimension. When we multiply query/key vectors by these complex numbers, we're rotating them by the corresponding angles.

The name "cis" comes from the mathematical notation cis(θ) = cos(θ) + i*sin(θ).

Great question! We apply different frequencies to different head dimensions because we want each dimension pair to encode positional information at a different "scale" or "resolution."

Think of it like this:

• High frequency dimensions (early dimensions) can distinguish between positions 1 vs 2, but positions 100 vs 101 might look similar after many rotations
• Low frequency dimensions (later dimensions) might not distinguish positions 1 vs 2 well, but can clearly distinguish positions 1 vs 100

This creates a hierarchy where:

• Some dimensions are good at encoding fine-grained position differences (nearby tokens)
• Other dimensions are good at encoding coarse-grained position differences (distant tokens)

By having this spectrum of frequencies across dimensions, the model can capture both local and global positional relationships simultaneously. It's like having multiple "clocks" running at different speeds - some for seconds, some for minutes, some for hours.

Without this variety, all dimensions would encode position at the same scale, which would be less expressive for capturing the full range of positional relationships in a sequence.

Yes, exactly! RoPE works on pairs of dimensions. Each pair gets rotated together:

• Dimension 0 & 1 share the same frequency
• Dimension 2 & 3 share the same frequency
• And so on...

So for head_dim=16, we only need 8 frequencies (for 8 pairs). The rotation is applied as:

[x0, x1] → [x0*cos(θ) - x1*sin(θ), x0*sin(θ) + x1*cos(θ)]

This is why we see torch.arange(0, dim, 2) - we're stepping by 2 to get frequencies for each pair, and [: (dim // 2)] ensures we get exactly half the dimensions worth of frequencies.

The complex representation with torch.polar elegantly handles this 2D rotation math for us.

It's called "rotary positional embedding" because it's embedding (encoding) positional information into the query and key vectors.

Traditional positional embeddings add position information to token embeddings before attention. RoPE is different - it embeds positional information directly into the attention mechanism by rotating the query and key vectors based on their positions.

The rotation angles encode where each token is positioned in the sequence. When two tokens at different positions have their vectors rotated by different amounts, their dot product (attention score) naturally reflects their relative distance.

So "embedding" here means "encoding positional information into" rather than the typical meaning of converting discrete tokens to vectors. It's embedding position into the geometry of the vector space through rotations.

Looking at this plot of GPT-2's absolute positional embeddings, I can tell you several key insights:

The plot shows three lines representing the first three principal components of GPT-2's positional embeddings. Research has found that GPT-2's 768-dimensional positional vectors mostly live in a low-rank linear subspace, with a three-dimensional subspace explaining 90% of the variance ^1 ^2 ^3.

When these positional embeddings are plotted in 3D PCA space, they form a helix structure ^4 ^5 ^6. The three lines in your plot are the projections of this helix onto each of the three main PCA dimensions.

Key observations from the plot:

1. Smooth, continuous patterns: The embeddings don't jump randomly but follow smooth trajectories, suggesting the model learned structured positional representations.

2. Different frequencies: Each PCA dimension oscillates at different rates - some change quickly (high frequency) while others change more gradually (low frequency).

3. Helical structure: The dark blue end represents early positions in the sequence, while the red end represents later positions ^7 ^8 ^9, forming a spiral pattern in 3D space.

4. Efficient encoding: This can be thought of as the neural network reserving 3 dimensions for positional information and 765 for semantic content ^10.

This helical structure is quite different from RoPE, which explicitly uses sinusoidal functions. GPT-2 learned this structure through training, discovering an efficient way to encode position that allows for both local and global positional relationships.

Yes, you can create a similar visualization for RoPE! Unlike GPT-2's absolute positional embeddings that are added to tokens, RoPE creates rotational patterns that you can visualize.

Here are a few ways to plot RoPE:

fig, axes = plt.subplots(2, 2, figsize=(15, 10))

# Plot 1: Rotation angles over positions for different dimensions
for i in range(min(8, freqs_cis.shape[1])):
    axes[0,0].plot(torch.angle(freqs_cis[:, i]), label=f'dim_{i}', marker='o')
axes[0,0].set_xlabel('Position')
axes[0,0].set_ylabel('Rotation angle (radians)')
axes[0,0].set_title('RoPE: Rotation angles by position')
axes[0,0].legend()
axes[0,0].grid(True)

# Plot 2: Real and imaginary parts
axes[0,1].plot(freqs_cis[:, 0].real, label='Real (cos)', marker='o')
axes[0,1].plot(freqs_cis[:, 0].imag, label='Imag (sin)', marker='s')
axes[0,1].set_xlabel('Position')
axes[0,1].set_ylabel('Value')
axes[0,1].set_title('RoPE: First dimension pair (cos/sin)')
axes[0,1].legend()
axes[0,1].grid(True)

# Plot 3: 2D rotation visualization
theta = torch.angle(freqs_cis[:, 0])
x_pos = torch.cos(theta)
y_pos = torch.sin(theta)
axes[1,0].plot(x_pos, y_pos, 'o-')
for i, pos in enumerate(range(len(theta))):
    axes[1,0].annotate(f'pos_{i}', (x_pos[i], y_pos[i]))
axes[1,0].set_xlabel('cos(θ)')
axes[1,0].set_ylabel('sin(θ)')
axes[1,0].set_title('RoPE: 2D rotation pattern')
axes[1,0].grid(True)
axes[1,0].set_aspect('equal')

# Plot 4: Heatmap of all rotation angles
angles = torch.angle(freqs_cis)
im = axes[1,1].imshow(angles.T, aspect='auto', cmap='hsv')
axes[1,1].set_xlabel('Position')
axes[1,1].set_ylabel('Dimension pair')
axes[1,1].set_title('RoPE: All rotation angles')
plt.colorbar(im, ax=axes[1,1])

plt.tight_layout()
plt.show()

tensor([[0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
         0.0000e+00, 0.0000e+00],
        [1.0000e+00, 3.1623e-01, 1.0000e-01, 3.1623e-02, 1.0000e-02, 3.1623e-03,
         1.0000e-03, 3.1623e-04],
        [2.0000e+00, 6.3246e-01, 2.0000e-01, 6.3246e-02, 2.0000e-02, 6.3246e-03,
         2.0000e-03, 6.3246e-04]])

GPT-2 has 768 dimensional states and 12 attention heads ^1 ^2 ^3. The head dimension is calculated as:

head_dim = embedding_dim / num_heads = 768 / 12 = 64

So each attention head gets 64 dimensions to work with. This is a design choice that divides the full embedding dimension evenly across all attention heads, allowing each head to focus on different aspects of the relationships between tokens.

torch.Size([1024, 32])