3D Attention That Knows Where to Look

The transformer's superpower is one simple idea: let every piece of the input look at every other piece.

A query asks what information it needs, while each key advertises what information it has. Comparing them produces relevance scores. The query then takes a weighted sum of every value using those scores as weights, effectively aggregating information from across the input. Do this for every query and you get attention.

For a sentence of a few hundred words, that's fine. For an image, it's already expensive: you first split the picture into patches, turning each patch into a token. We then ask every token to attend to every other token.

With N tokens of dimension d, attention forms an N × N score matrix (one number for every query–key pair) before mixing the values. That matrix is the cost center, and its size grows with N².

Hold onto that picture: a dense web of connections between every location and every other location. It's about to become the problem.

The catch is in that N × N matrix. Double the side of an image and you don't double the work, you roughly 16× it! Read that again: tokens go from N to roughly 4N and attention squares that number.

Putting real numbers on it:

1. 64 × 64 feature map → 4,096 tokens → ~16.8M attention entries.
2. 128 × 128 feature map → 16,384 tokens → ~268M attention entries.
3. 64 × 64 × 64 volume → 262,144 tokens → ~68.7 billion attention entries.

That last row is the whole story. The moment you go 3D (imagine a LiDAR occupancy grid) the token count is now cubic in resolution, and attention squares it. A modest 64³ volume now needs tens of billions of attention scores per head, per layer. It doesn't fit in memory, and it wouldn't finish in time if it did.

This is the honest answer to "why not just run a transformer over the voxels?" You can't. Dense global attention is already strained in 2D and becomes practically hopeless in 3D.

So we need attention that does not look everywhere.

What if every query could learn exactly where to look? This is the core idea behind deformable attention (Zhu et al., 2020).

Instead of comparing a query against all the keys, it lets every query predict a few locations worth looking at and read only from those. The dense web of connections collapses into a few deliberate arrows.

Two ideas do the heavy lifting:

1. Sampling, not scanning. Each query predicts a small set of offsets (these are directions and distances from a reference point) that say "look a little left, a little higher, and over there." A small fixed number of points (P) per query instead of thousands.
2. No query–key dot products at all. Classic attention computes its weights by comparing queries to keys. Deformable attention predicts the weights directly from the query with a small linear layer, then softmaxes them. There is no giant query–key similarity matrix, surprising but works!

The cost goes from N² to N × P. That's linear in the number of tokens, which is exactly why 3D becomes feasible. The quadratic wall is largely gone!

Usually a sampling point lands at a continuous location, say (7.3, 12.8). So how do you read a value from this?

The same way image resizing works: bilinear interpolation. Take the four pixels surrounding the point and blend them by how close each one is. The result is smooth and crucially, differentiable. This is what enables the offsets to learn.

The full recipe for one query, one head looks like:

1. Start at a reference point aka the query's normalized (x, y) in the feature map.
2. Add the predicted offsets, these make the set of sampling locations.
3. Bilinearly interpolate the value map at each location.
4. Weight and sum them by the predicted attention weights.

With this the gradients flow back into both the offsets (where to look) and the weights (how much to trust them).

One oversimplification we made is that in the real world objects are at wildly different sizes: a pedestrian, a car and a building can appear in the same frame. A single-resolution feature map can't serve them all well.

So usually vision backbones produce a feature pyramid: high resolution features for local details (like edges) and low-resolution features for broader semantic context.

The "Multi-Scale" in MSDA means each query samples a few points from every level of the pyramid at once. Fine levels supply local detail while coarse levels supply broader semantic context and the query blends them in a single weighted sum. With L levels and P points, a query reads just L × P samples total. Across all N queries, this is still relatively small at N × L × P.

Everything above was the 2D story. In 3D, the scaling benefits of deformable attention become impossible to ignore. Volumetric data explodes in size: occupancy grids for robotics, CT and MRI volumes in medical imaging, or video treated as a space-time block.

And the idea extends cleanly:

What changes from 2D to 3D, item by item:

1. Offsets — (x, y) becomes (x, y, z).
2. Interpolation — bilinear (4 neighbors) becomes trilinear (8 corners).
3. Spatial shapes — (H, W) becomes (D, H, W).
4. Sampling locations — (…, P, 2) becomes (…, P, 3).

The math does not change much. A sampling point now sits inside a cube of eight voxels instead of a square of four pixels. You just extend from interpolating 4 neighbors to interpolating 8 corner values.

Deformable attention has been around since 2020, but the fast kernels never left 2D. We needed it for 3D LiDAR scans, so we built it: a from-scratch Triton implementation of Multi-Scale Deformable Attention for both 2D and 3D.

If your data lives in a volume, this is the missing piece. Install, benchmarks, and the full API are on GitHub.