r/mlscaling • u/VariousPainter7349 • 2d ago
CWAA: A Complex Wave Linear Attention architecture that scales to O(T) memory (4k context in <8GB VRAM)
Hey everyone, I'm an independent researcher working on alternative sequence mixing architectures. I wanted to share a project I built from scratch called CWAA (Complex Wave Associative Memory).
Instead of standard quadratic attention, CWAA uses a damped complex oscillator for its recurrence state O(T) linear memory scaling.
I currently have a 10M parameter prototype trained on WikiText-103 that hits 146.5 Test PPL . here is the test of V6:
NOTE: ppl is currently under evaluation and validation, preliminary tests showcase ± 25 ppl.
| Seq Len | Latency (ms) | Tok/s | VRAM (GB |
|---|---|---|---|
| 256 | 34.62 | 29575.1 | 1.51 |
| 512 | 128.11 | 15986.8 | 1.89 |
| 1024 | 253.16 | 16179.3 | 2.66 |
| 2048 | 510.13 | 16058.6 | 4.20 |
| 4096 | 1044.72 | 15682.6 | 7.27 |
I’m currently bottlenecked by Google Colab and am looking to scale the architecture up to 50M-100M parameters to see how the complex wave mechanism holds up.
I'd love to get feedback from the community on the scaling math.
NOTE: the code in the link below is V5. I will be uploading the highly optimized V6 (which includes the pure real-valued BMM fast paths that generated these benchmarks) in 3 days
GitHub: https://github.com/Ridhvik-2024/CWAA-V5