README.md

---
license: apache-2.0
library_name: litert
pipeline_tag: mask-generation
tags:
  - litert
  - tflite
  - sam2
  - segment-anything
  - mask-decoder
  - interactive-segmentation
  - on-device
  - gpu
base_model: facebook/sam2.1-hiera-tiny
---

# 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**](https://huggingface.co/facebook/sam2.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**](https://huggingface.co/litert-community/SAM2.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** |
| Device | Pixel 8a — GPU-resident (358/358 LITERT_CL) but run on **CPU** for correct masks (see below) |
| Op set | `banned ops = NONE`, `>4-D tensors = 0` (BATCH_MATMUL ×15, SOFTMAX ×7, GELU ×2, CONV_2D ×2) |

> **⚠ Residency ≠ correctness.** This decoder fully delegates to the LiteRT GPU (358/358 LITERT_CL
> nodes), but on the Pixel 8a its **GPU fp16 output is numerically wrong** — a face tap that the CPU
> decoder segments at IoU ≈ 0.62 collapses to ≈ 0.10 with the mask on the background under the GPU
> delegate. The companion encoder's GPU output is fine (encoder-GPU + decoder-CPU matches all-CPU
> exactly). It is **not** LayerNorm (plain vs. overflow-safe give the same wrong GPU result); the
> offending op is still being localized with a per-op GPU dump. **Run this decoder on CPU** (it is tiny
> and fast); the heavy image encoder is the part that benefits from the GPU.

## 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:

1. **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`).
2. **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).
3. **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).
4. **LayerNorm (×9)**: scale-before-square **SafeLayerNorm** (fp16-overflow-safe, mathematically identical).
5. **Constants baked**: `image_positional_embeddings` and the no-mask dense prompt are baked as buffers.
6. **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.

## Usage (Android / LiteRT CompiledModel)

```kotlin
// once per image — encoder on GPU
val enc = CompiledModel.create(assets, "sam2_image_encoder_v2_fp16.tflite", Options(Accelerator.GPU), null)
// per tap — decoder on CPU (GPU-resident but fp16-incorrect on device; see the residency note above)
val dec = CompiledModel.create(assets, "sam2_mask_decoder_fp16.tflite", Options(Accelerator.CPU), null)
// dec inputs (by index): 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]  -> pick argmax(iou), upsample, threshold 0
```

## 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](https://github.com/facebookresearch/sam2) and
[paper](https://arxiv.org/abs/2408.00714) for full dataset details.

## License

Apache-2.0, inherited from the upstream [SAM 2.1](https://huggingface.co/facebook/sam2.1-hiera-tiny).
This is a format conversion; all credit to the original authors (Meta AI).