SwiftFormer: Efficient Additive Attention for Transformer-based Real-time Mobile Vision Applications

In plain words

Transformer models let every patch of an image look at every other patch. That is powerful, but the work grows with the square of the number of patches, and the big matrix multiplications behind it are slow on a phone. Convolutional networks are fast on phones but see only a local neighbourhood at a time.

SwiftFormer keeps the "every patch can look at the whole image" idea but changes how it is computed. First it condenses all the queries into one summary vector for the image, using learned weights that decide how much each patch contributes. It then multiplies that summary element by element with the keys, which are the per-patch descriptors, and passes the result through one small linear layer. Only element-wise products and linear layers are needed, so the cost grows in a straight line with the number of patches instead of with its square.

The network uses this block in every stage, combined with cheap convolutions that handle local detail. The authors report that the resulting models are more accurate than comparable phone-oriented networks at similar speed on an iPhone, and that they also transfer to detection and segmentation.

TL;DR

• Replaces the query-key matrix product of self-attention with a pooled global query, an element-wise product with keys, and a linear layer. Cost is linear in the number of tokens §3.2 eq. 4-6.
• A hybrid network: a convolutional encoder plus this additive attention block, applied in all four stages §3.3 fig. 3.
• Author-reported: SwiftFormer-S gives 78.5% ImageNet-1K top-1 at 0.8 ms on iPhone 14 (CoreML, Neural Engine) tab. 2.
• Relevant if you need global context in a vision backbone on a phone-class device. The paper reports no positional encoding and no attention bias.

Problem

Self-attention costs O(n²·d) in the number of tokens n, which limits it on high-resolution inputs and on resource-constrained devices §3.1. Hybrid designs keep attention only in late, low-resolution stages §2. Earlier linear variants (transpose attention, separable attention, low-rank factorisation) still contain matrix products or explicit key-value interactions that the authors find slow on mobile hardware §2 §3.1 fig. 2. The paper argues that FLOP and parameter savings do not always become faster phone inference, so it optimises measured latency instead §2 §3.

Method

Attention variants compared §3.1 fig. 2 eq. 1-3. Standard self-attention is softmax(Q K^T / sqrt(d)) V eq. 1. Transpose attention (EdgeNeXt) applies attention across channels instead of tokens eq. 2. Separable attention (MobileViT-v2) computes a per-token context score from Q, a context vector from K weighted by those scores, and multiplies it element-wise with V eq. 3.

Efficient additive attention §3.2 fig. 2 (d). The method drops the key-value interaction entirely and keeps only query-key interaction, followed by a linear layer. Steps:

1. Project tokens to queries and keys with two learned d × d matrices §3.2.
2. Normalise the query matrix; call the result Q-hat eq. 6. The paper does not state the norm.
3. Score each token with a learned vector w_a, scaled by 1/sqrt(d) eq. 4.
4. Pool the queries with these scores into one global query vector q of size d eq. 5.
5. Multiply the keys element-wise by q eq. 6.
6. Apply a linear transform T to that product and add Q-hat; the output is Q-hat + T(K * q) eq. 6.

The paper states that this is linear in the token count and needs no key-value matrix product §3.2. It also states that no positional encoding or attention bias is used app. F.

# Efficient additive attention  [§3.2] [eq. 4-6]
# Input:  x (B, N, d) token features
# Output: (B, N, d)
# Params: W_q, W_k (d, d); w_a (d,); T linear (d, d)

def efficient_additive_attention(x):
    # Query and key projections  [§3.2]
    queries = x @ W_q                                  # (B, N, d)
    keys = x @ W_k                                     # (B, N, d)

    # Normalised query matrix Q-hat; norm not named in the paper  [eq. 6]
    # The code also normalises keys (see "Paper vs code")
    queries = normalize(queries, dim=-1)               # (B, N, d)

    # One score per token from the query  [eq. 4]
    # Normalisation over tokens is not stated (see Pitfalls)
    scores = (queries @ w_a) / sqrt(d)                 # (B, N)
    weights = normalize(scores, dim=1)                 # (B, N)

    # Pool the queries into one global query per image  [eq. 5]
    global_query = sum(weights[..., None] * queries, dim=1)  # (B, d)

    # Element-wise query-key interaction  [eq. 6]
    interaction = keys * global_query[:, None, :]      # (B, N, d)
    # Linear transform T, plus the normalised query  [eq. 6]
    return queries + T(interaction)                    # (B, N, d)

Architecture §3.3 tab. 4 fig. 3. Four stages at strides 4, 8, 16 and 32. A patch embedding (two stride-2 3 × 3 convolutions) produces the stride-4 map. Between stages, a 3 × 3 stride-2 convolution downsamples and changes the width tab. 4. Each stage is a run of convolutional encoder blocks followed by one SwiftFormer encoder block as its last block §3.3 tab. 4. The stage layout for the four variants is below. Depths and widths are also in swiftformer.py.

The stem and downsampling widths in the table come from tab. 4. The per-stage block counts add up to the depths in swiftformer.py. The expansion ratio of 4 in the convolutional encoder is stated in app. A.

Convolutional encoder block §3.3 eq. 7. A 3 × 3 depth-wise convolution followed by BatchNorm, then a 1 × 1 convolution expanding to 4C, GELU, a 1 × 1 convolution back to C, and a skip connection eq. 7 app. A.

# Convolutional encoder block  [§3.3] [eq. 7] [app. A]
# Input/Output: x (B, C, H, W); hidden width 4C
# The code also applies a per-channel layer scale (init 1); see Pitfalls

def conv_encoder(x):
    # 3x3 depth-wise conv, then BatchNorm  [eq. 7]
    y = batch_norm(depthwise_conv3x3(x))               # (B, C, H, W)
    # 1x1 conv to 4C, GELU, 1x1 conv back to C  [eq. 7] [app. A]
    y = conv1x1(gelu(conv1x1(y)))                      # (B, C, H, W)
    # Skip connection  [eq. 7]
    return x + y                                       # (B, C, H, W)

SwiftFormer encoder block §3.3 eq. 8. Local mixing, then efficient additive attention, then a linear block with two 1 × 1 convolutions, BatchNorm and GELU §3.3 eq. 8. The code version is shown here, because it differs from eq. 8 in the local mixing step (see Paper vs code).

# SwiftFormer encoder block  [§3.3] [eq. 8]
# Input/Output: x (B, C, H, W)

def swiftformer_encoder(x):
    # Local representation: dw 3x3 conv, BN, 1x1 conv, GELU, 1x1 conv
    # Residual with layer scale init 1  [code: models/swiftformer.py#L211-L222]
    x = local_representation(x)                        # (B, C, H, W)
    # Efficient additive attention over the H*W tokens  [§3.2] [eq. 6]
    # Residual with layer scale init 1e-5  [code: models/swiftformer.py#L252-L265]
    tokens = flatten_to_tokens(x)                      # (B, H*W, C)
    x = x + layer_scale_1 * unflatten(efficient_additive_attention(tokens))
    # Linear block: BN, 1x1 conv to 4C, GELU, 1x1 conv back  [§3.3] [eq. 8]
    x = x + layer_scale_2 * mlp(x)                     # (B, C, H, W)
    return x

Training recipe for classification §4.1 app. B. Trained from scratch on ImageNet-1K for 300 epochs with AdamW, a cosine schedule and a linear warm-up over 5 epochs. Hard distillation from a RegNetY-16GF teacher. The full values are in the hyperparameter table below.

Hyperparameters

Values are from the paper unless marked as code. "not stated" means the paper does not give the value, so the code's value is shown separately.

Results

Author-reported. Setup: ImageNet-1K classification, top-1 accuracy with 224 × 224 input §4.1 §4.3. Latency is measured on an iPhone 14 Neural Engine (iOS 16), CoreML, batch size 1; throughput on an A100 at batch size 128 §4.1 tab. 2. The paper gives an error of less than ±0.1 ms for latency tab. 2.

Ablation on ImageNet-1K, built from EfficientFormer-L1 tab. 1 §4.2:

The text notes that the additive attention alone costs 0.2 points of accuracy against the baseline, and that placing the block in all stages recovers accuracy at unchanged latency §4.2. Appendix C reports, without a table, that removing the key-value interaction gives about 10% lower latency and 0.4 points higher top-1 app. C.

Comparison on ImageNet-1K tab. 2 §4.3:

Bold rows are SwiftFormer variants. Swin-T latency is reported as NA tab. 2.

Dense prediction tab. 3 §4.4 §4.5: backbone-only comparison, with the backbone features passed to Mask R-CNN (COCO) or semantic FPN (ADE20K).

Headline claims in the text: SwiftFormer-S is 1.6 × faster than EfficientNet-b0 and 1.4 points more accurate §4.3; SwiftFormer-L3 is 2.6 × faster than MobileViT-v2 ×2.0 with 1.8 points more accuracy §4.3. These are the authors' comparisons on their own measurement setup.

Paper vs code

Code: repository Amshaker/SwiftFormer at commit 4aa6cd67527baae00390b58134fa8954e06c5431, Apache-2.0 LICENSE. Owner is the first author. The repo has classification training and evaluation only; the README says so and gives no detection or segmentation code README.md.

• Query and key normalisation. Paper: only the query matrix is normalised, as Q-hat eq. 6. Code: both queries and keys are L2-normalised before use swiftformer.py.
• Score normalisation. Paper: α = Q·w_a/sqrt(d), with no normalisation over tokens stated eq. 4. Code: scales by token_dim^-0.5 swiftformer.py, then L2-normalises the scores over the token axis swiftformer.py. Because the last step is scale-invariant, the 1/sqrt(d) factor has no effect on the output in the code.
• Output layers. Paper: a single linear transform T in x̂ = Q-hat + T(K * q) eq. 6. Code: Proj on K * G, plus the query, then a further linear layer final on the sum swiftformer.py. Both layers have biases.
• Local representation. Paper: first line of the encoder is X̂ = Conv1(DWConv_BN(X̂)), with one point-wise conv and no residual eq. 8. Code: depth-wise conv, BN, 1×1 conv, GELU, 1×1 conv, with layer scale and a residual swiftformer.py.
• Stem. Paper: two stride-2 3×3 convolutions §3.3. Code: each convolution is followed by BN and ReLU swiftformer.py.
• Convolutional encoder. Paper: no layer scale eq. 7. Code: a learnable per-channel layer scale, initialised to 1, multiplies the branch swiftformer.py swiftformer.py. Hidden width 4C matches the paper's expansion ratio swiftformer.py.
• Layer scale, encoder block. Not stated in the paper. Code: initial value 1e-5 for both branches swiftformer.py swiftformer.py.
• Heads. Not stated in the paper. Code: a single head swiftformer.py.
• Architecture depths and widths. Match Table 4 for all four variants swiftformer.py tab. 4.
• Mixup and CutMix. Paper: "percentages of 1 and 0.8, respectively" for Mixup and CutMix app. B. Code defaults: Mixup 0.8 main.py, CutMix 1.0 main.py. The two values appear swapped relative to the paper's order.
• Augmentation per variant. Paper: one augmentation recipe for all models app. B. Released scripts: dist_train.sh disables RandAugment and Mixup/CutMix for XS, and Mixup/CutMix for S dist_train.sh dist_train.sh. The slurm script's note says the same slurm_train.sh.
• Learning rate. Paper: initial 1e-3 §4.1. Code: default 2e-3 main.py. No script in the repo passes --lr.
• Batch and GPUs. Paper: 8 × A100 with total batch 2048 §4.1 app. B. That implies 256 per GPU. Code: default 128 per GPU main.py. The slurm script requests 16 GPUs at 128 each, summing to 2048 slurm_train.sh slurm_train.sh slurm_train.sh. The dist script takes the GPU count as an argument dist_train.sh.
• Weight decay, EMA, gradient-clip mode. Not stated in the paper. Code: weight decay 0.025 main.py, model EMA on with decay 0.99996 main.py, clip mode agc main.py.
• Distillation. Paper: hard distillation app. B; weighting not stated. Code: loss = (1 − α)·CE(class head, labels) + α·CE(distillation head, teacher argmax) with α = 0.5 by default losses.py main.py.
• Teacher. Paper: RegNetY-16GF with 82.9% top-1 app. B. Code: timm name regnety_160 with a DeiT-hosted checkpoint main.py. The two are assumed to be the same network; the weights were not checked against the 82.9% figure.
• Evaluation. Not stated in the paper. Code: at test time the class and distillation head logits are averaged swiftformer.py.
• Detection and segmentation. Paper: Mask R-CNN and semantic FPN settings §4.4 §4.5. Code: no training or evaluation code for these tasks in the clone; the only mention is a comment about mmdetection swiftformer.py. These results cannot be checked here.
• Latency. Paper: iPhone 14 CoreML measurements tab. 2. Code: no latency or CoreML export script in the clone; the README points to the Xcode 14 benchmark tool README.md.
• Ablations. Table 1 and the Appendix C numbers have no ablation code in the repo.

Pitfalls

• Scale factor is a no-op in the code. The 1/sqrt(d) term in eq. 4 cancels against the token-wise L2 normalisation swiftformer.py. Keep it if you follow the paper literally; it changes nothing.
• Token normalisation is unstated. The paper gives α without a normalisation or softmax eq. 4. The code uses an L2 normalisation over tokens. A softmax over tokens would be a different operation. Implement the code's version when you need the released weights to load and match.
• Key normalisation is unstated. The paper normalises only the query. The code normalises the keys too swiftformer.py. Without it the released weights will not reproduce the reported behaviour.
• Two output layers. The released model has an extra linear layer after the residual sum swiftformer.py. Plan the parameter count accordingly; the paper's eq. 6 does not show it.
• Positional information. The discussion says no positional encoding or attention bias is used app. F. Spatial mixing comes from the depth-wise convolutions only.
• Layer scales are load-bearing. The encoder's attention branch starts at 1e-5, so early training behaves almost like the identity swiftformer.py. A reimplementation without layer scale will train differently.
• Stem normalisation. BN and ReLU inside the stem are not in the paper's text §3.3 swiftformer.py.
• Recipe conflicts. The Mixup and CutMix values are swapped against the code defaults, and the released scripts disable some augmentations for XS and S. Use the scripts' values for those variants, and treat the paper's blanket recipe as the default case only.
• Batch size and LR. The paper's recipe (2048, 1e-3) does not match the code defaults (128 per GPU, 2e-3). The effective batch depends on your GPU count; check the scaling before reproducing.
• Latency is not throughput. The A100 throughput is measured at batch size 128 tab. 2. It does not transfer to batch-1 phone latency, and the phone numbers come from a CoreML/Neural Engine pipeline that is not in the repo.
• Dense-task numbers are unverifiable here. The Mask R-CNN and semantic FPN training code is not in the repo; the settings in §4.4 and §4.5 come from the paper alone.
• Compute. The paper reports about 36 h per classification run on 8 × A100 app. B. Hardware and cost for detection and segmentation are not stated.
• Equation numbers above are counted from the LaTeX source in order; if they disagree with a rendered PDF, cite the section instead.

Lineage

• Builds on: EfficientFormer, which provides the baseline, the training recipe and the teacher model §3.3 §4.1. Related: additive attention in NLP (Fastformer, arXiv:2108.09084), whose form uses three sequential steps with key-value interaction; this paper drops that key-value interaction §3.2.
• Compared against: MobileViT-v2 (separable self-attention), MobileViT, EdgeViT, EdgeNeXt (transpose attention), PoolFormer, MobileNet-v2/v3, EfficientNet, DeiT, Swin and ConvNeXt tab. 2 §2.
• Superseded by: not stated in this card's sources.