SSOG: Near linear Visual-Attention that doesn't score but steers

I swapped the transformer's matchmaking service for a handful of Gaussians, and it beat content scoring without comparing a single pair of tokens.

Crack open any vision transformer and you'll find the same machinery humming at its core: scaled dot-product attention (SDPA). Every token asks every other token "how much do I care about you?", and the answer is a big $N \times N$ matrix of similarity scores, computed from the content of the tokens themselves.

We rarely second-guess this. I know, because I tried once before: a while ago I swapped the dot product for an RBF kernel and learned a lot about what similarity even means in these systems. This time I wanted to question something more fundamental:

Why does the network need to compute "where to look" from scratch, for every image, from content?

Think about how you read an image. When you process a patch on a bird's wing, you don't run a similarity search against all other patches. You already know where to look: a bit left, a bit up (is that the head?), further out (where does the wing end?). Where to look is mostly a function of geometry, not content. Content only fine-tunes it.

So I built attention that way. Each head owns a few Gaussians, a small handful of numbers in total, forming a fixed field over relative position. A learned habit of where to look. Then the trick that makes it actually work: a tiny content-conditioned nudge that lets each token shift its field. No query-key dot products anywhere. Content never scores, it only steers.

The results surprised me. The fixed field alone, completely content-blind, comes within one point of SDPA on ImageNet. Steering closes that gap entirely, and the full model beats the baseline. On small data the geometric prior is worth a ridiculous +17 points. It scales too: a 12M version reaches 72% on ImageNet, ahead of its SDPA twin while being 20% smaller and 30% cheaper to run. The kicker: since the field factorizes into two 1D filter passes, the $N \times N$ attention matrix never exists, so this is attention that scales near-linearly instead of quadratically. Best of all, "what did attention learn?" stops being a heatmap-and-a-shrug question. Every head is just a few blobs you can literally plot and read with a ruler.

(a) score every pair from content. (b) a fixed Gaussian field, same for every image. (c) content only steers that field. ★ = query.

Meet the Field

This is the entire attention mechanism. Each head owns a few Gaussian atoms. An atom is five numbers: a center offset (μy, μx), a width in each direction (σy, σx), and a weight λ. How many atoms per head is a free dial (even one gets you most of the way).

★ is the query. μ is where an atom looks, σ how wide it stares, λ how much it counts. Three atoms = fifteen numbers.

That's the head: a handful of numbers instead of an $N \times N$ matrix of content-dependent scores. The weight from "here" to "there" is just this mixture evaluated at the displacement between them:

$$A(p, q) = \mathrm{softmax}_q \Big( \tfrac{1}{\tau}\; s(p, q) \Big)$$

with each atom contributing its log-weighted Gaussian to the score:

$$s(p, q) = \mathrm{logsumexp}_r \big( \log \lambda_r + \log \mathcal{N}(p - q;\ \mu_r, \sigma_r) \big)$$

Same field for every image, every time. A cat photo and a car photo get the exact same attention pattern. I called the family SSOG (Separable Sum of Gaussians), because that's what it is.

But a purely fixed field only gets you so far. Images aren't all the same, and sometimes the bird is just not in the center. So the field needs to move. The fix is called mu_delta, with its siblings sigma_delta and lambda_gate. Each token gets one tiny linear layer (zero-initialized!) that predicts small residuals on the field parameters:

$$\begin{aligned} \mu &\leftarrow \mu_0 + s_\mu \cdot 4 \cdot \tanh(W_\mu x)\\ \sigma &\leftarrow \sigma_0 \cdot e^{s_\sigma \tanh(W_\sigma x)}\\ \lambda &\leftarrow \mathrm{softmax}\big(\log \lambda_0 + s_\lambda \tanh(W_\lambda x)\big) \end{aligned}$$

In words: content may shift where an atom looks, widen or tighten how hard it stares, and re-weight which atoms matter. All bounded, all starting at zero. I call it lookat. Tanh caps travel at $\pm$4 grid cells, and because the gates start at $\approx$ 0 the model begins as a frozen geometric animal and learns whether to open the content taps at all.

(Spoiler: every layer opens them. Plot coming.)

Build a head yourself: drag atoms, widen them, re-weight them, or watch content steer onto the bird. The presets are geometries the trained model actually converged to:

Error: Embedded data could not be displayed.

If this sounds crazy: it kind of is, and it kind of isn't. Synthesizer showed in 2020 that even random learned attention matrices work surprisingly well. I'm just forcing the matrix to be a smooth, translation-invariant geometric object — a spatial prior, like convolutions, but softer and longer-ranged.

The Separable Trick (a.k.a. How to Never Build the $N \times N$ Matrix)

One more piece, and the reason any of this is fast. A 2D Gaussian factors: $\mathcal{N}(\Delta y, \Delta x) = \mathcal{N}(\Delta y) \cdot \mathcal{N}(\Delta x)$. Applying the field is two 1D filter passes:

Row pass × column pass = one atom's 2D field. ★ = query. Dense attention builds an $N \times N$ map ($O(N^{2} d)$); separable SSOG never does ($O(N \sqrt{N}\, d)$ per atom).

Filter rows, then columns, once per atom; mix with λ; done. Three einsums per layer:

y = jnp.einsum("biwprj,bjwpd->biwpdr", ay, v) # down the rows, per atom y = jnp.einsum("biwprk,bikpdr->biwpdr", ax, y) # across the columns y = jnp.einsum("pr,biwpdr->biwpd", lam, y) # mix atoms