Lab 6-04: Promptable Segmentation — Build a mini-SAM From Scratch

Mission. You will build the Segment Anything Model's entire architecture at toy scale, in pure PyTorch, on synthetic data, offline, on a CPU: an image encoder that runs once per image, a prompt encoder that turns a mouse click into a token for almost zero FLOPs, and a fast mask decoder with genuine two-way attention that emits three candidate masks plus a self-assessed quality score for each. By the end you will understand — because you implemented every line — why SAM predicts multiple masks for one click, how a model can grade its own output without ground truth, and why the encoder/decoder asymmetry is the single most important serving decision in interactive vision systems.

Table of Contents

• Why this lab exists
• The architecture at a glance
• File map
• Setup and how to run
• The synthetic task and why it contains ambiguity
• Concepts under the hood
  • Fourier features and spectral bias
  • Two-way attention
  • Mask tokens and hypernetworks
  • BCE, dice, and focal losses
  • Min-loss-over-candidates and the IoU head
  • Toy-scale concessions and why real SAM does not need them
• TODO walkthrough
• Success criteria
• From mini-SAM to real SAM and SAM2
• Stretch goals
• Interview Q&A
• References

Why this lab exists

SAM (Segment Anything Model, Kirillov et al. 2023) is usually presented as a dataset story — 1.1 billion masks — but its architecture is first and foremost a systems design lesson, and that lesson survives intact at 1/10,000th the scale:

1. Amortize the expensive part. Segmenting interactively means answering many prompts about the same image. SAM splits the network so that everything prompt-independent (a 632M-parameter ViT-H) runs once and is cached, while everything prompt-dependent (a ~4M-parameter decoder) runs per click in ~10 ms. The user experiences real-time segmentation even though the full model is enormous. If you take away one idea, take this one: find the cut in your computation graph that separates "per-item" from "per-query" work, and cache at that cut.

2. Make prompting near-free. A click is two numbers. SAM encodes it with a fixed random-Fourier positional code plus one learned embedding — a few thousand FLOPs. No conv stack, no MLP tower. Cheap prompts are what make the amortization worthwhile.

3. Design for ambiguity instead of pretending it away. A single click is genuinely underspecified — click on a shirt logo and you might mean the logo, the shirt, or the person. A single-output model trained on such data averages the interpretations and produces a blurry compromise that is no valid answer. SAM outputs K=3 masks, trains only the best-matching one per sample (min-loss-over-candidates), and adds an IoU-prediction head so the model can rank its own candidates at inference, when no ground truth exists.

This lab rebuilds all three decisions small enough to train in ~90 seconds on a laptop CPU, and the tests verify the behaviors (ambiguity divergence, honest self-ranking), not just the shapes.

The architecture at a glance

                      ┌──────────────────────────────────────────────┐
   image (1×48×48) ──▶│ ImageEncoder: conv stack, stride 4           │  EXPENSIVE
                      │ (SAM: ViT-H, 632M params, run ONCE, cached)  │  once per image
                      └──────────────────┬───────────────────────────┘
                                         │ features (64×12×12)
                                         ▼
   click (x, y) ──▶ PromptEncoder ─── prompt token (1×64)             CHEAP
                      │  γ(p)=[sin 2πBp, cos 2πBp] + label embed      per prompt
                      │
                      ├─ dense PE  (144×64: same γ over grid cells)
                      └─ dense prompt "beacon" (cos-sim of click γ vs each cell γ)
                                         │
                                         ▼
                      ┌──────────────────────────────────────────────┐
   learned tokens ──▶ │ MaskDecoder: 2 × TwoWayAttentionBlock        │  FAST
   [iou][m0][m1][m2]  │   1. token self-attention                    │  per prompt
                      │   2. cross-attn tokens → image               │
                      │   3. MLP on tokens                           │
                      │   4. cross-attn image → tokens  ◀── the      │
                      │      (image learns about the prompt) 2nd way │
                      │  PE re-injected into q,k at EVERY step       │
                      └──────┬────────────────────────┬──────────────┘
                             │ mask tokens out        │ image stream out
                             ▼                        ▼
                      hypernetwork MLPs        upscale ×4 (ConvT ×2)
                      3 weight vectors (16-d)  16ch × 48 × 48
                             └───────┬────────────────┘
                                     ▼   per-pixel dot product
                        mask logits (3×48×48)      iou_pred (3), self-graded

File map

Setup and how to run

Everything runs offline on CPU with the repo virtualenv (Python 3.14, torch ≥ 2.2):

cd AI-Engineer/cv-engineer/phase-06-sota-models/lab-04-promptable-segmentation

# 1. Verify the reference solution passes (≈ 90 s, one shared training run):
/Users/s0x/src/10xdev/ai-enigneer/.venv/bin/python -m pytest test_lab.py -q

# 2. Work through lab.py, then test YOUR implementation:
LAB_MODULE=lab /Users/s0x/src/10xdev/ai-enigneer/.venv/bin/python -m pytest test_lab.py -q

# 3. Train + evaluate the reference end-to-end by hand:
/Users/s0x/src/10xdev/ai-enigneer/.venv/bin/python solution.py

Determinism: all randomness flows from explicit seeds (np.random.default_rng(seed) for data, torch.manual_seed for parameters, a private torch.Generator inside the prompt encoder), so runs are reproducible on a given machine.

The synthetic task and why it contains ambiguity

A sample is a triple (image, point-prompt, ground-truth mask) on a 48×48 canvas:

• Non-overlapping scenes (70%): 1–3 shapes (circle, square, triangle) of random size and intensity on a noisy background, placed by rejection sampling so they never touch. The prompt is a random pixel inside one chosen shape; the GT is that shape's mask. One click → one right answer.
• Overlapping scenes (30%): a small shape drawn fully inside a large one, with distinct intensities so both boundaries are visible. The prompt lands in the overlap region, and the GT is the small shape or the large shape with probability ½ each.

That second family is the point of the lab. The identical (image, click) input appears in the dataset with two different correct answers. A single-mask model minimizing expected loss on such data converges to the average of both masks — a half-confident blob that matches neither interpretation. A multi-mask model with min-loss-over-candidates training can dedicate one candidate to "the small thing under the cursor" and another to "the big thing under the cursor" and win on both. The test suite explicitly checks that your trained candidates disagree on overlap clicks (pairwise IoU between candidates far below 1).

Concepts under the hood

Fourier features and spectral bias

What. Instead of feeding raw coordinates (x, y) ∈ [0,1]² to the network, we project them through a fixed random matrix B ~ N(0, σ²) (shape 2 × D/2) and take γ(p) = [sin(2πBp), cos(2πBp)], giving a D-dimensional code.

Why raw coordinates fail. MLPs exhibit spectral bias: viewed through their neural tangent kernel, they fit low-frequency components of the target function orders of magnitude faster than high-frequency ones. A mask boundary is a high-frequency function of position, and two clicks one pixel apart (Δ=1/48) produce inputs that differ by 2%, so their embeddings — and therefore their masks — stay nearly identical for most of training.

How the fix works. After the random projection, a 1-pixel move rotates each (sin, cos) pair by an angle proportional to that row's frequency. With frequencies drawn at scale σ, nearby points become linearly separable in embedding space from step one; Tancik et al. (2020) show this converts the network's kernel into a stationary, tunable-bandwidth kernel. It is the same idea as the Transformer's sinusoidal positional encoding, with random Gaussian frequencies instead of a fixed geometric ladder. σ trades smoothness for resolution: too low and nearby prompts collapse together, too high and the encoding decorrelates between adjacent feature-grid cells. We use σ=4 for a 12×12 grid.

Crucially, the same basis γ encodes both the click and the feature-grid cell centers (the "dense PE"), so "is this cell under the click?" reduces to a dot product between matching signatures.

Two-way attention

What. Standard attention: every query builds a weighted average of values, with weights softmax(q·kᵀ/√d). Cross-attention lets one sequence (queries) read from another (keys/values). A standard transformer decoder is one-way: object queries read from the image; the image is never updated.

Why SAM goes both ways. SAM's decoder must produce a pixel-dense output (a mask), not just token-level answers. Its trick: masks are decoded as dot products between token vectors and per-pixel features (see hypernetworks below). For that dot product to cut out "the object under the click", the per-pixel features themselves must already encode "how do I relate to this prompt?" — so the decoder needs a path for information to flow into the image stream, not just out of it.

How one block works (implemented in TwoWayAttentionBlock, residual + LayerNorm after every step):

1. Token self-attention — the iou/mask/prompt tokens negotiate among themselves; this is where mask tokens coordinate who covers which interpretation.
2. Cross-attention, tokens → image — tokens gather visual evidence (what intensity is under the click, how far does that region extend).
3. MLP on tokens — per-token nonlinear mixing.
4. Cross-attention, image → tokens — the second way: each of the 144 image cells attends over the 5 tokens and updates itself with prompt-conditioned information.

One faithful detail matters enormously in practice: positional encodings are re-added to queries and keys at every attention call (values stay raw), exactly as in SAM's TwoWayTransformer. Residual mixing dilutes a once-added PE layer by layer; if the click-to-cell alignment signal fades from the attention logits, the decoder degenerates into prompt-ignoring "segment every object" behavior — we verified this failure mode empirically while building the lab.

Mask tokens and hypernetworks

What. The decoder prepends 4 learned embeddings to the prompt token: one IoU token and K=3 mask tokens (descendants of DETR's object queries — learned vectors that each "own" one output slot). A hypernetwork is a network whose output is the weights of another network.

How masks are decoded. After the two-way blocks, each mask token is passed through its own small MLP to produce a 16-dimensional weight vector w_k. The image stream is reshaped back to a map and upscaled ×4 by two transposed convolutions (a transposed convolution reverses a stride-2 convolution's geometry: each input cell scatters a learned 2×2 kernel into the output, doubling resolution). The candidate mask is then simply

mask_k[y, x] = ⟨ w_k , features[:, y, x] ⟩

— a per-pixel linear classifier whose weights were generated from the token. All K candidates share every pixel feature and differ only in K tiny vectors, which is why emitting 3 masks costs barely more than emitting 1.

BCE, dice, and focal losses

Binary cross-entropy (BCE) treats each pixel as an independent binary classification: −[t·log p + (1−t)·log(1−p)] averaged over pixels. It gives smooth, well-scaled gradients everywhere but has a structural flaw for segmentation: it weights every pixel equally, and masks are tiny relative to the background (a 30-px shape on a 2304-px canvas is ~1% positive). A model can score excellent average BCE by confidently predicting "background everywhere".

Dice loss fixes the imbalance by scoring region overlap instead of per-pixel accuracy: 1 − 2|P∩T| / (|P|+|T|) computed on sigmoid probabilities (an ε in numerator and denominator keeps it defined for empty masks). Because both terms scale with mask size, a small object contributes as much as a large one. Its weakness is the mirror image of BCE's: gradients are poorly conditioned when predictions are near-empty. Using BCE + dice together — as this lab does — is the standard pairing: BCE supplies stable pixel-level gradients, dice supplies scale-invariant region pressure.

Focal loss (used by real SAM, weighted 20:1 with dice) is BCE reshaped by (1−p_t)^γ: pixels the model already classifies confidently are down-weighted, focusing gradient on hard pixels — which, in segmentation, are precisely the boundary pixels. At our resolution plain BCE converges just as well, so we keep the simpler pairing; swapping in focal is a stretch goal.

Min-loss-over-candidates and the IoU head

Training the K candidates. Every candidate is scored against the single GT mask (BCE + dice), and the mask loss is backpropagated only through the arg-min candidate. On an ambiguous prompt, whichever candidate is currently closest to this sample's interpretation gets pulled further toward it — a winner-take-all dynamic that makes candidates specialize into distinct interpretations instead of all regressing to the blurry mean. This is SAM's exact trick, and it is the multiple-choice-learning literature's answer to multimodal targets.

Training the IoU head. At inference there is no ground truth, so something must rank the 3 candidates. The IoU head is trained with MSE to predict each candidate's actual IoU with the GT (computed on the binarized masks, detached so this loss never reshapes the masks themselves — it only teaches the judge). Supervising all K candidates (not just the winner) is what lets the head rank.

The honest metric. evaluate_iou reports the IoU of the candidate the head ranked first — never the candidate closest to ground truth. Selecting by GT would leak the answer into the metric; head-selected IoU is what a user actually experiences, and it is strictly ≤ the oracle. The gap between the two is exactly "how good is the model at judging itself", a quantity worth monitoring in any deployed system.

Toy-scale concessions and why real SAM does not need them

Three places where this lab deliberately deviates from the paper — each is a lesson about what scale buys:

1. Dense prompt beacon. We add cos-sim(γ(click), γ(cell)) as an explicit channel to the decoder's image stream. Real SAM's decoder learns this alignment inside its attention logits over ~10⁵ ViT-scale steps; a 2-block decoder given a few hundred CPU steps does not, and without the beacon it collapses to segmenting every object. (SAM itself adds dense mask-prompt embeddings to the image stream through exactly this pathway, so the mechanism is canonical — only its use for points is our concession.)
2. Mask-quality statistics for the IoU head. We feed the head four GT-free cues per candidate — area, mean confidence, boundary contrast against the image, inside-mask intensity variance — alongside the token outputs. Real SAM's head reads only the IoU token and learns equivalent evidence end-to-end. The cues are all computable at deployment, so the metric stays honest.
3. ε-relaxed winner-take-all (relax=0.02 × mean candidate loss). Pure WTA gives losing candidates zero gradient on samples they lose, so their off-role behavior drifts arbitrarily and becomes unpredictable — for the IoU head, ungradeable. The tiny relaxation keeps losers sane without stopping specialization. SAM's data diversity provides this stabilization for free.

One more empirical finding worth internalizing: with only a few hundred training scenes this model memorizes the training set (train IoU climbs while held-out IoU stalls). Because the generator is nearly free, the fix is simply more data (default 4096 scenes) — a miniature of the SA-1B lesson that data volume is the regularizer.

TODO walkthrough

Work through lab.py in numeric order; each TODO's docstring carries the detailed hints. The order follows the dependency chain:

1. make_dataset — TODO(1). Rejection-sample non-overlapping scenes; build nested overlap pairs; sample the prompt inside the target (or overlap) region; GT flips small/big on ambiguous samples. Single default_rng(seed) for everything — a test checks bit-identical replay.
2. PromptEncoder — TODO(2). Buffer freq and the label embedding from a private seeded torch.Generator (so two encoders with equal seeds match exactly); fourier, forward, dense_pe (same basis on cell centers), dense_prompt (the beacon).
3. ImageEncoder — TODO(3). Five convs, two of them stride-2: (1, 48, 48) → (64, 12, 12). No prompt anywhere in this class — that is the contract.
4. Attention — TODO(4). By hand: q/k/v/out projections, head split, softmax(q·kᵀ/√d_head), merge. No nn.MultiheadAttention.
5. TwoWayAttentionBlock — TODO(5). The four steps, post-LN residuals, and PE re-injection into q and k at every call.
6. mask_quality_stats — TODO(6). Morphology via max_pool2d (dilation) and −max_pool2d(−·) (erosion) to get 1-px inner/outer boundary rings; ε-guard all divisions so empty masks give zeros, not NaNs; everything under no_grad.
7. MaskDecoder — TODO(7). Token concatenation ([iou, m0, m1, m2, prompt]), input_norm neck (without it the conv activations' scale drowns the unit-scale PE in the attention logits), beacon injection, blocks, final tokens→image read, upscale, hypernetwork MLPs, einsum masks, stats-conditioned IoU head with sigmoid.
8. MiniSAM — TODO(8). Wire the three parts, but keep encode_image / decode_prompt as separate methods — the split is the serving lesson.
9. mask_iou + dice_loss — TODO(9). Reductions over the last two dims so both broadcast over arbitrary leading dims.
10. sam_loss — TODO(10). Per-candidate BCE+dice → argmin → gather winner; + relax × mean; + iou_weight × MSE(iou_pred, detached actual IoU of all candidates).
11. train_mini_sam + evaluate_iou — TODO(11). Adam + cosine decay, seeded minibatch sampler, loss history; evaluation selects each sample's mask by iou_pred.argmax and reports mean IoU vs GT.

Success criteria

Tests passing is necessary, not sufficient. You are done when you can:

• Explain why clicking the same pixel can legitimately require two different masks, and trace how min-loss-over-candidates converts that ambiguity into candidate specialization (rather than averaging).
• State what breaks if you (a) remove PE re-injection from the attention calls, (b) remove the LayerNorm neck, (c) shrink the dataset to 300 scenes — and why, mechanically, in each case.
• Defend the honest metric: why the eval selects masks with the IoU head instead of picking the candidate closest to GT, and what the oracle-vs-head gap measures.
• Sketch the latency budget of an interactive SAM service and say which components run per-image vs per-click, and where the cache sits.
• Quantitative bars (the tests enforce them): loss falls > 40% (reference: ~80–87%); held-out head-selected mean IoU > 0.5 (reference: ≈ 0.65); on overlap clicks the 3 candidates' mean min pairwise IoU < 0.5 (reference: ≈ 0.03–0.14) with at least 2 non-trivial masks.

From mini-SAM to real SAM and SAM2

What changes at scale — and what does not:

• Image encoder. Ours: 5 convs, ~100K params, 12×12 features. SAM: a ViT-H pre-trained with MAE, 632M params, windowed attention on 1024×1024 inputs producing 64×64×256 embeddings — >99% of total compute, which is precisely why the once-per-image split exists. The contract (prompt-independence) is identical.
• Data engine, not just a dataset. SA-1B's 1.1B masks were produced by a three-stage flywheel: assisted-manual (annotators correct model output), semi-automatic (model proposes, humans fill gaps), fully automatic (a 32×32 point grid prompts the model, masks filtered by the IoU head's own confidence and stability). The IoU head you built is what makes stage three possible — a model that can grade itself can label data for its successor.
• Prompt types. We built foreground points. SAM adds background points (second label embedding), boxes (two corner embeddings), rough masks (conv-embedded dense prompts — the pathway our beacon reuses), and demonstrated text prompts via CLIP embeddings (the ancestor of open-vocabulary pipelines like Grounded-SAM).
• SAM2 (2024) extends promptable segmentation to video: a streaming architecture where each frame's features are enriched by memory attention — cross-attention into a FIFO memory bank of encoded past frames and past mask predictions, plus object pointer tokens. Prompt once on frame 1, and the memory carries the object identity through occlusions across the clip, at ~44 fps. It is the same amortization lesson upgraded to time: encode each frame once, let cheap attention against cached memory replace re-prompting.
• Try the real thing (outside this air-gapped environment): pip install "git+https://github.com/facebookresearch/sam2.git", download a checkpoint (e.g. sam2.1_hiera_small.pt), then SAM2ImagePredictor.set_image(img) once followed by many predict(point_coords=..., multimask_output=True) calls — you will find the API mirrors encode_image / decode_prompt from this lab almost method-for-method. Ultralytics and HuggingFace transformers (facebook/sam2-hiera-small) ship wrappers as well.

Stretch goals

1. Box prompts. Encode a box as two corner tokens (each = γ(corner) + corner-role embedding), concatenate with the point token, and extend make_dataset to emit boxes. Boxes are near-unambiguous — verify the 3 candidates converge for box prompts while still diverging for overlap points.
2. Iterative point refinement. SAM's interactive loop: after the first prediction, sample a correction click from the error region (foreground click in false-negative area, background click in false-positive area), feed both points back, repeat. Requires a background-point label embedding. Measure IoU vs number of clicks — it should climb monotonically.
3. Distill the encoder. Train a 2-conv "student" encoder to match your trained encoder's features (MSE in feature space), freeze the decoder, and measure the IoU cost of the ~4× cheaper encoder. This is exactly the MobileSAM/EfficientSAM recipe (they distill TinyViT from ViT-H).
4. Focal loss. Replace BCE with (1−p_t)^γ-weighted BCE (γ=2) at SAM's 20:1 focal:dice weighting and compare convergence speed and boundary sharpness.
5. Remove a load-bearing part. Delete PE re-injection (pass zeros as token_pe/image_pe after block 1) or the beacon, retrain, and watch held-out IoU collapse — the fastest way to convince yourself why each piece exists.

Interview Q&A

Q1: Why does SAM predict three masks and an IoU score for each, instead of a single mask?

Because a point prompt is inherently ambiguous — a click on a logo is consistent with the logo, the shirt, and the wearer — and ambiguity breaks single-output regression. If one output head is trained on data where the same input has several valid targets, the loss-minimizing prediction is the average of those targets: a blurry mask that is none of the valid answers. SAM instead emits K=3 candidates (empirically enough to cover whole/part/subpart nesting) and, during training, backpropagates the mask loss only through the candidate with the lowest loss on that sample. This winner-take-all assignment lets different mask tokens specialize into different interpretations — the multimodality moves from "impossible for one head" to "one head per mode". But it creates an inference problem: with no ground truth available, something must choose among the 3. That is the IoU head's job: trained with MSE to predict each candidate's actual IoU with ground truth, it lets the model rank its own outputs at deployment (return the argmax, or all three with scores). The same self-grading signal is what filtered SA-1B's fully-automatic masks — the head is both a serving component and a data-engine component. In this lab you can see all of it in miniature: on overlap clicks the trained candidates disagree (pairwise IoU ≈ 0.1), and held-out evaluation selects masks by the head alone.

Q2: Why is the encoder/decoder asymmetry the key serving insight in SAM?

Because it converts an interactive workload from "run a 632M-parameter model per click" into "run it once per image, then run a ~4M-parameter decoder per click". The design commits to a strict information contract: the image encoder never sees the prompt (so its output is a reusable function of the image alone), and the prompt encoder is trivially cheap (Fourier code + one embedding). All prompt-conditioned computation is pushed into a small decoder that fuses cached image features with prompt tokens via two-way attention in ~10 ms. The serving consequences compound: the embedding is computed once and cached (even precomputed offline for whole galleries); the heavy encoder can run server-side on GPU while the decoder runs client-side — in-browser SAM demos ship the decoder as an ONNX model; latency per click is decoder-only, which is what makes the tool feel real-time; and throughput scales with clicks-per-image, which in annotation workloads is high (SA-1B prompted a 32×32 grid = 1024 prompts against one cached embedding). SAM2 extends the same amortization through time: frames are encoded once into a memory bank, and cheap memory attention replaces re-prompting every frame. The general principle — locate the cut between per-item and per-query computation and cache at that boundary — transfers to retrieval (document embeddings vs query), recommendation (item towers vs user request), and LLM serving (KV-cache vs new tokens).

References

• Kirillov, A., Mintun, E., Ravi, N., Mao, H., Rolland, C., Gustafson, L., Xiao, T., Whitehead, S., Berg, A., Lo, W-Y., Dollár, P., Girshick, R. Segment Anything. arXiv:2304.02643 (2023).
• Ravi, N., Gabeur, V., Hu, Y-T., et al. SAM 2: Segment Anything in Images and Videos. arXiv:2408.00714 (2024).
• Lin, T-Y., Goyal, P., Girshick, R., He, K., Dollár, P. Focal Loss for Dense Object Detection. arXiv:1708.02002 (2017).
• Milletari, F., Navab, N., Ahmadi, S-A. V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation (the dice-loss paper). arXiv:1606.04797 (2016).
• Tancik, M., Srinivasan, P., Mildenhall, B., et al. Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains. arXiv:2006.10739 (2020).
• Vaswani, A., Shazeer, N., Parmar, N., et al. Attention Is All You Need. arXiv:1706.03762 (2017).
• Carion, N., Massa, F., Synnaeve, G., et al. End-to-End Object Detection with Transformers (DETR — origin of learned query tokens). arXiv:2005.12872 (2020).
• Guzmán-Rivera, A., Batra, D., Kohli, P. Multiple Choice Learning: Learning to Produce Multiple Structured Outputs (the min-loss-over-candidates ancestry). NeurIPS 2012.