Rotary Position Embedding (RoPE) › Integration with Self-Attention
Language AI HandbookAll chapters Rotary Position Embedding (RoPE) › Integration with Self-AttentionOne section · Sections in this chapter12 / 19 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
With efficient RoPE implementation in hand, let's see how it fits into a complete self-attention layer. The integration is remarkably clean: RoPE slots in between the QKV projections and the attention computation, requiring no architectural changes to the transformer.
Here's the complete flow:
- Project input embeddings to Q, K, V using learned weight matrices
- Apply RoPE to Q and K (but not to V)
- Compute scaled dot-product attention as usual
- Return the attention output
The critical detail is step 2: we rotate queries and keys but leave values untouched. Let's implement this:
class RoPEAttention:
"""Self-attention with Rotary Position Embedding."""
def __init__(self, d_model, d_k, base=10000, seed=None):
"""Initialize RoPE attention layer.
Args:
d_model: Input embedding dimension
d_k: Query/Key/Value dimension (must be even for RoPE)
base: Base for frequency computation
seed: Random seed for weight initialization
"""
if d_k % 2 != 0:
raise ValueError("d_k must be even for RoPE")
if seed is not None:
np.random.seed(seed)
# Projection matrices
scale = np.sqrt(2.0 / (d_model + d_k))
self.W_q = np.random.randn(d_model, d_k) * scale
self.W_k = np.random.randn(d_model, d_k) * scale
self.W_v = np.random.randn(d_model, d_k) * scale
# RoPE frequencies
self.frequencies = compute_rope_frequencies(d_k, base)
self.d_k = d_k
def forward(self, x):
"""Compute RoPE attention.
Args:
x: Input embeddings of shape (seq_len, d_model)
Returns:
output: Attention output of shape (seq_len, d_k)
attention_weights: Weights of shape (seq_len, seq_len)
"""
# Project to Q, K, V
Q = x @ self.W_q
K = x @ self.W_k
V = x @ self.W_v
# Apply RoPE to Q and K
Q_rope = apply_rope_batch(Q, self.frequencies)
K_rope = apply_rope_batch(K, self.frequencies)
# Scaled dot-product attention
scores = Q_rope @ K_rope.T / np.sqrt(self.d_k)
# Softmax
scores_stable = scores - scores.max(axis=1, keepdims=True)
attention_weights = np.exp(scores_stable) / np.exp(scores_stable).sum(
axis=1, keepdims=True
)
# Aggregate values (V is NOT rotated)
output = attention_weights @ V
return output, attention_weightsNote that RoPE is applied only to queries and keys, not to values. This is because:
- Queries and keys determine attention patterns: The dot product between them computes compatibility. RoPE makes this compatibility position-aware.
- Values carry content: They should not be position-encoded because the content itself doesn't depend on position, only how much weight it receives.
The asymmetry matters. Attention works in two stages: first decide how much to attend to each token (queries and keys), then aggregate what content to retrieve (values). Position belongs only to the first stage. A value vector encodes what a token means, and that meaning should not be rotated away just because the token appears at a different position. If you were to rotate values as well, the final attention output would carry a positional artifact mixed into the content representation, complicating what the next layer needs to do. Keeping values position-free separates the two functions: the attention distribution is position-aware, but the content it aggregates is not.
# Test the RoPE attention layer
seq_len = 8
d_model = 16
d_k = 8
x = np.random.randn(seq_len, d_model)
attention = RoPEAttention(d_model, d_k, seed=123)
output, weights = attention.forward(x)RoPE Attention Layer Test: Input shape: (8, 16) Output shape: (8, 8) Weights shape: (8, 8) Row sums (should be 1.0): [1. 1. 1. 1. 1. 1. 1. 1.]
The output confirms correct behavior: input of shape (8, 16) produces output of shape (8, 8) after projection to the query/key/value dimension. The attention weights form an 8×8 matrix where each row sums to exactly 1.0, confirming proper softmax normalization. The RoPE transformations are applied internally to queries and keys, making attention position-aware without changing the external interface.
The integration is truly a drop-in modification. An existing transformer implementation that computes Q, K, V = linear(x) followed by scores = Q @ K.T / sqrt(d_k) needs only one additional step: apply RoPE to Q and K before the matrix multiply. No changes to the attention score formula, no changes to the softmax, no changes to the value aggregation. This is why RoPE was adopted so quickly: it required almost no refactoring of existing transformer codebases.