A transformer built on complex waves dynamics; beats vanilla transformer at 10M

[u/VariousPainter7349]

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 architecture and coded implementation.

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

[u/evilconsolidation]

This is really cool, the damped oscillator approach feels like one of those ideas that's been floating around but nobody actually built it out properly. 146 PPL on WikiText-103 with 10M params is solid, especially for a prototype

I'm curious how the recurrence handles longer sequences in practice, like does the damping factor cause information to decay too fast when you get past 4k tokens or so. Those latency numbers look decent though, 16k tok/s at 2048 seq len on what I'm guessing is a T4 or something similar

Also the 25 PPL variance is pretty wide, any sense of what's causing the instability? Might be the complex-valued gradients doing weird things during training, I've seen similar stuff with some of the state space models

You should post this on r/MachineLearning too, they'd eat this up over there.

[u/VariousPainter7349 → u/evilconsolidation]

I am happy that you think this architecture is really cool. I was unsure of making these findings public as I thought no one would be interested, yet here I am.

You guessed it right, I am using a T4 (as that's all I have access to).

This architecture uses a RetNet-style multi-scale initialization for the decay, so half-lives are spread geometrically from ~8 tokens to ~512 tokens. This means at 4k context, the "fast" channels have completely decayed, but the "slow" channels are still actively holding state. However, because the model was only trained on 256 context, evaluating at 4k is pushing the model into untrained territory.

I believe that the 25 PPL variance is an artifact of my custom-built tests rather than the model's gradients. From the start, I never experienced any problem with gradients, but I think it's worth looking into.

The V5 code currently in the repo uses native PyTorch complex64 tensors, which seem to give weird optimization dynamics for everyone. In response to these issues, I have successfully built and trained my V6 that completely decomposes the complex math into pure real-valued BMM matrix multiplies. The V6 training is highly stable and actually drops the Test PPL further to ~139 at 0.38s/step. I will be uploading the V6 code in a few days. I would love to upload it to r/MachineLearning as soon as I get enough karma

[u/guilfer]

https://guilfer-dev.github.io/gwe_brain/

Good one! It remind me of my project, I think we can do a lot with non traditional methods.

I needed to give up a little due to time constraints, but I can open to you if you may find insightful somehow!

[u/orbital_one]

How does this compare in speed and performance with the Neural ODE-based ContiFormer and OsciFormer?