Attention Meets Pixels — Windows, Axes and the First Fully Attentional Networks

The transformer block of Attention Is All You Need — The Transformer, Block by Block reads a sentence of a few dozen words. An image, read pixel by pixel, is a sentence of fifty thousand words. The problem statement of this read is one line: the attention matrix grows with the square of the number of tokens, and an image has too many of them. Four papers from 2018–2020 answered it, and each answer is a different reply to one question: which keys is one query allowed to look at?

Brush up before you start ;)

• How a query scores keys, the softmax, √d and the $O(n^2)$ argument: Query, Key, Value — Dot-Product Attention from Scratch. Here I only substitute $N=HW$.
• The full block, residuals, LayerNorm, sinusoidal positions: Attention Is All You Need — The Transformer, Block by Block.
• The non-local block, global pairwise attention inside a CNN: Attention Before Transformers — Score, Squash, Multiply.

This one is mostly about counting :)

1. Why n² Kills Images

Let's devise the numbers step by step. One attention layer takes $N$ tokens and forms an $N\times N$ matrix of scores. On a feature map of height $H$, width $W$ and $d$ channels, a token is a pixel, so:

• $N$: the number of tokens (the $n$ of Query, Key, Value — Dot-Product Attention from Scratch, capitalised from here on because it is a pixel count); $H,W$: the feature-map size; $d$: the channels every query and key carry.
• The quadratic is in $N$, and $N$ is itself quadratic in the image side. Doubling the side multiplies the attention cost by 16, not 4. [WANG20] write the same thing as $O(h^2{w}^2)$.

The global attention matrix is a seating chart with a column for every pixel. A 224×224 image needs 2.5 billion cells before anyone sits down.

Per layer, counting only the entries of the attention matrix (multiply by $d$ for multiply-adds):

Table. 1 Attention-matrix entries per layer for the three schemes of this read, N = S² pixels. The formulas are Eq. 1, Eq. 3 and Eq. 5; the k = 7 and span = side choices are those of [RAMA19] and [WANG20].

• At 224² pixels, fp32, one head: $2.52\times {10}^9\times 4$ bytes ≈ 10 GB. With 8 heads, ≈ 80 GB. For one layer. That is the whole argument in one number.
• It is also why the non-local block of Attention Before Transformers — Score, Squash, Multiply only ever sat on small, late feature maps: global attention "can only be used after significant spatial downsampling has been applied to the input because it is computationally expensive" ([RAMA19], §2.2).

The three rows of Table. 1 are the three curves below. Slide the window $k$ and the span $m$ and watch the two cheap curves cross; the tiles read the costs off at the 224² mark.

window k (local)

7

span m (axial, capped at S)

256

Everything that follows makes one of the two cheap curves above the one you pay for.

2. Blocks: Attention Inside a Memory Block

[PARM18] (Image Transformer, 2018) wanted to generate images pixel by pixel. A 32×32×3 image is a sequence of 3072 positions, "rendering attending to all positions impractical" (§3.3); only their 8×8 super-resolution encoder, with 192 positions, attends globally. Their fix is the first in the series, and they say where it came from: "Inspired by convolutional neural networks we address this by adopting a notion of locality, restricting the positions in the memory matrix M to a local neighborhood around the query position."

2.1 Query Blocks and Memory Blocks

• Flatten the image in raster order and cut it into non-overlapping query blocks of length $l_q$.
• Give each block one memory block: itself plus the $l_m$ positions before it. Every query in the block shares that memory matrix, so attention stays two matrix multiplies per block, all blocks in parallel.
• 1-D local: blocks are contiguous in the flattened sequence, so in the picture they are a run of rows (Fig. 1, second panel). 2-D local: rectangular blocks of $l_q=w_q·h_q$ pixels with the memory extended up by $h_m$ and sideways by $w_m$.

The time of one layer becomes

• $l_m$: the memory positions one query scores. Compare with Eq. 1: the $HW$ that multiplied itself is now multiplied by a constant you choose. The cost is linear in the number of pixels.
• Their best models: 1-D local, $l_q=256$, a total memory of 512 positions, 12–14 layers (§5.1).

2.2 What the Memory Block Buys

Table 4 of [PARM18], CIFAR-10 bits per dimension (lower is better), 8 layers, same parameter count throughout:

• Memory block 8 → 4.06, 16 → 3.47, 64 → 3.13, 256 → 2.99. The full model reaches 2.90, and 3.77 on ImageNet 32×32.
• A bigger memory block keeps helping, and for free in one sense: "the number of parameters used by local self-attention is independent of the size of the receptive field" (§5.1). A convolution cannot say that.

NOTE: "local" in [PARM18] means local in the flattened sequence, not a square window, and because generation is causal the memory only looks backwards in raster order. That is why the block schemes in the figure of §4 never cover the grid, however many layers you stack.

3. Windows: Stand-Alone Self-Attention

[RAMA19] (2019) turned the question around: not "where can we add attention to a CNN?", but can attention be the primitive, and the convolution the add-on?

3.1 The Layer

Let's build it in two steps. First the content half, Eq. 2 of [RAMA19], for a pixel $x_{ij}\in {ℝ}^{d_{in}}$:

• ${𝒩}_k(i,j)=\{(a,b):|a-i|\le k/2,|b-j|\le k/2\}$: the k×k memory block, the paper's name for the neighbourhood.
• $q_{ij}=W_Q{x}_{ij}$, $k_{ab}=W_K{x}_{ab}$, $v_{ab}=W_V{x}_{ab}$, with $W_Q,W_K,W_V\in {ℝ}^{d_{out}\times d_{in}}$. The softmax runs over the $k^2$ logits of one neighbourhood only.
• Put it next to a convolution (Eq. 1 of the same paper): a conv sums the same neighbourhood with fixed weights $W_{i-a,j-b}$; this layer sums it with weights that depend on the content.
• Cost: $O(hw·k^2·d)$, linear in the pixels; $l_m=k^2$ is Eq. 2 with the memory block made square.

3.2 Relative Position in Two Dimensions

Eq. 3 has no idea where a key sits in the window: swap two neighbours and the output does not change. [RAMA19] fix this with a learned term per offset, their Eq. 3:

• $a-i$ is the row offset and $b-j$ the column offset (their Fig. 4). Each indexes its own learned table of $2\lfloor k/2\rfloor +1$ rows and $d_{out}/2$ columns; the two halves are concatenated to length $d_{out}$.
• Reading the logit: "how similar is the content" plus "how much does this query like that offset".
• Offsets, not absolute positions, so the layer "enjoys translation equivariance, similar to convolutions" (§2.2), which the sinusoidal positions of Attention Is All You Need — The Transformer, Block by Block do not give you.

Table 5 of [RAMA19] is the cleanest ablation in the paper: no position 77.6 %, absolute (sinusoidal) 78.2 %, relative 80.2 % ImageNet top-1. And Table 6: the $q·r$ term alone is only 0.5 % below $q·k+q·r$.

NOTE: "so local attention is a convolution with extra steps" is the wrong reading of Table 6. Even q·r alone is a dot product with the query under a softmax, so the weights still change from pixel to pixel; a conv's do not. And q·r is a vector per offset scored against the query, not the per-offset scalar bias of Swin Transformer — Windows, Shifts and a Pyramid for Dense Prediction, which is a later object.

3.3 Where to Keep Attention?

• They swap every 3×3 convolution in the bottlenecks of ResNet-26/38/50 for this layer, $k=7$, 8 heads; where the conv had stride 2, a 2×2 average pool with stride 2 follows.
• The stem is the exception. Raw RGB pixels "are individually uninformative and heavily spatially correlated" (§3.2), so content matching struggles to learn an edge detector there; they use a separate "spatially-aware value" attention stem, or simply a conv stem.
• Table 3 settles the order: convolutions in stage group 1 (or 1–2) and attention in the rest is best, 80.7 %; attention everywhere is 80.2 %; attention early and convolutions late lands at 79.6–79.9 %, barely above the all-convolution 79.5 %. Remember that the next time someone says "attention everywhere".

3.4 The Numbers

Table. 2 Stand-Alone Self-Attention on ImageNet, Tables 1 and 4 of [RAMA19]. Full Attention is the abstract's "12 % fewer FLOPs and 29 % fewer parameters" than ResNet-50.

• 25.6 M → 18.0 M parameters is −29 %, 8.2 → 7.2 B FLOPs is −12 %, and accuracy goes up 0.7 %.
• COCO with RetinaNet (Table 2; vocabulary in Object Detection — A Quick Read): fully attentional 110 B FLOPs / 22.0 M / 36.6 mAP against 182 B / 33.4 M / 36.5.
• Why $k=7$ is affordable at all: at $d_{in}=d_{out}=128$, a 3×3 convolution costs the same as this layer with $k=19$ (§2.2).

4. Axes: Attention One Axis at a Time

Windows cap the cost, but they also cap the view: a 7×7 window is 49 of 3,136 pixels on a 56×56 map. [HO19] (Axial Attention, 2019) wanted global context and a sub-quadratic cost, without custom kernels. Their starting point is ours, "a small 32×32×3 image has 3072 dimensions. Sequences such as these are too long for self-attention" (§2), and their answer is a reshape.

4.1 The Trick

• Keep the tensor as a tensor and run standard attention along one axis only, "by transposing all axes except k to the batch axis, calling standard attention as a subroutine, then undoing the transpose" (§3.1): every row of the image becomes a separate sentence in the batch.
• A column pass followed by a row pass gives every pixel a path to every other pixel in two layers.
• Position embeddings are "additively factorized" the same way: an $H\times 1\times D$ row table plus a $1\times W\times D$ column table (§3.2.1).

My own picture for it: axial attention is reading a page with a ruler. One pass along the ruler, one pass sliding it down, and the whole page is covered without ever holding all of it at once.

4.2 The Cost

• $S$: the image side; $N$: the pixels; $m$: the span, how far along the axis a query may look ([WANG20] add this knob, next section).
• The saving over global attention is a factor $\sqrt{N}=S$: 56× on a 56×56 map, 224× on 224×224 pixels; "an O(√N) savings", and $O(N^{(d-1)/d})$ for a $d$-dimensional tensor.
• Against the window: local costs $N{k}^2$, two axial passes $2NS$. They tie at $S=k^2/2$, a 24-pixel side for $k=7$; above that axial is dearer and buys the full field in return.

Cost is one side of the trade; reach is the other. Pick a scheme in the figure below and move the query around the 16×16 grid: the left grid is what one layer sees, the right grid two layers, and the last tile counts how many layers it takes to see everything.

scheme (0 global, 1 local, 2 1-D block, 3 2-D block, 4 axial)

4

query position (raster index, row = index ÷ 16)

121

• One axial layer sees $2S-1$ pixels (12 % of a 16×16 grid); the second sees all of them. The 7×7 window needs $\lceil (S-1)/(k-1)\rceil =3$ layers for the same grid, and sees 66 % after two.

NOTE: axial attention is not a separable convolution, even though both factor 2-D into a row pass and a column pass. The axial weights come from q·k per pair and change with the content; a separable conv's are fixed. And one axial layer spans the whole row, not a 1×k kernel.

4.3 The Numbers

Table 2 of [HO19]: ImageNet 32×32, 3.76 bits/dim against Image Transformer's 3.77; ImageNet 64×64, 3.44, level with the Sparse Transformer's 3.44. Global context at the cost of two cheap passes, and no custom CUDA.

5. Axial-DeepLab: Replacing the Convolutions One by One

[WANG20] (2020) is the recognition-side cousin of [HO19] (different papers, same word in the title). Their complaint about the window of §3: a local square "sacrifices the global connection, making attention's receptive field no larger than a depthwise convolution with the same kernel size", and is still quadratic in the region side $m$ (§3.2). Two changes.

5.1 Position-Sensitive Attention

The $q·r$ term of Eq. 4 depends only on the query; the output never records where a value came from. [WANG20] add a key term and a value term (their Eq. 3):

• $o=(i,j)$: the output position; $p=(a,b)$: a key position in the $m\times m$ square around it; $r^q,r^k\in {ℝ}^{d_q}$: relative embeddings for query and key; $r^v\in {ℝ}^{d_{out}}$: one added to the value. All three tables are shared across heads.
• Three position terms instead of one: the query says where it likes to look, the key says where it is, and the value carries its offset out, so the next layer still knows.

5.2 Then Do It Along an Axis

Swap the $m\times m$ square for a $1\times m$ span along the width, and the same again along the height (their Eq. 4):

• Only the neighbourhood changed: ${𝒩}_{1\times m}(o)$ is a 1-D span of $m$ positions along one axis. Cost $O(hwm)$ per pass (the last line of Eq. 5), and $m$ may be the whole axis.

• The span is the whole 56×56 input for conv-stem ImageNet, $m=15$ for the stand-alone version, $m=65$ for segmentation. This is the last step of "replace the convs one by one": a conv stem, axial attention everywhere else.

5.3 The Numbers

Table. 3 Axial-DeepLab on ImageNet, Table 1 of [WANG20]. The two "+" rows are the paper's ablation ladder on top of the conv-stem SASA model. M-Adds count a multiply-add as one; Table. 2's FLOPs count it as two, which is why ResNet-50 is 4.1 B here and 8.2 B there.

• Full Axial-Attention: half the parameters of ResNet-50 (12.5 M vs 25.6 M), fewer multiply-adds, +1.2 % top-1. The position-sensitive terms alone are worth +0.4 % on the ladder.
• COCO panoptic (Table 2): Axial-DeepLab-S, 12.1 M params / 110.4 B M-Adds, 41.8 PQ against Panoptic-DeepLab's 46.7 M / 274.0 B / 39.7.
• The span knob (Table 7, Cityscapes): $m=5\to 65$ takes PQ from 59.1 to 64.2 for +11 % M-Adds (617.4 → 687.4 B). Reach is cheap when it is one-dimensional.

NOTE: the paper's "3.8× fewer parameters and 27× fewer M-Adds" is measured against the multi-scale Panoptic-DeepLab row (3081.4 B M-Adds). Single-scale, the compute ratio is 2.5×, which is still the right direction.

6. The Four Answers Side by Side

Table. 4 The four answers to "which keys?" against the non-local baseline. Costs are attention-matrix entries per layer, N = S² pixels.

Two things to carry forward: every scheme keeps the ResNet skeleton and a convolutional stem (the best numbers of [RAMA19] and [WANG20] both come from that mix), and every scheme buys its saving by shrinking the key set per query, never the number of queries.

7. What Came Next

The first transformer that mattered in vision did not take any of these exits. DETR — Detection as Set Prediction, No Anchors and No NMS (May 2020) keeps the ResNet, runs the unmodified encoder-decoder of Attention Is All You Need — The Transformer, Block by Block on its last feature map, and uses the cost arithmetic of §1 to decide how coarse that map has to be; what it changes is not the attention but the loss. Five months later An Image Is Worth 16×16 Words — ViT, Patch by Patch took the other exit: keep global attention exactly as written and shrink the token count instead, by cutting the image into 16×16 patches so that 50,176 pixels become 196 tokens. No windows, no axes, no relative terms, and no ResNet. The windows of §3 do return later, with a shift and a hierarchy, in Swin Transformer — Windows, Shifts and a Pyramid for Dense Prediction. Stay tuned.

8. References

1. [PARM18] Parmar, Vaswani, Uszkoreit, Kaiser, Shazeer, Ku, Tran. Image Transformer. ICML 2018, arXiv 1802.05751.
2. [RAMA19] Ramachandran, Parmar, Vaswani, Bello, Levskaya, Shlens. Stand-Alone Self-Attention in Vision Models. 2019, arXiv 1906.05909.
3. [HO19] Ho, Kalchbrenner, Weissenborn, Salimans. Axial Attention in Multidimensional Transformers. 2019, arXiv 1912.12180.
4. [WANG20] Wang, Zhu, Green, Adam, Yuille, Chen. Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation. ECCV 2020, arXiv 2003.07853.
5. [VASW17] Vaswani et al. Attention Is All You Need. 2017, arXiv 1706.03762.