I Derived a New Attention Formula (Cosine-RoPE Attention) and Tested It. Here's What Happened

Yesterday morning I was staring at the attention formula -

Something bothered me about it, not in a "this is broken" way, but more like, why is magnitude here at all?

Attention is supposed to answer one question: which tokens should look at which other tokens? that's a question about relevance, relationship, not about how loud a token is. But we know dot product is magnitude-sensitive. A high-norm token can easily dominate attention just by being large. And $\sqrt{d_k}$ is fixed forever. In result, the model has literally zero control over it. It's like a number is chosen by someone and never gets changed again. That felt counterintuitive to me.

I started thinking about what we actually want from an attention score. We want: given token at position m and token at position n, how relevant are they to each other?

That means, relevance should be about direction. About what the token means, where it points in embedding space. not how large its vector is.

So I thought about cosine similarity, which is magnitude-free, purely angular and bounded in $[-1,1]$. And RoPE (Rotary Position Embedding) which already rotates query and key vectors to encode relative postion. It is orthogonal transformation and it preserves norms, makes position implicit in the geometry. Wait. What if you put them together? RoPE rotates. cosine ignores magnitude. What if you normalize first, then rotate, and compare the resulting unit vectors?

Initially I tried adding Jaccard similarity into it. But I killed it quickly, because gradients die through discrete set operations. So it was not worth it.

What was left after some manipulation and removing everything that didn't fit:

where $\hat{q} = q/\|q\|$ and $\hat{k} = k/\|k\|$ are unit vectors.

Only two things determine this score: what the tokens mean (directio) and how far apart they are (relative position $n-m$). There is no longer magnitude and absolute position. Here $\tau$ is learned per head. Each attention head can find its own sharpness. While standard attention can't do this, because its $\sqrt{d_k}$ is fixed at initialization and never moves.

I called it Cosine-RoPE Attention.

then i thought about softmax

Here's the thing I almost missed that cosine lives in $[-1,1]$ which is always small. And in the other side, softmax on small numbers becomes flat. And the model can't focus sharply focus on anything. But wait, standard attention also goes through softmax and softmax only cares about relative differences between scores. Multiplying all scores by a constant changes nothing...
So what does magnitude in $QK^T$ actually do? It gives non-uniform magnitude. Some tokens grow large norms, others stay small. That creates score differences that softmax can act on. But it also creates instability by sinking attention, explosion of magnitude and dying heads.

$\tau$ solves this cleanly. Instead of uncontrolled magnitude creating implicit sharpness, $\tau$ creates explicit sharpness. And that is learned, controllable.

small $\tau$ → scores spread wider → sharp attention.
large $\tau$ → scores compressed → diffuse attention.

This is not initialized, but model decides it.

the experiment

I didn't want to just theorize. I built two transformers from scratch in PyTorch. Both were identical in every way except the attention mechanism.

architecture - 3 layers, 4 heads, d_model=128, ~13M parameters each.

data - MiniPile. a 22-domain mix of Wikipedia, books, StackExchange, GitHub, web text, arXiv. They were not just grammar but real diverse text with reasoning, facts, code, everything. 59,617 chunks of 128 tokens. 61MB total.

training - AdamW, lr=3e-4, 2000 steps, same random seed for both. fair comparison.

results:

The loss curves nearly sat on top of each other for the 2000 steps. Nearly identical with small difference

what this actually means

My first reaction was disappointment, because the loss was the same, the speed was the same, and it was hard to see the point. But that's the wrong way to read the result.

Identical curves mean the formula is sound. The model trained with no regression, converging at the same rate with the same stability and the same gradient behavior, just with a cleaner geometric foundation underneath. Working in pure angular space costs nothing. That's worth knowing.

And then there's $\tau$.

It initialized at $\sqrt{32} \approx 5.656$, which is the same value standard attention uses as its fixed scale, and after 2000 steps it converged to somewhere between $4.8$ and $5.2$ across all heads. Every head drifted downward in the same direction, which suggests the model wanted slightly sharper attention than the initialization provided. What's interesting is that the heads didn't diverge from each other much and all settled into a narrow band together.

At 2000 steps, the heads haven't yet discovered that they could specialize at different sharpness levels. Standard attention can't express this at all because its scale is permanently fixed at initialization and never moves. My formula can express it, and whether that actually happens at 10k+ steps is the next experiment.

what's still open

The formula works and it matches standard attention at small scale, but equivalence was never really the goal. The goal is finding what it can do that standard attention structurally cannot, and that question isn't answered yet.

Three things stay open. First, does $\tau$ diverge meaningfully across heads at longer training, where different heads settle into genuinely different temperatures, some sharp and some diffuse? That would be expressivity standard attention simply cannot achieve. Second, does pure angular geometry help on tasks requiring fine-grained positional precision, which hasn't been tested at all yet.

The third one is what I keep coming back to. The real problem in attention isn't the scoring formula, it's the $N \times N$ matrix, where every token attends to every other token and most of those scores collapse to near zero after softmax anyway, meaning we computed the majority of them for nothing. If cosine similarity between normalized embeddings can predict which pairs will score high before the full matrix multiply runs, you get sparse attention with a principled geometric reason behind it and the cost drops from $O(N^2)$ to $O(Nk)$, where $k$ is how many tokens each token actually needs to attend to.

The formula came from subtraction, not addition. Magnitude removed, absolute position removed, fixed scaling removed. What's left is clean enough to be worth pursuing, and that's usually how real things start.

Thank you for reading :)