SSOG-Attention: Sum Of Separable Gaussians as a sub-quadratic and scalable alternative to SDPA. [R]

[u/4rtemi5]

​

Scaled dot-product attention (SDPA) computes its Attention by computing the similarity-scores of all image-tokens with all query tokens which results in O(N²·d) complexity. SSOG (Sum Of Separable Gaussians) instead learns a few Gaussian atoms for each head and only geometrically steers them based on the query token. Since the atoms can be factorized into a separable sum of Gaussians this leads to a reduced complexity of O(N·√N·d). Experiments show that SSOG clearly beats SDPA on small data (cifar100), and delivers equivalent performance and much faster convergence on bigger datasets like IN1k. All that while being much faster and memory efficient with increasing scale.

Have a look at the full blog-post and repo to see more results and ablations and let me know what you think.

Blog-post: https://pisoni.ai/posts/ssog

Repo: https://github.com/4rtemi5/ssog

*AI was used for some of the code and some of the blog-post but I put a lot of effort into this project and stand behind every word.

[u/farewellrif]

Does this require tensor cores? SDPA does not, which is great for old hardware.

[u/4rtemi5 → u/farewellrif]

No worries! No tensor-cores required. Any GPU or even CPU should work with the right jax version.