Beau

Okay, so last time we really got into the weeds with the scaled dot-product attention. The whole QK transpose operation, using einsum for clean batching, and that causal mask... that triangular matrix with negative infinity to stop the model from cheating and looking into the future.

Jo

Right, the autoregressive property. But that brings up a pretty big problem. The core attention calculation, that dot product, is permutation-invariant. It doesn't care about order. You could shuffle all the past tokens, and the raw dot products would be the same. The mask imposes a limit, but it doesn't tell the model that token 5 comes *right after* token 4.

Beau

So... we have no sense of sequence. It knows what it's allowed to see, but not the arrangement of what it sees. So we need to inject positional information somehow. I've seen the classic approach, just adding a fixed sinusoidal encoding to the input embeddings.

Jo

Exactly. And that works, to an extent. But it's an absolute encoding. It mixes token information and position information additively, right at the beginning. For modern LLMs where we want to generalize to different sequence lengths and really understand relative positioning, we need something more... integrated. That's where Rotary Positional Embeddings, or RoPE, come in.

Beau

RoPE. Okay. So how is it different? If it's not adding, is it multiplying? What's the operation?

Jo

It's a rotation. That's the key. Instead of adding a vector, we are rotating the Query and Key vectors themselves. The angle of rotation depends on the token's position in the sequence.

Beau

A rotation... huh. So it's like we're encoding the position into the vector's direction, not just its magnitude or value in some dimension. So a query vector for the word 'the' at position 5 will be geometrically different from the same query vector for 'the' at position 50.

Jo

Precisely. And the genius of it is how this rotation affects the dot product. Because of the properties of rotation matrices, the dot product between a rotated query at position 'm' and a rotated key at position 'n' only depends on the original vectors and their relative distance, 'm minus n'.

Beau

Wait, that's the goal. You're saying the attention score inherently becomes a function of relative distance. That's incredibly elegant. It's not learning relativity; it's baked into the geometry of the operation.

Jo

That's it. So, how do we implement this? We can't rotate a 512-dimensional vector with a single angle. The trick is to view the embedding dimension in pairs. We take our d-dimensional vector and treat it as d-over-2 pairs of 2D vectors.

Beau

Okay, so you're chunking it. Like, dimensions 0 and 1 are a pair, 2 and 3 are a pair, and so on. And each pair gets its own rotation?

Jo

Exactly. Each pair is treated like a complex number, x_i + jx_{i+1}, and we apply a 2D rotation. The rotation angle is determined by two things: the token's position, 'm', and a frequency, theta_i, which is different for each pair of dimensions. High-frequency rotations happen for the earlier pairs, and low-frequency for the later ones.

Beau

Like a Fourier transform, breaking down the position signal into different frequencies across the embedding dimension. Okay, so in PyTorch... how do we build this? First, we need to pre-compute the frequencies, right?

Jo

Yes. We define a base, say 10,000. The frequencies, theta, are calculated as 1 over base raised to the power of 2i over d, where 'i' is the index of the dimension pair. You can compute this once. So you'd have a tensor of shape (d/2).

Beau

Got it. Then we need the actual positions, a tensor from 0 to sequence length minus 1. And we combine them?

Jo

Right. We use `torch.outer` on the position tensor and the frequency tensor. That gives us a matrix of shape (seq_len, d/2) where each entry is `m * theta_i`. This is our matrix of rotation angles. From there, we compute the cosine and sine of these angles. Now we have a `cos` tensor and a `sin` tensor, both of that shape.

Beau

Okay, so now we have the values we need to perform the rotation. But applying it to the Q and K tensors, which have a shape like (batch, heads, seq_len, d_k), seems tricky. How do we apply this 2D rotation to our paired dimensions without a slow loop?

Jo

This is the clever tensor manipulation part. We can't just multiply. We need to grab the even dimensions (0, 2, 4...) and the odd dimensions (1, 3, 5...) of our Q tensor separately. The rotation formula for a pair (x_i, x_{i+1}) is (x_i*cos - x_{i+1}*sin) for the new x_i, and (x_i*sin + x_{i+1}*cos) for the new x_{i+1}.

Beau

Ah, so we can implement that with slicing and broadcasting. We take our original Q tensor, create a copy of it, and then do some targeted assignments. So, something like... Q rotated even dims equals Q's original even dims times cos minus Q's original odd dims times sin.

Jo

Exactly that. You create a new empty tensor for the output. Then you'd slice `q_out[..., 0::2]` and assign it the result of that first equation. Then slice `q_out[..., 1::2]` for the second. You have to be careful with the shapes to make sure the `cos` and `sin` tensors broadcast correctly, but it's entirely vectorized. No loops.

Beau

And we do the same operation, with the same pre-computed sin and cos values, for the K tensor. But not for V, right? The values shouldn't be position-dependent in this way, they just get aggregated based on the attention scores.

Jo

Correct. RoPE is only applied to Queries and Keys, because its entire purpose is to modify the dot product calculation for the attention scores. Once those scores are computed using the rotated Q and K, the rest of the mechanism proceeds as usual.

Beau

So the performance trade-off... we're adding some computation before the main QK matrix multiplication. But we're gaining a much more robust and flexible sense of position that isn't tied to a maximum sequence length we trained on.

Jo

That's the main benefit. Since it encodes relative position, it extrapolates much better to longer contexts than absolute encodings do. The cost is a few element-wise tensor operations, which on a GPU is marginal compared to the massive matrix multiplications in the attention block itself. It's an efficient way to buy a lot of representational power.

Beau

So, we've replaced a simple vector addition with a more complex, but geometrically-grounded, rotation operation that directly engineers relative position into the attention scores. That's a serious upgrade.