Deploying Attention-Based Vision Transformers on the ANE

Part of the ANE Knowledge Base. Source: Apple ML Research, "Deploying Attention-Based Vision Transformers to Apple Neural Engine" (2024), code in references/ml-vision-transformers-ane.

This is the follow-up to the 2022 transformer work, extending the four principles to vision transformers (ViT) and their extra challenges: 2D feature maps, high-resolution inputs, window partitioning, and position embeddings. Apple cites Photos search, RoomPlan, and ARKit semantic features as production users of such models.

Two architectures are studied:

• DeiT/16-tiny — a conventional ViT, as baseline.
• MOAT (tiny-MOAT-0/1/2) — a hybrid Mobile-conv + attention network: MBConv (inverted residual) stages first, attention stages later. Mobile-friendly; higher ImageNet accuracy than DeiT at similar parameter count.

The repo ships WindowAttention, window_partition/window_reverse, the MOAT implementation, and export.py.

---

1. High-resolution inputs: local (window) attention#

Problem. Attention is quadratic in token length. At 512×512 input, global attention over the full feature map is prohibitively slow.

Solution. Partition the feature map into rectangular windows and attend within each window (as in Swin/MOAT). Cross-window information flow is recovered by other means — Swin uses shifted windows, MOAT uses its depthwise convolutions (the MBConv half of each block), which is convenient on the ANE because shifted windows have no native ANE support (WindowAttention explicitly "supports only non-shifting window attention").

In the repo (vision_transformers/model.py):

• attention_mode="global": window = whole feature map (fine at low resolution / late stages).
• attention_mode="local": fixed window (e.g., 8×8; default candidates 6/8/10 chosen to divide the feature size; sizes constrained to 6–16).
• MOAT stages: ("mbconv", "mbconv", "moat", "moat") — attention only enters at strides 16/32, where token counts are manageable.

2. The 6D-tensor problem: "relay" window partitioning in ≤5D#

Problem. Standard window partition/reverse goes through a 6D tensor (N, C, Nh, Nw, Hw, Ww) (windows-per-axis × window-size-per-axis). ANE supports at most 5D tensors.

Solution. "Relay" the partition, factoring out one dimension at a time — first height, then width — so no intermediate exceeds 5D. From vision_transformers/attention_utils.py:

def window_partition(x, window_size):            # x: (B, H, W, C)  — NHWC!
    B, H, W, C = x.shape
    x = x.reshape((B, H // wh, wh, W, C))         # 5D: split H
    x = x.reshape((B * H // wh, wh, W, C))        # 4D: fold into batch
    x = x.reshape((B * H // wh, wh, W // ww, ww, -1))  # 5D: split W
    x = x.permute((0, 2, 1, 3, 4))
    windows = x.reshape((-1, wh, ww, C))          # (B·num_windows, wh, ww, C)
    return windows

window_reverse mirrors the same staged reshapes back. Comment in the code shows the shape relay: 1,12,16,160 → 1,2,6,16,160 → 2,6,16,160 → 2,6,2,8,160 → ...

Generalizable trick: any op that "needs" rank > 5 can usually be staged as fold-into-batch reshapes plus one small permute.

3. Tensor layout for partitioning: NHWC, not NCHW#

Problem. ANE processes the last dimension in 64-byte batches and pads it to 64 bytes if smaller. Window sizes are small: with NCHW and a 7×7 window, the last dim holds 7 FP16 values = 14 bytes → 50 bytes padding. Worst case (singleton last dim) is a 32× effective slowdown in FP16.

Solution. For the partition/reverse path, use NHWC: the channel dim (usually a multiple of 32) sits last and aligns naturally. The MOAT block transposes NCHW→NHWC once before LayerNorm + partition and back once after window reverse — "instead of looping on each partitioned window":

# vision_transformers/model.py — MOATBlock.forward
output = self._mbconv(inputs)                 # NCHW
shortcut = output
output = output.permute((0, 2, 3, 1))         # NHWC, once
assert output.shape[-1] % 32 == 0, "ANE buffer not aligned"
output = self._attn_norm(output)              # nn.LayerNorm on last dim (=C)
...window_partition → WindowAttention → window_reverse...
output = output.reshape((N, H, W, C)).permute((0, 3, 1, 2))  # back to NCHW, once
output = shortcut + output

Note the interplay of constraints: NHWC also lets plain nn.LayerNorm work (it normalizes the last dim), and _build_model(channel_buffer_align=True) rounds all hidden sizes up to multiples of 32 channels to keep the last dim 64-byte aligned.

Takeaway: the "right" layout is op-dependent. Sequence transformers put S last (BC1S); window partitioning puts C last (NHWC). What is invariant is the rule: the last axis must be large and 64-byte-friendly.

4. Attention internals: same recipe as NLP, plus split_softmax#

Inside WindowAttention (vision_transformers/attention_utils.py) the 2022 recipe reappears:

• Separate 1×1 Conv2d projections for Q, K, V — the code comments "Use separate conv1x1 projection to avoid L2 cache hit" (vs. one fused QKV projection).
• Per-head torch.split + einsum bchq,bkhc->bkhq / bkhq,bchk->bchq, one transpose on K ("transposeThenSplit is more efficient than the other way around").
• split_softmax / split_head=True (default): softmax applied per head on dim=1. "Splitting on the softmax… increases the chance of L2 residency and parallelizes the computation for the softmax layer" — flagged as the key latency win, since softmax is quadratic in token length. (Linear-attention alternatives like CosFormer exist but trade off accuracy.)

5. Position embeddings: file size vs. token length#

For ViTs the position-embedding choice materially affects model file size and latency on device (PEType enum in the repo):

PE type	Extra parameters	Growth	Notes
RPE (Swin-style relative PE)	num_heads × (2Wh−1)(2Ww−1) table, gathered to token_len² per head	quadratic in token length	Significant file-size and latency overhead at large windows
SINGLE_HEAD_RPE	one shared table across heads	quadratic ÷ num_heads	"Reduces the file size of the positional embedding to 1/num_heads of the original RPE"
LePE_ADD (Locally-enhanced PE)	3×3 depthwise conv on V (3·3·dim per block) + absolute PE 1 × token_len × dim	linear in token length	Depthwise conv encodes locality into the value tensor; "significantly smaller than RPE"

Implementation notes from WindowAttention:

• RPE bias is looked up from relative_position_bias_table via a precomputed relative_position_index and added per head before the split softmax — with single-head RPE, all heads add the same table (rpe_idx = 0).
• LePE runs LePE_for_Value (depthwise 3×3, groups=dim, padding="same") on the 2D value map and adds it per head after attn @ v; the absolute PE is added to the input tokens.

6. Results#

• Optimized tiny-MOAT-1 is "multiple times faster" on ANE than the third-party open-source MOAT implementation.
• At high resolution (512×512), optimized MOAT is also much faster than optimized DeiT/16-tiny — locality wins when token counts explode.
• Tiny-MOAT-1 additionally has higher ImageNet accuracy than DeiT-tiny at similar parameter count.
• Figure 3 of the article / repo assets show Xcode device measurements across iPhone models.

7. Checklist for putting a ViT on the ANE#

1. Prefer a hybrid conv+attention backbone (MOAT-like); keep attention in low-resolution stages.
2. High-res inputs → local window attention, window sized to divide the feature map (6–16 per side); no shifted windows.
3. Implement partition/reverse with the 5D relay pattern; NHWC during partitioning; transpose once in, once out.
4. Keep channels a multiple of 32 (channel_buffer_align).
5. Use separate 1×1 conv Q/K/V, per-head split, einsum matmuls, split softmax.
6. Prefer LePE (or single-head RPE) over full RPE for large token lengths.
7. Export with torch.jit.trace + ct.convert(convert_to="mlprogram") and check the Xcode performance report (doc 05).

Next: 05 — Deployment Workflow.