SAM 2.1 (Hiera-Tiny) mask decoder — LiteRT GPU

On-device LiteRT / TFLite conversion of the prompt-conditioned mask decoder of SAM 2.1 Hiera-Tiny (Meta, Apache-2.0), running fully on the mobile GPU via the LiteRT CompiledModel API (ML Drift / LITERT_CL delegate).

This is the lightweight, per-click half of the SAM 2 image path. Pair it with the SAM 2.1 Hiera-Tiny image encoder (run once per image, ~7 ms): the encoder produces the multi-scale feature pyramid, and this decoder turns a point prompt into segmentation masks per tap (a few ms each) — interactive "tap to segment".


Task	Mask decoder for promptable segmentation (SAM 2 image path)
Architecture	2-layer two-way transformer (token↔image cross-attention) + mask up-sampler
Inputs	`image_embeddings [1,256,64,64]`, `sparse_prompt [1,2,256]`, `feat_s1 [1,64,128,128]`, `feat_s0 [1,32,256,256]`
Outputs	`pred_masks [1,3,256,256]` (logits, 3 multimask candidates), `iou_scores [1,3]`
Precision / size	FP16, 17 MB
File	`sam2_tiny_mask_decoder_v2_fp16.tflite` (recommended). `sam2_tiny_mask_decoder_fp16.tflite` is the earlier build, kept for reference — see the note below.
Device	Pixel 8a — fully GPU (LITERT_CL), correct masks, ~7 ms/tap
Op set	`banned ops = NONE`, `>4-D tensors = 0` (BATCH_MATMUL ×15, SOFTMAX ×7, GELU ×2, CONV_2D ×2)

⚠ Residency ≠ correctness — and why v2 exists. The first build (sam2_tiny_mask_decoder_fp16.tflite) fully delegated to the GPU (358/358 LITERT_CL nodes, banned ops = NONE, >4-D = 0, desktop parity corr 1.0) yet returned silently wrong masks on the Pixel 8a GPU (corr 0.265 vs CPU; a face tap at IoU ≈ 0.62 on CPU collapsed to ≈ 0.10 with the mask on the background).

The cause was found by device A/B bisection: its attention was written with the batch dim collapsed (q/k/v shaped [heads, N, d], rank 3). The GPU delegate mis-computes that form. It is not an fp16 problem (forcing fp32 GPU compute still gives corr 0.473) and not LayerNorm (plain and overflow-safe LN give the same wrong result). The mask head's rank-2 matmul is innocent.

v2 keeps the leading batch dim (rank-4 SDPA, [1, heads, N, d]). Host numerics are identical (eager cos 0.999999); on the Pixel 8a GPU it restores corr 0.9998 / binary-IoU 0.999 vs CPU and is ~20 % faster (6.8 ms vs 8.5 ms). Inputs and outputs are unchanged, so v2 is a drop-in replacement. Note the companion encoder's rank-3 SDPA is GPU-correct — a healthy sibling graph proves nothing; only a numeric GPU-vs-CPU check on device catches this.

Pipeline (how the inputs are produced)

RGB image ──> image encoder (run once) ──> image_embeddings[1,256,64,64], feat_s1[1,64,128,128], feat_s0[1,32,256,256]
tap (x,y in 1024-space) ──> prompt encode (host-side, see below) ──> sparse_prompt[1,2,256]
                                              │
       image_embeddings + feat_s0/s1 + sparse_prompt ──> THIS decoder ──> 3 masks + 3 IoU
       pick argmax(IoU) ──> upsample 256×256 logits to image size ──> threshold > 0 ──> overlay

The decoder uses the encoder variant that already folds conv_s0 / conv_s1 + no_memory so its outputs are directly decoder-ready (no host reshaping between the two models).

Host-side prompt encoding (single positive point)

The tiny point→token step (a sin/cos positional encoding) is done on the host to keep the GPU graph sin/cos-free. For a positive click (x, y) in 1024×1024 model space, with the bundled constants posmat [2,128], point_embed[1] [256], not_a_point [256]:

c      = (([x, y]) + 0.5) / 1024          # normalize, half-pixel shift
c      = 2*c - 1
coord  = 2*pi * (c @ posmat)              # [128]
token0 = concat(sin(coord), cos(coord)) + point_embed[1]   # the positive point
token1 = not_a_point                      # the padding point
sparse_prompt = [[token0, token1]]        # [1, 2, 256]

This matches the upstream Sam2PromptEncoder to ~3.7e-7.

GPU-clean conversion (what was re-authored)

Converted with litert-torch, model-side rewrites only — no converter patch, each weights-faithful:

Two-way attention (×7): re-expressed as 3-D batched SDPA [heads, N, d] (a 4-D SDPA makes the delegate emit a BROADCAST_TO).
Mask up-sampler ConvTranspose2d (×2): replaced with the exact zero-stuff + Conv2d identity (TRANSPOSE_CONV is rejected on Pixel 8a; this is numerically identical, not a bilinear approximation).
Mask head: the hyper_in @ upscaled mask projection is kept ≤4-D (the upstream [1,1,4,256,256] 5-D tensor is collapsed; batch/point-batch are 1).
LayerNorm (×9): scale-before-square SafeLayerNorm (fp16-overflow-safe, mathematically identical).
Constants baked: image_positional_embeddings and the no-mask dense prompt are baked as buffers.
Multimask path: static slice [1:] of the 3 candidate masks — no dynamic-stability argmax / gather / where.

Fidelity (honest)

Eager re-authoring is numerically exact (cos = 1.000). End-to-end through the two FP16 tflite models (encoder → host prompt-encode → decoder) vs the PyTorch reference, for a center click:

Metric	value
mask logits cosine	0.999999
binary mask IoU (threshold 0)	0.99964
IoU-score head	ref `[0.936, 0.022, 0.399]` vs got `[0.936, 0.022, 0.399]`

The deepest 64×64 image embedding drifts slightly on the GPU (true-fp16 deep attention; see the encoder card). Mask boundaries are carried by the near-exact high-resolution features, so mask quality holds.

Minimal usage

Android (Kotlin, CompiledModel GPU)

// once per image - encoder on GPU (decoder-ready v2 variant from the companion repo)
val enc = CompiledModel.create(context.assets, "sam2_tiny_image_encoder_v2_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
// per tap - decoder on GPU (v2: rank-4 attention, GPU-correct; see the note above)
val dec = CompiledModel.create(context.assets, "sam2_tiny_mask_decoder_v2_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
// dec inputs: 0 image_embeddings[1,256,64,64], 1 sparse[1,2,256],
//             2 feat_s1[1,64,128,128], 3 feat_s0[1,32,256,256]
// dec outputs: pred_masks[1,3,256,256] logits, iou_scores[1,3] -> argmax(iou), threshold 0

Python (desktop verification)

MEAN = np.array([0.485, 0.456, 0.406], np.float32)
STD  = np.array([0.229, 0.224, 0.225], np.float32)
import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter

# 1) encode once (decoder-ready v2 encoder from the companion encoder repo)
img = Image.open("photo.jpg").convert("RGB").resize((1024, 1024))
x = ((np.asarray(img, np.float32) / 255 - MEAN) / STD).transpose(2, 0, 1)[None]
enc = Interpreter(model_path="sam2_tiny_image_encoder_v2_fp16.tflite"); enc.allocate_tensors()
enc.set_tensor(enc.get_input_details()[0]["index"], x); enc.invoke()
eo = {tuple(d["shape"]): enc.get_tensor(d["index"]) for d in enc.get_output_details()}

# 2) host prompt-encode one positive tap (px, py) in 1024-space (constants: this repo)
px, py = 512, 384
posmat, pe1, nap = np.split(np.fromfile("prompt_encode_const.bin", np.float32), [256, 512])
coord = 2 * np.pi * ((2 * (np.array([px, py], np.float32) + 0.5) / 1024 - 1) @ posmat.reshape(2, 128))
tok0 = np.concatenate([np.sin(coord), np.cos(coord)]) + pe1
sparse = np.stack([tok0, nap])[None].astype(np.float32)               # [1,2,256]

# 3) decode masks
dec = Interpreter(model_path="sam2_tiny_mask_decoder_v2_fp16.tflite"); dec.allocate_tensors()
feed = {(1,2,256): sparse}; feed.update(eo)                           # match inputs by shape
for d in dec.get_input_details(): dec.set_tensor(d["index"], feed[tuple(d["shape"])])
dec.invoke()
o = {len(d["shape"]): dec.get_tensor(d["index"]) for d in dec.get_output_details()}
masks, iou = o[4], o[2]                                               # [1,3,256,256], [1,3]
best = masks[0, iou[0].argmax()] > 0                                  # [256,256] binary mask
Image.fromarray(best.astype(np.uint8) * 255).resize(Image.open("photo.jpg").size).save("mask.png")

Training data & PII

SAM 2 was trained by Meta on SA-1B (licensed photos) and SA-V (licensed videos) with model-in-the-loop mask annotation. No new training was performed for this conversion — it is a weights-faithful format change of the public facebook/sam2.1-hiera-tiny checkpoint. Because the source data is real-world imagery it may incidentally contain people, faces, vehicles, signage and other PII; no PII was deliberately collected and this conversion adds none. Apply your own content/PII filtering as appropriate. See the SAM 2 release and paper for full dataset details.

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean.

Runtime	Backend	Graph on GPU	Latency
TFLite `benchmark_model` (`TfLiteGpuDelegateV2`) — `sam2_tiny_mask_decoder_v2_fp16.tflite`	GPU (OpenCL)	71 / 425	247.6 ms
TFLite `benchmark_model` (`TfLiteGpuDelegateV2`) — `sam2_tiny_mask_decoder_fp16.tflite`	GPU (OpenCL)	294 / 378	123.7 ms
TFLite `benchmark_model` — `sam2_tiny_mask_decoder_v2_fp16.tflite`	CPU (XNNPACK, 4 threads)	—	88.1 ms
TFLite `benchmark_model` — `sam2_tiny_mask_decoder_fp16.tflite`	CPU (XNNPACK, 4 threads)	—	86.8 ms

Any on-device figure recorded when this model shipped came from a different runtime. It was taken through LiteRT's own CompiledModel accelerator (logcat reports it as LITERT_CL), which is the path the Kotlin sample app and the LiteRT API use, and it appears elsewhere on this card. The rows above are the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. The two are not comparable, so read the rows above as a reproducible floor rather than as this model's speed on LiteRT.

On this delegate the CPU is the faster choice for sam2_tiny_mask_decoder_v2_fp16.tflite (88.1 ms on CPU against 247.6 ms on GPU), sam2_tiny_mask_decoder_fp16.tflite (86.8 ms on CPU against 123.7 ms on GPU) — worth knowing before you reach for the GPU on a mid-range phone.

Note that the GPU does not take the whole graph here (71 / 425 in sam2_tiny_mask_decoder_v2_fp16.tflite, 294 / 378 in sam2_tiny_mask_decoder_fp16.tflite); the remainder runs on the CPU and the split costs a per-partition round trip.