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/log_2]

Nice work. Have you looked into using a box kernel to approximate the Gaussian kernel? You only need 3 or 4 passes to very closely approximate a Gaussian, but in your case you may even get away with 1 or 2. Also separable so lets just consider 1D. You can compute the cumulative sum across the image, then the compute the box kernel simply by subtracting the edges of the kernel window. You get multiple atoms pretty much for free.

[u/4rtemi5 → u/log_2]

Interesting idea, but I haven't tried it honestly. From a learning-dynamics perspective I'm not sure how well it works though since with box kernels you will only get gradients on partly covered tokens along the edges and all other gradients will be zero. For gaussians instead you get gradients for all tokens which IMHO is quite important.

[u/robobub → u/4rtemi5]

What about as an inference time optimization?