DL0178 SAM: Promptable Segmentation

SAM splits segmentation into one expensive image step and a very cheap prompt step. A ViT-H image encoder (MAE-pretrained, about 632M parameters) takes a 1024×1024 image, converts it into patch tokens, and emits a single image embedding of shape that is computed once and cached. A tiny prompt encoder maps whatever the user supplies (clicks, a box, a coarse mask, text) into 256-dimensional tokens, and a two-block two-way cross-attention mask decoder reads the cached embedding together with those tokens to produce three mask logits plus three predicted IoU scores in roughly 50 ms. Generalization to unseen categories comes from the task definition rather than from any category machinery: the objective is “return a valid mask for this prompt”, there is no classification head and no fixed label set, so the model only ever learns a geometry-to-region mapping, supervised by 1.1B masks over 11M images (SA-1B) that a model-in-the-loop data engine collected without class labels. The one thing a category-free interface cannot resolve is ambiguity, and SAM handles it by predicting three nested candidates and back-propagating only the lowest-loss one.

(1) Asymmetric Compute: the ViT encoder runs once per image (about 0.15 s on a GPU) while the decoder runs once per prompt (about 50 ms, light enough for a browser), which is what makes interactive clicking feasible.
(2) Prompts Become Tokens: points and boxes are Fourier positional encodings of their coordinates plus a learned type embedding; a coarse mask prompt is instead embedded by a small convolutional stack and added elementwise to the image embedding.
(3) Two-Way Attention Decoder: prompt tokens attend to image features and image features attend back to the tokens, so a handful of tokens can reshape a dense 4096-token feature map at negligible cost.
(4) Dynamic Mask Head: each mask token is turned by an MLP into the weights of a linear classifier that is dotted with the 4x-upscaled embedding, so the mask is produced by a per-prompt hypernetwork rather than by fixed per-class filters.
(5) Ambiguity-Aware Output: three masks (subpart, part, whole) with a min-over-masks loss and an IoU-prediction token for ranking, which prevents the averaging that destroys single-output models on ambiguous clicks.
(6) Composable Interface: because the prompt is only geometry, any upstream module can drive SAM (a detector’s box, a CLIP text embedding, or a regular grid of points for fully automatic mask generation) with no retraining.

Figure 1: The whole design is an amortization argument. Everything expensive is a function of the image alone and is cached; everything that depends on the prompt is a few million parameters of decoder. Note the absence of any classification branch: the output is a region and a quality score, never a category.

The reason this interface transfers to categories that never appeared in training is that the prompt carries no semantics at all. A click is a coordinate, a box is two coordinates, and the encoder was pretrained to represent generic visual structure, so segmenting a species of coral or an industrial part reuses exactly the same computation as segmenting a cat. What must be learned instead is the notion of “object-ness at the scale the prompt implies”, and SA-1B provides that at massive scale precisely because its data engine was label-free: annotators clicked regions, then a partially trained SAM proposed masks that annotators only corrected, then the final stage generated masks automatically from a point grid with NMS and stability filtering. Ambiguity is the residual difficulty. A single click on a shirt pocket is consistent with the pocket, the shirt, and the person, so a single-output model trained with an averaged loss learns a blurred compromise. SAM emits three masks and only the one with the lowest loss receives gradient, which forces the three slots to specialize by scale.

Figure 2: One prompt, three valid answers. The min-over-masks loss back-propagates only through whichever slot best matches the ground-truth region, so slot 1 drifts toward subparts, slot 2 toward parts, and slot 3 toward whole objects. The IoU-prediction token then supplies the ranking the interface needs when no human is in the loop.

Mathematical Formulation:

Where:

• is the cached image embedding, the ViT-H encoder applied to the 1024×1024 image , and an optional coarse mask prompt whose convolutional embedding is added elementwise.
• is the token for sparse prompt , built from a random Fourier positional encoding of its coordinates plus a learned embedding for foreground point, background point, box corner, or text.
• is the decoder’s token sequence: prompt tokens plus three learned mask tokens and one IoU token .
• is the updated image feature map and the updated tokens after two blocks of two-way cross-attention; is a pair of transposed convolutions giving a map.
• is mask ‘s logit map, formed by dotting the dynamically predicted weights with every spatial location of , then bilinearly resized to 1024×1024.
• and are combined in a 20:1 ratio, and only the argmin slot receives gradient; the IoU head is trained separately with an MSE loss against the realized IoU.

Latency Accounting For Prompts On One Image:

The quadratic term belongs entirely to the encoder over tokens, while the decoder is linear in the tokens it carries, which is why 32 clicks cost 1.75 s with caching but 6.4 s if the ViT is re-run each time. Automatic mask generation exploits the same asymmetry: a grid of 1,024 point prompts, run in batches against one cached embedding, yields a full mask inventory for the image without ever re-encoding it.

Figure 3: Why the encoder/decoder split is the architecture, not an optimization. Caching turns the marginal cost of a prompt from 200 ms into 50 ms, and the gap widens with every additional click, which is exactly the regime an interactive annotation tool or a 1,024-point automatic sweep lives in.