Coordinate Attention for Efficient Mobile Network Design

In plain words

Mobile image models often use a small attention module that asks "which feature channels matter for this image?" but never asks "where in the image?". The usual version squeezes each channel into one number for the whole picture, so the location of a feature is thrown away.

This paper's coordinate attention keeps the location. It averages along one axis at a time: first across each row, then across each column. The two resulting strips are mixed by one small shared layer, then split back into a row part and a column part. Each part becomes a gate between zero and one, and the input is scaled by the row gate and the column gate together. The result is a cheap block that can be dropped into the same places as the existing attention modules and that lets the network emphasise a region by its row and its column.

It is aimed at phones and other small models, where the extra cost has to stay close to zero. The authors report gains on image classification and on the detection and segmentation tasks built on those backbones.

TL;DR

Coordinate attention is a light channel-attention block that also encodes position. It pools features separately along height and width eq. 4 eq. 5, mixes them with a shared 1x1 conv and sigmoid gates eq. 6 eq. 7 eq. 8, and multiplies the input by both gates eq. 9. Inserted into MobileNetV2 with reduction ratio 32, it adds about 0.45M parameters and 10M M-Adds over the baseline, and 0.06M parameters over SE, in the authors' ImageNet table tab. 1. Its reported gains are larger on dense tasks such as segmentation §4.4.2. All numbers are author-reported.

Problem

Attention modules that help large networks are often too costly for mobile networks, which have limited size and compute §1. Squeeze-and-Excitation (SE) is cheap and widely used in mobile networks, but it computes channel weights from a global average over the whole map, so it keeps no positional information §1 §3.1. CBAM adds a spatial branch, but that branch reduces channels to one and uses a 7x7 convolution, so it sees only local relations and loses information §1 §4.3. The paper's goal is an attention block that captures channel relations and long-range positional information at a cost comparable to SE §1 §3.2.2.

Method

The block takes any intermediate feature tensor X with shape C x H x W and returns Y of the same shape §3.2. It has two steps: coordinate information embedding and coordinate attention generation §3.2.

1. Direction-aware pooling. For each channel, average over the width to get a height-indexed vector, and average over the height to get a width-indexed vector §3.2.1 eq. 4 eq. 5. Unlike SE's single vector per channel, these keep one spatial axis each.
2. Shared transform. Concatenate the two vectors along the spatial axis and apply one shared 1x1 conv F1 followed by a non-linearity delta §3.2.2 eq. 6. The output has C/r channels, where r is the reduction ratio. Split it back into a height part f^h and a width part f^w §3.2.2.
3. Two gates. Apply two separate 1x1 convs F_h and F_w to f^h and f^w to return C channels each, then apply a sigmoid to get gates g^h and g^w §3.2.2 eq. 7 eq. 8.
4. Apply. Multiply the input by both gates at each position: y(i, j) = x(i, j) * g^h(i) * g^w(j) §3.2.2 eq. 9.

The authors read each element of the two gate maps as whether the object of interest is present in that row or column §3.2.2, Discussion.

Choosing r: the block description gives 32 as an example reduction ratio §3.2.2; Table 1 uses r = 24 for the SE and X/Y variants and r = 32 for coordinate attention tab. 1. Table 4 compares r = 24 and r = 32 against r = 12 and r = 16 tab. 4.

Placement: the paper says Figure 3 shows how the block is plugged into the MobileNetV2 inverted residual block and the MobileNeXt sandglass block §3.3 fig. 3. The text does not give the exact position in words; see Paper vs code.

Pseudocode for the block (features are one sample, no batch axis):

# Coordinate attention block  [§3.2] [eq. 4-9]
# Input:  features X of shape (C, H, W)
# Output: Y of shape (C, H, W), the same shape as X

def coordinate_attention(features, reduction):
    # Average over the width, keeping one value per row  [eq. 4]
    pooled_h = mean_over_width(features)             # (C, H, 1)
    # Average over the height, keeping one value per column  [eq. 5]
    pooled_w = mean_over_height(features)            # (C, 1, W)
    # Stack both strips on one spatial axis for a shared transform  [eq. 6]
    joined = concat_spatial(pooled_h, transpose(pooled_w))  # (C, H+W, 1)
    # Reduce channels by the ratio r, then apply the non-linearity  [eq. 6]
    mixed = delta(conv1x1(joined, out_channels=C // reduction))  # (C/r, H+W, 1)
    # Split the mixed map back into its row part and column part  [§3.2.2]
    mixed_h, mixed_w = split_spatial(mixed, sizes=[H, W])
    # Expand each part back to C channels and squash to (0, 1)  [eq. 7] [eq. 8]
    gate_h = sigmoid(conv1x1_h(mixed_h))             # (C, H, 1)
    gate_w = sigmoid(conv1x1_w(transpose(mixed_w)))  # (C, 1, W)
    # Scale each position by its row gate and its column gate  [eq. 9]
    return features * gate_h * gate_w

Hyperparameters

Training values are from the paper text unless a row says otherwise. The repository contains model code and weights only, not a training script README.md.

Results

All results are author-reported. ImageNet numbers use MobileNetV2 or MobileNeXt as the baseline, trained for 200 epochs §4.1. The EfficientNet-b0 run follows the original EfficientNet paper for other settings §4.3. Detection and segmentation numbers come from the transfer setups in §4.4.

ImageNet, MobileNetV2 baseline tab. 1:

ImageNet, MobileNetV2 at three width multipliers tab. 2:

ImageNet, MobileNeXt at three width multipliers tab. 3:

Reduction ratio, MobileNetV2 tab. 4:

EfficientNet-b0 with SE replaced by CA tab. 5: 76.3% (SE, 5.3M params, 390M M-Adds) vs 76.9% (CA, 5.4M params, 400M M-Adds). The table also lists other methods at comparable size; the best of them is MNasNet-A3 at 76.7%.

Object detection, SSDLite320 on COCO val tab. 6:

Object detection, Pascal VOC 2007 test mAP tab. 7: MobileNetV2 71.7; + SE 71.7; + CBAM 71.7; + CA 73.1.

Semantic segmentation, DeepLabV3 on Pascal VOC 2012 val, mIoU tab. 8:

Semantic segmentation, DeepLabV3 on Cityscapes val, mIoU at output stride 8 tab. 9: MobileNetV2 71.4; + SE 72.2; + CBAM 71.4; + CA 74.0.

Paper vs code

The official repository is https://github.com/Andrew-Qibin/CoordAttention. Its owner is the first author, and its LICENSE is MIT and names Qibin (Andrew) Hou LICENSE. The repository holds the block, a MobileNetV2 model and a pretrained MobileNetV2 + CA checkpoint. It does not hold training, detection or segmentation code: the README points to pytorch-image-models for ImageNet training, to ssdlite-pytorch-mobilenext for detection and to SPNet for segmentation README.md README.md README.md. Those three repositories were not checked here.

• Core equations match. Pooling, the shared transform, two sigmoid gates and the multiplication follow eq. 4 to eq. 9 coordatt.py. The default reduction is 32 coordatt.py, as in the paper's CA rows tab. 1.
• Bottleneck width has a floor. The paper gives C/r §3.2.2. The code uses max(8, C // r) coordatt.py mbv2_ca.py.
• Non-linearity. The paper calls delta a "non-linear activation function" §3.2.2. The code applies BatchNorm and then h-swish coordatt.py coordatt.py. The README calls this activation SiLU and says it works better than ReLU6 README.md. The code's h-swish is a piecewise approximation of SiLU, not SiLU itself.
• Placement in MobileNetV2. The text does not say where the block sits. The code puts it after the depthwise conv's activation and before the pointwise projection mbv2_ca.py. It is not used in the expansion-ratio-1 block mbv2_ca.py.
• Biases. The paper does not state whether the 1x1 convs have biases. The code uses the default bias=True on conv1, conv_h and conv_w coordatt.py coordatt.py.
• Initialization and head. The code initialises convs with normal(0, sqrt(2/n)), BatchNorm weight 1 and bias 0, and linear layers with normal(0, 0.01) mbv2_ca.py. It adds Dropout(0.1) before the classifier mbv2_ca.py. The paper states neither.
• Naming. In the MobileNetV2 file the reduction argument is called groups mbv2_ca.py. It is the reduction ratio, not a convolution group count.
• README table. The README's SSDLite320 table has an extra MobileNetV2 row at 5.0M params and 22.0 AP that the paper does not list README.md.
• Not checked. Parameter counts, M-Adds and accuracy numbers were not recomputed by running the code. The checkpoint was not loaded.

Pitfalls

• Bottleneck floor changes small layers. With max(8, C // r) and r = 32, the bottleneck is 8 for hidden widths below 256. That is 8 where C/r would be 3, 4 or 6 for widths 96, 144 and 192. For MobileNetV2 blocks with hidden widths 384, 576 and 960, the bottleneck is C/32, i.e. 12, 18 and 30 coordatt.py. A reimplementation that uses C/r alone will differ from the released code in these layers.
• Two sigmoid gates shrink the output. The input is multiplied by two sigmoid gates, each in (0, 1), so the block's output is smaller in scale than its input. The paper does not discuss how this changes activation scale. Check the gated output's scale in your own setup. (Inference from the equations.)
• Use the code's activation. The paper leaves delta open. Using ReLU6, or a plain SiLU, is a different block from the one that produced the numbers.
• Batch statistics. The bottleneck has a BatchNorm layer coordatt.py. Train and eval mode behave differently, as in any BatchNorm layer.
• Weight decay "0.9". The VOC detection text says "the weight decay is set to 0.9" §4.4.1. This is almost certainly a typo, and no value is given elsewhere. Do not use 0.9 as weight decay. The optimizer text ("decay and momentum of 0.9") is also ambiguous §4.1.
• Confounded comparison. Table 1 compares SE and X/Y attention at r = 24 against coordinate attention at r = 32 tab. 1. Parameters differ slightly (3.89M vs 3.95M). Table 4 has SE and CBAM at r = 24 and r = 12, and CA at r = 32 and r = 16; CA is never run at the same r as SE or CBAM tab. 4.
• Training recipe is not in the repo. ImageNet numbers depend on a training pipeline that is not in this repository README.md. Expect to rebuild augmentation and schedule from the paper text.
• Latency is one device. The latency column comes from a Google Pixel 4 tab. 1.
• Segmentation and detection code are elsewhere. Those numbers depend on the external repositories named above, which were not checked.

Lineage

• Builds on: the MobileNetV2 inverted residual block §3.3 fig. 3 and the MobileNeXt sandglass block §3.3 fig. 3. Both are named in the paper; this card does not use their content.
• Compared against: Squeeze-and-Excitation §1 tab. 1, CBAM §1 tab. 2, and TA (named in the related work §2.2).
• Superseded by: not stated in this card's sources.