EdgeTAM: On-Device Track Anything Model

TL;DR

EdgeTAM makes SAM 2 run on mobile by attacking what the authors identify as the real latency bottleneck — memory attention, not the image encoder. They replace the dense per-frame memory feature map with a 2D Spatial Perceiver: a small set of learnable latents (split into global and 2D-windowed groups) that compresses the densely stored memories into ~16× fewer tokens while preserving 2D spatial structure. Paired with a two-stage SAM-2 → student distillation pipeline, the resulting model runs at 16 FPS on iPhone 15 Pro Max (22× faster than the smallest SAM 2 variant) and matches or beats prior VOS state-of-the-art on SA-V val/test, with a comparable image-segmentation profile to SAM and SAM 2.

Key claims

The latency bottleneck of SAM 2 on mobile is the memory attention blocks, not the image encoder — even replacing Hiera with RepViT-M1 and reducing memory attention from 4 to 2 blocks only reaches 2.5 FPS on iPhone 15 Pro Max [§3.2, Fig. 2].
Reducing the number of memory attention blocks “almost linearly” cuts decoder latency, and inside each block the cross-attention (current frame features × memory keys/values) is the dominant cost [§3.2, Fig. 2].
A naïve global Perceiver that compresses memories into a small set of learned vectors hurts segmentation quality because video object segmentation is a dense prediction task that needs explicit positional information; the Global Perceiver alone underperforms the baseline (62.0 / 60.6 vs 63.5 / 62.1 J&F on SA-V val/test) [§3.2, Tab. 3(b)].
The 2D Spatial Perceiver — window-partitioned memory with learnable latents that each attend to one non-overlapping local patch, with positional embeddings applied on output rather than input — recovers and exceeds the dense-memory baseline (63.1 / 62.4 on SA-V val/test, 15.7 FPS) [§3.2, Tab. 3(b)].
Combining 256 global + 256 2D-spatial latents beats either alone (64.4 / 62.5 on SA-V val/test, same 15.7 FPS) — global latents capture frame-level summaries, 2D latents preserve dense spatial structure [Tab. 3(b)].
Self-attention between the latents inside each Perceiver stage is necessary; removing it costs 1.8 J&F on SA-V val at no speed gain [Tab. 3(d)].
The two-stage distillation pipeline (MSE feature alignment on image-encoder output in image pre-training, plus MSE alignment on the memory-attention output during video training, on top of task losses) lifts SA-V val by 1.3 and SA-V test by 3.3 J&F at zero inference cost [§3.3, Tab. 3(a)].
End-to-end, RepViT-M1 + 2 memory-attention blocks + 2D Spatial Perceiver + distillation reaches 70.0 / 87.7 / 72.3 / 71.7 J&F on MOSE / DAVIS 2017 / SA-V val / SA-V test, running at 150.9 FPS on A100 and 15.7 FPS on iPhone — versus SAM 2.1-B+ at 76.6 / 90.2 / 76.8 / 77.0 at 0.7 FPS on iPhone (22× slower) [Tab. 2].
On the Segment Anything benchmark (SA-23, 23 datasets), EdgeTAM reaches 81.7 mIoU at 5 clicks vs SAM-H’s 81.3 — comparable image-segmentation performance despite being trained on a smaller mix without SAM 2’s internal datasets [Tab. 1].
On high-end GPUs (A100) the streaming-multiprocessor utilization is low; via Torch profile, CUDA kernel launch overhead from the CPU is the bottleneck — so the authors explicitly recommend judging EdgeTAM by edge-device latency, not GPU FPS [§4.5].

Method

EdgeTAM keeps SAM 2’s overall four-component shape (image encoder, mask decoder, memory encoder, memory attention) and changes two things. (1) Architecturally, the image encoder is RepViT-M1 (selected over ViT-Tiny for the speed/accuracy tradeoff) with 2 memory-attention blocks (down from 4). Each frame’s dense memory feature map M_t (size H×W×C) is replaced before memory attention by a much smaller token set produced by a 2D Spatial Perceiver module. The module has two parallel halves sharing parameters: a Global Perceiver where 256 learnable latents do single-head cross-attention over the full memory feature, followed by self-attention, producing global frame-summary vectors; and a 2D Spatial Perceiver that window-partitions M_t into non-overlapping patches, restricts each latent to attend only to its window, and applies positional embeddings (2D-RoPE) to the output rather than the input — this preserves spatial structure. Outputs of both halves are concatenated. The module is stacked twice. Net effect: memory-attention complexity drops from O(THW · HW) to O(THW · k) with k ≈ 0.25·HW, roughly 4× speed-up per block while improving SA-V accuracy.

(2) Training is the standard SAM 2 two-stage curriculum (SA-1B image pretraining, then video segmentation on SA-V + 10% SA-1B + DAVIS + MOSE + YTVOS), with three extra MSE losses for distillation against a SAM2-HieraB+ teacher: image-encoder feature alignment in both stages, plus memory-attention output alignment in stage 2. Progressive fine-tuning on 16- then 32-frame samples follows SAM 2.1; inference memory-bank size is held fixed at 7 frame-level memories + 16 object pointers.

Results

Headline: 22× faster than SAM 2 on iPhone, comparable accuracy on the hardest VOS benchmark.

iPhone 15 Pro Max latency: 15.7 FPS vs SAM 2-B+‘s 0.7 FPS (22.4×) [Tab. 2].
A100 throughput: 150.9 FPS vs SAM 2-B+‘s 64.8 FPS [Tab. 2].
SA-V val/test J&F: 72.3 / 71.7 vs SAM 2-B+‘s 73.6 / 74.1 — within 1.5–2.4 J&F at 22× the speed [Tab. 2].
MOSE val: 70.0 vs SAM 2-B+ 75.8 (5.8 J&F gap, still beats every non-SAM2 baseline including Cutie-base+ at 71.7) [Tab. 2].
DAVIS 2017: 87.7 vs SAM 2-B+ 90.9; YTVOS 2019: 86.2 vs 88.4 — small gaps on legacy benchmarks [Tab. 2].
Promptable Video Segmentation (PVS) with 5 clicks: 75.5 J&F vs SAM 2 75.4 and SAM+Cutie 72.2 — at 5 clicks EdgeTAM matches the SAM 2 teacher [Tab. 5].
Segment Anything (SA-23 image benchmark): 81.7 mIoU at 5 clicks vs SAM-H 81.3, SAM 2-B+ 83.7 — comparable to dedicated image-segmentation models despite being a tracking-first design [Tab. 1].
Ablation on memory-attention block count + backbone: RepViT-M1 + 2 blocks gives the best Pareto point (65.7 / 65.8 SA-V, 15.7 FPS); ViT-Tiny + 2 blocks scores higher on SA-V (67.9 / 66.0) but at 7.4 FPS [Tab. 3(c)].
Per-component contribution: 2D Spatial Perceiver alone gives 64.4 / 62.5 at 15.7 FPS; adding distillation gives 65.7 / 65.8 at the same speed (1.3 / 3.3 J&F gain free at inference) [Tab. 3(a)].

Why it’s interesting

EdgeTAM’s diagnosis — “the latency cost has migrated from the image encoder to memory attention” — is the same shape of finding that Radial Attention: O(n log n) Sparse Attention with Energy Decay for Long Video Generation makes for long video generation and that MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens makes for 100M-token memory models: once you bolt a long-history memory bank onto a transformer, the cross-attention against that bank, not the per-step network, dominates. Their fix — a learned 2D-windowed Perceiver whose latents preserve spatial structure — is a vision-task analogue of Raven Part 1 — Memory as a Set of Slots‘s “memory as a set of slots” line: replace a dense per-frame KV bank with a small fixed pool, and pay attention to what the pool’s geometry has to encode. The paper is the first wiki-filed instance of pushing a SAM-2-class video model onto a phone CPU/NPU with no quantization, and a useful counter-example to the bias in [[llm-inference-efficiency]] that on-device wins must come from weight quantization (1-bit Bonsai) or speculative decoding — here the win is purely architectural + distillation.