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

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

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).

Results

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.

Links

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.