import numpy as np
import matplotlib.pyplot as plt
plt.style.use("../figure_style.mplstyle")
rng = np.random.default_rng(22)
positions, d = 24, 8
base = 10000.0
theta = base ** (-2 * np.arange(d // 2) / d)Demo: When Do RoPE Scores Form Diagonal Stripes?
RoPE rotates queries and keys by absolute position while making their dot products depend on relative offset. We will compare repeated content with content that changes from position to position.
Each pair of coordinates rotates by position times its frequency. The loop below writes the two-coordinate rotation explicitly.
Rotating Each Coordinate Pair
def rope(vectors):
rotated = np.empty_like(vectors)
for position in range(len(vectors)):
for pair in range(d // 2):
angle = position * theta[pair]
cosine = np.cos(angle)
sine = np.sin(angle)
x, y = vectors[position, 2 * pair:2 * pair + 2]
rotated[position, 2 * pair] = cosine * x - sine * y
rotated[position, 2 * pair + 1] = sine * x + cosine * y
return rotatedUse one query and key at every position for the first sequence. Draw fresh content at each position for the second.
q = rng.standard_normal(d)
k = rng.standard_normal(d)
Q_repeated = np.tile(q, (positions, 1))
K_repeated = np.tile(k, (positions, 1))
Q_varied = rng.standard_normal((positions, d))
K_varied = rng.standard_normal((positions, d))Question: When Are RoPE Scores Constant along Diagonals?
With repeated content, every pair at the same relative offset starts from the same query and key. With varied content, the offset may match while the vectors do not. Predict which score matrix will show diagonal stripes.
scores_repeated = rope(Q_repeated) @ rope(K_repeated).T / np.sqrt(d)
scores_varied = rope(Q_varied) @ rope(K_varied).T / np.sqrt(d)
vmax = max(np.abs(scores_repeated).max(), np.abs(scores_varied).max())
fig, axes = plt.subplots(1, 2, figsize=(8, 3.5), layout="constrained")
for ax, scores, label in zip(
axes,
[scores_repeated, scores_varied],
["Repeated content", "Position-dependent content"],
):
im = ax.imshow(scores, cmap="RdBu", vmin=-vmax, vmax=vmax)
ax.set_xlabel("Key Position")
ax.set_title(label)
ax.set_ylabel("Query Position")
fig.colorbar(im, ax=axes, fraction=0.045, pad=0.03, label="Attention Score")
plt.show()
A diagonal contains pairs with the same relative offset. Measure how much the scores vary within each diagonal.
for name, scores in [("Repeated content", scores_repeated), ("Varied content", scores_varied)]:
largest_deviation = 0.0
for offset in range(-positions + 1, positions):
diagonal = scores.diagonal(offset)
deviation = np.max(np.abs(diagonal - diagonal.mean()))
largest_deviation = max(largest_deviation, deviation)
print(f"{name}: largest diagonal deviation = {largest_deviation:.2e}")Repeated content: largest diagonal deviation = 6.66e-16
Varied content: largest diagonal deviation = 3.36e+00
Answer. Repeated content produces a Toeplitz score matrix up to floating-point error: each diagonal corresponds to one relative offset. Varying content breaks the visible stripes because every score depends on both content and offset.
RoPE makes the positional part of the comparison depend on the relative offset. With varied content, scores at the same offset can still differ.