Replacing Linear Attention Projections With Learned Rotations
Grant Sanderson (3Blue1Brown) posted a video about taking the logarithm of an image. Conformal maps in complex space. Squares stay squares. Geometry bends but doesn’t shear. I watched it and thought: what if attention worked like that?
Transformers project embeddings through big learned matrices W_Q and W_K. For d_head=64, that’s 4,096 parameters per projection per head. But conformal maps preserve angular relationships between vectors and are parameterized by a tiny number of values.
If attention heads are just learning geometric transforms, most of those 4,096 parameters are wasted.
That thought pulled me through Clifford algebras, Möbius transforms, quaternion poles, and N-dimensional conformal groups. Most of the theory was wrong. But something useful fell out.
I’m calling it Affine Möbius Attention.
The Möbius decomposition
A Möbius transform is the most general conformal map of N-dimensional space. It decomposes into:
- Translation. Shift the vector. (d_head params)
- Rotation. Rotate in paired orthogonal planes. (d_head/2 angle params)
- Inversion. Map z → z/|z|². (0 params, creates a pole)
- Scaling. Uniform scale. (1 param)
Total: ~1.5 × d_head + 1 parameters. For d_head=64 that’s 97 params versus 4,096 for linear. A 42x reduction.
The inversion is the mathematically interesting part. Theory says the pole is harmless in high dimensions. The theory is wrong. But I’m getting ahead of myself.
GPT-2’s heads aren’t geometric
Before writing training code I pulled the QK circuit matrices from every head in GPT-2 Small. Polar decomposition. Measured conformality.
Mean conformality: 0.44 out of 1.0. Singular value ratios of 100-5000x. Zero heads above 0.8.
GPT-2’s heads amplify some dimensions thousands of times more than others. They use the full expressivity of their linear projections. Not conformal. Not close.
The foundational hypothesis was dead. I built it anyway.
Inversion destroys everything
Standard attention baseline: 2.41 BPB. Full Möbius with inversion: 3.39 BPB. Forty percent worse. Every variant with inversion failed.
Remove the inversion, keep rotation + translation + scale? Loss came right down to baseline.
The thing that makes Möbius transforms mathematically special is the thing that makes them useless for attention. “Measure zero” doesn’t mean “gradient-stable.” I dropped the inversion and never looked back.
Fewer parameters, better generalization
With dropout, Affine Möbius hit 2.250 val BPB versus 2.316 for standard attention. 95.4% fewer Q/K parameters. Better loss.
| Model | Q/K Params | Val BPB |
|---|---|---|
| Standard + dropout | 67,584 | 2.316 |
| Affine Möbius + dropout | 3,136 | 2.250 |
Learned rotation planes? Worse than fixed adjacent pairs. Hybrid (linear early layers + Möbius later)? Worse than pure Möbius everywhere.
The constraint is structural regularization. The model can’t memorize token co-occurrences through arbitrary projections. It has to learn coherent geometric transforms. Those generalize better.
The actual transform:
cos_a, sin_a = cos(angles), sin(angles)
x1, x2 = x[..., :half], x[..., half:]
x_rot = cat([x1 * cos_a - x2 * sin_a,
x1 * sin_a + x2 * cos_a], dim=-1)
Q = exp(log_scale) * x_rot + translation
Per head: d_head/2 angles, d_head translations, 1 scale. For GQA with 8 query heads and 4 KV heads: 1,164 params/layer versus 393,216 for linear Q/K.
H100 scale
Integrated into the Parameter Golf SOTA stack. 512d, GQA, Flash Attention 3, GPTQ, 8xH100, 10-minute window. The parameter savings let me scale from 11 to 13 layers within the 16MB budget.
| Config | Steps | Sliding BPB | Artifact |
|---|---|---|---|
| SOTA (11L, linear Q/K) | ~6,900 | 1.1147 | 15.86MB |
| 13L Affine Möbius | 5,344 | 1.1463 | 15.95MB |
Gap: 0.032 BPB. SOTA wins.
But the Möbius ops are 4.4% of CUDA time. The gap comes from layer count. Two extra layers cost ~1,500 training steps over 10 minutes. Better per-step learning, fewer steps. H100 tensor cores are so fast at matmuls that O(d) vs O(d²) doesn’t matter. The matmuls aren’t the bottleneck.
Quantization surface area
Affine Möbius has far fewer parameters exposed to quantization. The rotation params themselves are tiny (97 per head) and stored as float passthrough. The bulk of quantization pressure falls on V, output, and MLP weights, which are identical between architectures.
In early 1xH100 experiments with undertrained models, the standard baseline showed much worse GPTQ degradation. But this turned out to be an artifact: both models’ Hessians were degenerate from insufficient training, and the baseline simply had more weight matrices hitting the quantization fallback path. On the full 8xH100 runs, both architectures quantize cleanly:
- Affine Möbius 13L: 1.1676 → 1.1719 post-quant (tiny degradation)
- SOTA: comparable post-quant behavior
Whether fewer Q/K parameters translates to a real quantization advantage at scale is an open question. The reduced quantization surface area is structurally real. But I don’t have the evidence to make a strong claim yet.
What this is
After all the Clifford algebra theory, what survived fits in a tweet:
Replace W_Q and W_K with learned rotations in fixed dimension pairs, plus translation and scale.
The affine subgroup of Möbius transforms. No inversion, no poles, no complex numbers. Just angles and offsets.
Closest existing work:
- RoPE : pairwise rotations, but position-dependent fixed angles after linear projections. Ours replaces the projections. Both compose naturally.
- MöbiusAttention (Halacheva et al. 2024): full Möbius in complex space per dimension pair (arXiv:2409.12175). We drop the inversion they kept and work in real space with cross-dimension coupling.
- qGOFT / BOFT : Givens rotations for fine-tuning. They modify existing projections. We replace them.
What if this scaled?
I don’t have the compute to test frontier scale. But the arithmetic is interesting.
A 128-layer model with 96 heads at d_head=128: standard Q/K projections use ~400M parameters. Affine Möbius: ~2.4M. That’s 397M freed parameters you could spend on more layers, wider MLPs, or larger embeddings.
During inference, Q/K projection weights get loaded from memory every token. Affine Möbius weights are ~170x smaller. Less memory traffic, lower latency, higher throughput.
In distributed training, cutting Q/K params by 99% makes attention layers nearly free in the allreduce communication graph. The training bottleneck shifts entirely to MLP and embedding layers.
The catch: my Phase 0 showed GPT-2 heads are highly non-conformal. If that holds at scale, some heads might need more expressivity. The fix could be Affine Möbius plus a small low-rank residual for heads that need it. Even Möbius + rank-4 residual would be dramatically smaller than full linear projections. There’s also an open question about per-head dimension slicing: each head sees a fixed 64-dim slice of the full model dimension. At wider models this constraint might bite.
Someone with a cluster should try this.
The path
3Blue1Brown video → complex holomorphic attention → Clifford algebras → Möbius transforms → Phase 0 reality check (GPT-2 isn’t conformal) → built it anyway → inversion fails → affine subgroup works → beats standard on small data → competitive at H100 scale.
Most of the math was unnecessary. But every wrong turn narrowed the search space. Sometimes the right answer is the simplest piece of a complicated idea.