Fast Vision Transformers with HiLo Attention

In plain words

Vision transformers look at an image as a grid of small pieces and let every piece look at every other piece. That is powerful but slow on large images, because the work grows with the square of the number of pieces. Many fast designs cut the work, but some of the cuts run slowly on real GPUs and CPUs even when the arithmetic count is low.

This paper's idea is to treat the attention layer as two jobs. One group of attention heads looks only at small neighbourhoods, which is where fine details live. The other group looks at the whole image, but only at a shrunken, averaged version of it, which is where the broad layout lives. The two outputs are glued back together.

The authors build a complete image backbone, LITv2, around this layer. They also drop the position-tagging scheme used by their earlier model and let a small convolution inside each feed-forward block carry position information. The authors choose their designs by how fast they run on real hardware, not by operation count alone. The claim is that the result is faster than similar models while staying competitive in accuracy.

TL;DR

HiLo is a drop-in replacement for multi-head self-attention. A fraction alpha of the heads does global attention with keys and values average-pooled over s x s windows, and the rest do self-attention inside non-overlapping s x s windows §4.1 eq. 6. LITv2 uses it in stages 3 and 4, with depthwise-convolution feed-forward blocks in stages 1 and 2 §4.3 app. A. Use this card to implement the layer or the backbone. The headline ImageNet, COCO and ADE20K numbers are author-reported §5.

Problem

Standard multi-head self-attention costs about 2N^2 D operations for the attention part, where N is the number of tokens and D the width §3. At high resolution this is the dominant cost, and the cost in memory access is higher than the operation count suggests §1. Efficient alternatives such as shifted windows, spatial reduction attention and multi-scale windows lower the operation count, but some of them need operations that do not run well on GPUs or CPUs, such as window shifting, multi-scale windows, recursion or dilated windows §1 §4.1 §5.4. The paper therefore uses measured throughput on GPU and CPU, not only FLOPs, as its design target §1.

The fixed relative position encoding of LITv1 also slows down dense prediction, because it has to be interpolated for each input resolution §1 §4.2.

Method

Overview §4.1 fig. 2. The authors split attention by spatial frequency. High frequencies capture local detail and are handled by windowed self-attention. Low frequencies capture global structure and are handled by attention to a pooled grid. Each layer splits its heads between the two branches. Both branches are cheap, so the layer is cheap §4.1.

# HiLo attention  [§4.1] [eq. 6]
# Input:  tokens (B, N, D) on an H x W grid, N = H * W
# Output: tokens (B, N, D)

def hilo_attention(tokens, H, W, s, alpha, num_heads):
    # Lo-Fi takes a fraction alpha of the heads; Hi-Fi takes the rest  [§4.1]
    lo_heads = floor(alpha * num_heads)
    hi_heads = num_heads - lo_heads
    # High frequencies: self-attention inside each s x s window  [§4.1]
    hi_out = hi_fi(tokens, H, W, s, hi_heads)        # (B, N, hi_heads*d_h)
    # Low frequencies: global attention to window-pooled keys and values  [§4.1]
    lo_out = lo_fi(tokens, H, W, s, lo_heads)        # (B, N, lo_heads*d_h)
    # Concatenate the two groups along the channel axis  [eq. 6]
    return concat([hi_out, lo_out], axis=-1)         # (B, N, D)

High-frequency branch (Hi-Fi) §4.1. Partition the grid into non-overlapping s x s windows, without shifting. Run standard attention inside each window. The paper uses s = 2 for ImageNet §5.1.

# Hi-Fi: window self-attention  [§4.1]
# Assumes H and W are divisible by s in the classification code path
# [code: classification/models/attentions.py#L193-L210]

def hi_fi(tokens, H, W, s, heads):
    # Cut the grid into non-overlapping s x s windows  [§4.1]
    windows = partition_windows(tokens, H, W, s)     # (B * N / s^2, s^2, D)
    # One linear map gives queries, keys and values for these heads  [§4.1]
    queries, keys, values = project_qkv(windows, heads)
    # Softmax attention inside each window  [§4.1]
    weights = softmax(queries @ transpose(keys) / sqrt(head_dim))
    out = weights @ values
    # Put the windows back in grid order, then project  [§4.1]
    out = merge_windows(out, H, W, s)                # (B, N, heads*d_h)
    return linear_out(out)                           # (B, N, heads*d_h)

Low-frequency branch (Lo-Fi) §4.1. Average-pool the tokens over each s x s window. The queries still come from every full-resolution token. The keys and values come from the pooled grid, which has N / s^2 entries. Standard attention follows.

# Lo-Fi: global attention to pooled keys and values  [§4.1]
# [code: classification/models/attentions.py#L212-L230]

def lo_fi(tokens, H, W, s, heads):
    # Low-pass the grid by averaging each s x s window  [§4.1]
    pooled = average_pool(tokens, H, W, s)           # (B, N / s^2, D)
    # Queries come from the full-resolution tokens  [§4.1]
    queries = project_q(tokens, heads)               # (B, N, heads, d_h)
    # Keys and values come from the pooled tokens  [§4.1]
    keys, values = project_kv(pooled, heads)         # (B, N/s^2, heads, d_h)
    weights = softmax(queries @ transpose(keys) / sqrt(head_dim))
    out = weights @ values                           # (B, N, heads, d_h)
    return linear_out(flatten(out))                  # (B, N, heads*d_h)

Head split §4.1. The paper writes the split as (1 - alpha) N_h heads for Hi-Fi and alpha N_h for Lo-Fi, so the total stays N_h. Splitting the heads rather than doubling them keeps the cost low. The paper also notes that the output projection can be split into two smaller matrices, which saves parameters §4.1. The code uses floor for the Lo-Fi count (see Paper vs code).

Complexity §4.1 app. B. With alpha = 0.5 and a square grid, the paper gives the following costs. Hi-Fi costs (7/4) N D^2 + s^2 N D. Lo-Fi costs (3/4 + 1/s^2) N D^2 + (1/s^2) N^2 D. The second term of Lo-Fi is still quadratic in N, so Lo-Fi becomes expensive at high resolution unless s grows §4.1. The paper therefore recommends a larger window for high-resolution inputs, and evaluates window-4 variants for detection §5.2.

Block structure §4.3 eq. 4 eq. 5. LITv2 has four stages, as in LITv1. Stages 1 and 2 have no attention at all, only a feed-forward block. Stages 3 and 4 use HiLo in place of the multi-head self-attention. Each block uses pre-norm residual connections, the same as the equations for LITv1 §3. Stage 4 uses window size 1, which makes HiLo a standard MSA app. A.

# One LITv2 block  [§4.3] [eq. 4] [eq. 5]
# Stages 1-2: no attention branch.  Stages 3-4: HiLo attention branch.
# [code: classification/models/blocks.py#L73-L124]

def block(x, H, W, has_attention, drop_path_prob):
    # Attention branch, skipped in stages 1-2  [eq. 4] [§4.3]
    if has_attention:
        normed = layer_norm(x)
        x = x + drop_path(hilo_attention(normed, H, W, s, alpha, heads),
                          drop_path_prob)
    # Feed-forward branch with a 3x3 depthwise conv  [eq. 5] [§4.2]
    normed = layer_norm(x)
    x = x + drop_path(conv_ffn(normed, H, W), drop_path_prob)
    return x

Feed-forward with depthwise convolution §4.2 §4.3. The paper adds one 3 x 3 depthwise convolution with zero padding to each feed-forward block, and drops the relative position bias from all attention layers. The convolution supplies the position information through its zero padding, and it also gives the early blocks a larger receptive field §4.2 §5.4. The code places the convolution between the two linear layers, before the GELU [code: classification/models/blocks.py#L46-L53].

# ConvFFN: feed-forward with a depthwise 3x3 conv  [§4.2]
# [code: classification/models/blocks.py#L18-L53]

def conv_ffn(tokens, H, W):
    # Expand the width by the MLP ratio, which is 4  [app. A]
    hidden = linear_in(tokens)                       # (B, N, 4D)
    # Mix neighbours with a 3x3 depthwise conv, stride 1, padding 1  [§5.1]
    hidden = depthwise_conv3x3(hidden, H, W)         # (B, N, 4D)
    hidden = gelu(hidden)
    return linear_out(hidden)                        # (B, N, D)

Stage layout app. A tab. I. The depths and widths below are the values in the paper's architecture table.

Stage 1 patch embedding has patch size 4. Each later stage is entered by a deformable token-merging step with patch size 2 app. A [code: classification/models/patch_embed.py#L80-L115].

Stochastic depth. The drop-path rate rises linearly from 0 at the first block to the configured maximum at the last block, across all blocks of the network litv2.py.

Hyperparameters

Classification unless marked otherwise. "Code" means the official repository at code_commit.

Results

Author-reported. Each row names the dataset, the model and the baseline. Throughput on RTX 3090 uses batch 64 and is averaged over 30 runs §5.1 §5.4. Detection FPS is measured with mmdetection on one RTX 3090 §5.2.

ImageNet-1K classification, 224 x 224 tab. 1.

The paper reports 13%, 27% and 36% faster inference than LITv1 for S, M and B §5.1. Focal-Tiny has higher top-1 (82.2) but runs at 384 img/s tab. 1.

COCO object detection and instance segmentation tab. 2. Single 1x schedule.

ADE20K semantic segmentation, Semantic FPN tab. 3.

Attention mechanisms on LITv2-S tab. 4. HiLo is the only row that wins on speed, FLOPs and ImageNet top-1 among the efficient rows. Plain MSA has the best top-1 but the lowest throughput among these rows.

Architecture modifications on LITv1-S tab. 5.

Other reported results. Alpha 0 performs badly; alpha 0.9 is best, with other non-zero values within about 0.2 top-1 §5.4 fig. 4. Pure Lo-Fi (alpha 1.0) is competitive on ImageNet, but on ADE20K alpha 0.9 gains 0.6 points over it (44.3 against 43.7) §5.4. On ImageNet throughput across GPUs, LITv2-S runs at 1,874 img/s on A100, 1,304 on V100, 928 on RTX 6000 and 1,471 on RTX 3090 tab. 6. Against CPUs and single-layer benchmarks, HiLo is 1.4x faster than SRA and 1.6x faster than local window attention on CPU §5.4 fig. 6.

Paper vs code

Compared with code_commit 7501c86 at https://github.com/ziplab/LITv2. The owner name matches the lab in the copyright line of the top-level hilo.py, "Zhuang Intelligent Processing Lab", and the paper's corresponding author is Bohan Zhuang, who is the corresponding author on the paper's title page. The file names Zizheng Pan as the writer hilo.py. Licence: Apache-2.0 LICENSE. The detection/ and segmentation/ folders carry OpenMMLab and MMSegmentation Apache-2.0 notices [code: detection/LICENSE] LICENSE.

• Classification base learning rate. Paper: 1e-3 §5.1. Code: default 5e-4 config.py. No LITv2 YAML overrides it litv2-small.yaml. main.py applies no linear learning-rate scaling that would change this value, as far as the grep of config.py and main.py shows.
• Warmup, clipping, label smoothing, mixup and augmentation. The paper says these are "the same as in LITv1" §5.1, and it gives no values. The code sets them as listed in the Hyperparameters table [code: classification/config.py#L53-L138].
• Stochastic depth, classification. Paper: not stated. Code: 0.2, 0.3 and 0.5 for S, M and B litv2-small.yaml litv2-medium.yaml litv2-base.yaml. Segmentation values match the paper §5.3.
• Lo-Fi head count. Paper: alpha N_h heads, which is not an integer for alpha 0.9 and N_h 12 §4.1. Code: floor, so 12 heads give 10 Lo-Fi and 2 Hi-Fi heads attentions.py. The split of the 16-head base stage is 14 and 2.
• Stage 4. Paper: alpha 1.0 app. A. Code: alpha 0.9 is passed in, but window size 1 overrides it and assigns all heads to Lo-Fi [code: classification/models/attentions.py#L171-L176]. Behaviour matches the paper.
• Depthwise conv position. Paper: one 3 x 3 depthwise conv in each FFN §4.2. Code: between fc1 and the activation [code: classification/models/blocks.py#L46-L49]. The paper does not state the order.
• Downsampling module. Paper: calls it the deformable token merging module of LITv1 and gives no normalisation app. A. Code: a 2 x 2, stride-2 deformable convolution, then BatchNorm and GELU [code: classification/models/patch_embed.py#L100-L114]. The deformable op needs the bundled mm_modules/DCN extension [code: classification/mm_modules/DCN/setup.py].
• Padding. Paper: pads feature maps that do not divide by s, as in Swin, for detection and segmentation app. C. Code: the classification HiLo requires divisible sizes and does not pad [code: classification/models/attentions.py#L193-L199]. The detection backbone pads before attention [code: detection/mmdet/models/backbones/litv2.py#L285-L294].
• Two HiLo copies. The top-level hilo.py is a standalone module whose forward takes H and W as arguments hilo.py. The classifier uses its own copy inside attentions.py, which infers H = W from the token count attentions.py. Use the copy that matches your input.
• Fine-tune to 384. Paper: Table 1 marks results "finetuned at 384", with no recipe tab. 1. Code: 30 epochs, LR 2e-5 and weight decay 1e-8 [code: classification/configs/litv2-base-finetune-384.yaml].
• Reproduced accuracy. The classification README reports 82.044 top-1 for LITv2-S README.md, consistent with the 82.0 in the paper.
• Detection and segmentation. Detection uses the same 1x-style schedule and window choices as the paper [code: detection/configs/litv2/mask_rcnn_litv2_s_fpn_1x_coco.py#L23-L28]. The segmentation FPN head width of 128 is in the base config [code: segmentation/configs/base/models/fpn_r50_lit.py]. Neither is in the paper.
• Not checked. The CPU and GPU throughput benchmark code, the Figure 6 benchmark scripts and the FLOP counts in the paper's tables were not traced in code.

Pitfalls

• Heads must divide the width. The code asserts dim % num_heads == 0 attentions.py. The head dimension is C / H, so stage 3 of LITv2-S has 32-dimensional heads.
• Floor on the Lo-Fi split. Use floor(alpha N_h) to match the code. The paper's fractional formula does not say how to round.
• Divisibility at inference. The classification HiLo needs H and W divisible by s for the Hi-Fi reshape attentions.py. The 224 and 384 inputs satisfy this at stage 3 (14 and 24 tokens per side). Other sizes need the padding from the detection backbone.
• Square-only classification path. The classifier infers H = W from the token count attentions.py and blocks.py. Non-square inputs need the detection-style path with explicit H and W.
• Speed is hardware-specific. The paper's throughput comparison is measured on RTX 3090 and CPUs with batch 64. Speed rankings on other GPUs differ, as the A100, V100 and RTX 6000 rows show tab. 6.
• Accuracy is not the headline. Plain MSA scores 82.3 top-1 against HiLo's 82.0 tab. 4. The gain is speed and FLOPs, not accuracy.
• Window size and resolution. The paper says a window of 4 improves speed at high resolution with a small accuracy loss §5.2 app. C. The starred detection rows show a drop of about 0.3 AP^b for RetinaNet S tab. 2.
• Position information is implicit. No position embedding or relative bias is used litv2.py. Position comes from zero padding in the depthwise conv, as the paper argues §4.2. Dense tasks depend on that choice.
• Stage-level attention. Stages 1 and 2 have no attention, which is the design the paper inherits from LITv1 app. A. A reimplementation that adds attention early will change the cost and the results.
• Pooled-query variant. Computing queries from the pooled grid is cheaper on paper, but it needs interpolation and was slower on GPU in the authors' test (1,084 against 1,471 img/s) app. D tab. III.
• Separate output projections. The Hi-Fi and Lo-Fi branches each have their own output projection attentions.py [code: classification/models/attentions.py#L191]. The paper says these can be split into smaller matrices, which is what the code does.
• Training recipe is missing from the paper. Warmup, clipping, label smoothing, mixup, cutmix and the 384 fine-tune recipe are only in the code (see the Hyperparameters table). Reproducing Table 1 needs them.
• Compile step. Stage transitions use DeformConv2dPack from the bundled mm_modules/DCN package patch_embed.py. That package has CUDA and C++ sources with a setup.py [code: classification/mm_modules/DCN/setup.py], so it may need building on the target machine.

Lineage

• Builds on: LITv1 (cited as [lit] in the paper), which removes early MSA layers and uses relative position bias and deformable token merging §1 §4.2 app. A. LITv2 keeps the widths and depths of LITv1 §4.3.
• Adds from other work: the depthwise-conv position idea from the cited zero-padding study §4.2; average pooling as a low-pass filter §4.1.
• Compared against: PVT (SRA), Swin (W-MSA and shifted windows), Twins (T-MSA), Focal, QuadTree, Performer, Linformer, CSWin, DAT, HorNet and VAN §5.4 tab. 4 tab. 6 fig. 6.
• Superseded by: not known from content on this wiki.