Scaling Local Self-Attention for Parameter Efficient Visual Backbones

In plain words

Convolutional networks look at small neighbourhoods of pixels and reuse the same filter everywhere. Attention lets every position weigh every other position, which is flexible but very expensive on large images. This paper asks whether attention can work well for image backbones if it is restricted to local windows.

The authors make the local windows practical to compute. Instead of building a separate window for every pixel, they cut the image into blocks, give each block one shared window that covers the block plus a border, and run attention on all blocks in parallel. The windows are not centred on each pixel, so the model is no longer exactly shift-invariant per pixel, but the authors found accuracy improved anyway. They also add a strided attention layer for downsampling, which works like a strided convolution.

Assembled into a ResNet-like family called HaloNet, the models match well-tuned efficient convolutional networks at similar parameter counts. Training and inference speed are the weak point: pure attention models are slower to train than the comparable convolutional ones on the authors' hardware.

TL;DR

• Local 2D self-attention can replace convolutions in an image backbone, with each pixel attending to a window instead of a fixed grid of filter taps §3.1.
• The key speed change is blocked local attention with haloing: one gathered memory per block of queries, with no masking §3.2.
• Strided attention downsampling replaces attention followed by pooling and cuts downsampling FLOPs by 4× §3.2.
• Author-reported: H7 gets 84.9% top-1 on ImageNet from scratch, and the family matches EfficientNet at equal parameters §4.1 fig. 4.
• Relevant if you need a parameter-efficient vision backbone, or want to know how to implement local attention efficiently on matrix hardware.

Problem

• Global attention over a high-resolution image has quadratic cost in the number of positions, which is infeasible for multi-scale backbones §3.2.
• The standard way to implement local attention gathers a k × k window for each pixel, which multiplies memory by about k² §3.2 tab. 1.
• Stand-alone self-attention (SASA, Ramachandran et al. 2019, no card yet) uses blocks plus attention masks to emulate centred windows. The masks waste computation §3.2.
• Self-attention image models had lagged behind state-of-the-art convolutional models in ImageNet classification §1.
• The authors assume the training data and augmentations make strict pixel-level translation equivariance less important. They test this rather than assume it §3.2 §4.3.

Method

Local attention as a spatially varying filter §3.1. The authors write a local pooling function where the weight for each neighbour is a matrix that depends on the relative offset. For convolution it is fixed per offset and shared across positions eq. 1. For dot-product self-attention, the weight is a softmax over the window of the query–key score plus a query–relative-position term, applied to a shared value projection eq. 2 eq. 3. The projections W_Q, W_K and W_V are shared across positions, so the weights vary with content.

Blocked local attention with haloing §3.2 fig. 1. Steps, for one attention layer with query block size b and halo h:

1. Split the H × W input into non-overlapping b × b blocks. The blocks give the queries.
2. For each block, gather a memory of (b + 2h) × (b + 2h) pixels, the block plus a band of h pixels on every side. Pixels outside the image are padded §3.2.
3. Compute keys and values from the memory and queries from the block.
4. Run attention for every block in parallel. There is no attention mask, so every query sees its full memory window.
5. Concatenate the outputs back to H × W.

The receptive field of a pixel is (b + 2h) × (b + 2h), and it is not centred on the pixel §3.2. The authors state that the block-level gather makes a shared neighbourhood cheap in memory and FLOPs tab. 1.

# Blocked local self-attention with haloing, stride 1  [§3.2] [fig. 1]
# x: (B, H, W, c); block size b; halo h; f heads (per-head dim 16 [app. C.1])

def halo_attention(x, b, h, heads):
    # Cut the image into non-overlapping b x b query blocks  [§3.2]
    blocks = block(x, b)                          # (B, H/b, W/b, b*b, c)
    queries = linear_q(blocks)                    # (B, H/b, W/b, b*b, c_qk)

    # Gather one (b+2h) x (b+2h) memory per block; outside the image the
    # padding is zero-filled  [§3.2]
    # [code: raw/src/2103.12731/haloing_code.py#L46-L86]
    memory = halo_gather(x, b, h)                 # (B, H/b, W/b, (b+2h)^2, c)
    keys = linear_k(memory)                       # (B, H/b, W/b, (b+2h)^2, c_qk)
    values = linear_v(memory)                     # (B, H/b, W/b, (b+2h)^2, c_v)

    # Content-geometry logits from relative embeddings, shared across
    # heads; each query-to-memory offset indexes one embedding  [app. A] [eq. 2]
    rel_logits = relative_logits(queries, b, h)   # (B, H/b, W/b, b*b, (b+2h)^2)

    # Content-content logits, then softmax over the memory window with no
    # mask  [§3.2]
    logits = queries @ transpose(keys) + rel_logits
    weights = softmax(logits)                     # over the (b+2h)^2 memory
    out = weights @ values                        # (B, H/b, W/b, b*b, c_v)
    return unblock(out)                           # (B, H, W, c_v)

Strided attention downsampling §3.2 fig. 2. Downsampling is a single attention layer that subsamples the queries with stride 2, like a strided convolution. It uses the same neighbourhood as the stride-1 layer. The authors replace the earlier design of attention followed by strided average pooling, and say this does not change accuracy while cutting downsampling FLOPs by 4× §3.2. How the block size applies after striding is not spelled out in the text (see Pitfalls).

Relative position embeddings app. A. Relative row and column embeddings are learned and factorised. Each takes half of the channels of a head. All heads share them. For window size k the parameter count is 2(k − 1) × 16 per layer, which the authors call small next to the projection matrices app. A.

Global-attention cost comparison §3.2 tab. 1. Table 1 lists neighbourhood memory and FLOPs per pixel for four schemes. The per-pixel and SASA rows are reproduced below. The global row is omitted because its cell formatting is not reliably legible in the source.

The authors also say that global attention on a 128 × 128 map with 64 channels costs about 28× the FLOPs of a 3 × 3 convolution with 64 channels §3.2, footnote.

HaloNet family §3.3 tab. 2 app. C.2. The ResNet-50 structure is kept with the following changes:

• (a) a final 1 × 1 convolution before pooling in larger models, as in EfficientNet;
• (b) bottleneck width factor that varies by model instead of the fixed 4;
• (c) output width multiplier of the spatial operation that varies, instead of 1;
• (d) three blocks in the third stage instead of four §3.3.

Attention heads are fixed per stage at (4, 8, 8, 8) §3.3. The four stages run at strides 4, 8, 16 and 32, and each attention layer uses the stage's block size and halo tab. 2. The attention window is between 14 × 14 (b = 8, h = 3) and 18 × 18 (b = 14, h = 2) §3.3.

The eight configurations H0 to H7 scale the hyperparameters listed in tab. A2. Each is matched in parameter count to the EfficientNet model with the same index app. C.2.

Implementation optimisations app. D. The authors keep the blocked layout (batch, H/b, W/b, b², channels) until the end of the network, so batch norm and other ops work on the blocked format, and they implement haloing as a fixed one-hot 3D convolution rather than slicing and concatenation app. D.

Hyperparameters

Values for ImageNet training from scratch (the H0–H7 family). The paper does not state a separate table for this recipe, so the rows below come from the sections cited.

Model configurations tab. A2 (reproduced from the appendix):

Legend (from tab. A2): b query block size, h halo size, r_v attention output width multiplier, r_b bottleneck output width multiplier, l3 blocks in the third group, d_f final 1 × 1 conv width.

Ablation setup for HaloNet-50 and ResNet-50 §4.2 app. C.1:

Detection and instance segmentation §4.6 app. C.4:

ImageNet-21k pretraining and transfer §4.5 app. E:

Results

Author-reported. The setups are noted in each row. Figures are not reproduced; they are cited by number.

ImageNet from scratch §4.1 fig. 4. H7 achieves 84.9% top-1 on ImageNet validation and 74.7% on ImageNet V2 §4.1. The authors describe the family as matching or slightly beating EfficientNet at equal parameters fig. 4. Training used images up to 600 × 600 §4.1 tab. A2. The paper gives no numbers for the EfficientNet curve in the text.

Component ablation, HaloNet-50 vs ResNet-50 §4.2.1 tab. 3. Both trained for 150 epochs at 256 px §4.2.

HaloNet-50 has about 18.0 M parameters and ResNet-50 about 25.5 M §4.2. The ResNet-50 baseline is 77.6%, which the authors say is higher than the commonly reported 76.3% §4.2.

Image size §4.2.2 fig. 6. Accuracy keeps improving with larger images, and the HaloNet-50 advantage over ResNet-50 is kept. No numbers are stated in the text.

Window size and block size §4.3 fig. 7. Enlarging the window from 6 × 6 to 12 × 12 gives a 1.3% gain §4.3. For a fixed window, block size does not change accuracy, and dropping the halo (non-overlapping blocks) lowers accuracy unless blocks are large. Block 4 with halo 1 beats block 8 with halo 0 §4.3.

Translation equivariance §4.3 fig. 5. Unmasked HaloNet-50 with b = 8, h = 3 beats the masked 7 × 7 window of SASA, and the gap widens with more regularisation §4.3. The paper gives no numbers in the text.

Scaling study app. B tab. A1. Scaling one dimension of HaloNet-50 to about 30 M parameters:

The baseline row is commented out in the source, so the baseline of 80.5 is implied by the deltas tab. A1.

Convolution–attention hybrids §4.4 tab. 4. H7 with convolutions in some stages. Training time is normalised to EfficientNet-B7 tab. 4.

ImageNet-21k transfer §4.5 tab. 6. Pretraining on ImageNet-21k, then finetuning on ImageNet. Parameters are in millions. Step time is TPUv3 compute per batch of 32 per core. Inference is images per second per core, given as 384 px / 512 px tab. 6 caption.

The authors say their models beat BiT and ViT on accuracy and speed in this setting §5. The H4 row here is a different, wider configuration from the H4 in tab. A2.

Detection and instance segmentation (COCO, Mask R-CNN) §4.6 tab. 5. Only the last three convolutional layers of ResNet-50 or ResNet-101 are replaced, with 2 halo layers (b = 8, h = 3), plus one ResNet-50 variant with b = 32 §4.6. AP is box (bb) and mask (mk).

Speed is for the backbone only, batch 32, on 2 TPUv3 cores tab. 5 caption. The authors note the b = 8 variant closes about half of the R50–R101 detection gap and 70% of the segmentation gap §4.6.

Speed and training cost §5 fig. 8. Pure attention HaloNets train more slowly than the matched EfficientNets on the authors' TPUs §5 fig. 8. The hybrids have a speed–accuracy trade-off close to EfficientNet's §5.

Paper vs code

No official code.

What was searched:

• Links in the paper and its source. The arXiv source contains no repository link for HaloNet. The only GitHub link is the TPU detection codebase, which the detection experiments use, not this method §4.6 app. C.4. The footnote in §3.2 says code and checkpoints "will be made available", but no link is given.
• Project page. None is linked in the source.
• GitHub search (repositories). halonet returned 15 repositories, none owned by the authors or by Google. Two unofficial implementations exist: lucidrains/halonet-pytorch (MIT) and qwopqwop200/halonet-tf2 (MIT, describes itself as unofficial). The others are unrelated or unlicensed.
• GitHub search within organisations. halonet+org:google-research and local self-attention haloing org:tensorflow returned zero results. standalone self-attention org:google-research also returned zero.

A TensorFlow snippet ships inside the arXiv source as raw/src/2103.12731/haloing_code.py. It is not a repository and its licence is not stated, so this card does not quote it and does not use it as code. It implements the haloing gather as a one-hot 3D convolution app. D and shows zero-filled padding outside the image haloing_code.py. It does not include the attention computation, so it does not show whether padded positions are masked.

Discrepancies within the paper.

• Learning rate: §4.1 uses 1.6 with batch 4096 and linear warmup. app. C.3 describes cosine annealing "with learning rate 0.1". Both may describe different runs, but the paper does not say which one the H0–H7 models used.
• Weight decay: app. C.3 gives 2e-5 for classification. app. C.1 gives 4e-5 or 8e-5 for the ablations. §4.1 does not state a value.
• Architecture: §4.2 says HaloNet-50 "exactly matches" the ResNet-50 architecture. §3.3 lists four changes that apply to the H0–H7 family. Both can be true, but the reader must know which applies where.

Pitfalls

• Non-centred windows. Output at a pixel depends on which block the pixel falls into, so the layer is equivariant only to shifts that are multiples of the block size §3.2. The authors chose this knowingly and report better accuracy §3.2 fig. 5. Do not assume SASA-style centred masking when you port the code.
• Padding at the image border is not specified. The text says "padding at boundaries" §3.2. The bundled snippet zero-fills out-of-image positions haloing_code.py. The snippet does not include the attention code, so it does not show whether those zeros are masked out of the softmax. If they are not masked, they take part in the softmax as ordinary keys.
• Image sizes that are not divisible by block size. The paper does not say how this is handled. H7 uses 600 px with b = 10 tab. A2. At the stride-8 stage the size is 75, which is not divisible by 10. At stride 16 it is 37.5 and at stride 32 it is 18.75. Padding or cropping is needed, and the paper does not state what it does.
• Strided attention and block size. The text says the strided layer uses the same neighbourhood as stride 1 §3.2. The snippet multiplies the query block by the stride when building the memory haloing_code.py. The paper does not state which block size applies after striding, so check the snippet and your design against each other.
• Block size versus halo cost. For a fixed window, smaller blocks with larger halos use more memory, and the choice of block size does not change accuracy in the paper's test §4.3. Memory and speed favour larger blocks §4.3.
• Heads and channels. Heads per stage are (4, 8, 8, 8) §3.3, and channels per head are 16 app. C.1. Stage widths follow the table, so the channel counts per stage follow from r_v and r_b tab. A2. The paper does not list per-stage channels.
• Hidden training cost. Pure attention models are slower to train than the matched EfficientNets §5 fig. 8. Speed claims are for TPUv3 and do not carry over to other hardware.
• Label smoothing and dropout values. Training from scratch uses label smoothing and dropout §4.1 §4.2, but the paper gives no value for either in that setting. Use "not stated" rather than a default.
• Checkpoints. The footnote promises checkpoints. No release is linked in the source, and the card does not claim one exists.
• Detection uses hybrids. The detection results use attention in only the last three convolutional layers §4.6. They do not show that pure attention helps detection. The authors leave that to future work §4.6.
• Pretraining results are from a different H4. The 21k experiments use a wider H4 with a different base width and block layout app. E, not the H4 in tab. A2.

Lineage

• Builds on: Stand-alone self-attention (SASA; Ramachandran et al., 2019; no card yet). This paper keeps SASA's local-window idea but drops its masks and changes the downsampling layer §1 §3.2.
• Builds on: ResNet bottleneck structure §3.3 and EfficientNet scaling and regularisation recipes §4.2 §4.1. Neither has a card yet.
• Compared against: EfficientNet B0–B7 at matched parameter counts §4.1 fig. 4 tab. A2; ResNet-50 and ResNet-101 §4.2 §4.6; ViT-L/16 and BiT-M on ImageNet-21k transfer §4.5 tab. 6; Mask R-CNN baselines on COCO tab. 5. BoTNet is named as concurrent work that uses global attention in ResNet hybrids §4.6.
• Superseded by: not known from wiki content.