SAM-Med3D: Towards General-purpose Segmentation Models for Volumetric Medical Images

In plain words

A medical scan is a stack of many slices, like a loaf of bread, and a doctor often needs the outline of one organ or tumour through the whole stack. Earlier tools for this were built for flat photos and worked one slice at a time, which meant the user had to click on many slices. They also lost track of how one slice connects to the next.

This paper builds a model that works on the whole 3D stack at once. The user clicks a single point inside the target, and the model returns the 3D outline of that structure. The model is trained on a large, mixed collection of public and private 3D scans with labelled structures, and it is trained from scratch rather than being adapted from a flat-image model.

The authors report that one click is enough to match or beat the older slice-by-slice tools on their validation set, which covers many organs, lesions and scan types. Those results come from the authors' own tests, so treat them as their claims until someone else reproduces them.

TL;DR

• A promptable 3D segmentation model for CT, MR and ultrasound volumes: a 3D ViT image encoder, a 3D point-prompt encoder and a 3D mask decoder §4.1 fig. 3.
• Trained from scratch on SA-Med3D-140K, which the authors say has 22K volumes and 143K 3D masks in 245 categories §1 §3.1.
• Training has two stages: pre-training on all 131K training masks for 800 epochs, then fine-tuning on about 75K filtered masks §4.2.
• Author-reported: 76.27% overall Dice with one 3D point, against 36.83% for SAM-Med2D with N points tab. 3. Ten 3D points give 80.71% tab. 3.
• The released code has no sliding-window inference and does not reproduce the paper's data cleaning or stage-2 selection. Several training details differ from the text. See Paper vs code and Pitfalls.

Problem

Flat-image segmentation models do not handle volumes well. The 2D Segment Anything Model (SAM) is not usable directly on volumetric medical images, because it does not model 3D structure and the domain differs from natural photos §1. The two common workarounds both have costs. Slice-by-slice methods need a prompt for every slice and do not model how slices connect §1. Adapter methods keep a frozen 2D encoder and learn only a new dimension, and they train on small, narrow datasets §1 §2.

The paper argues that models built for one task generalise poorly to other organs and modalities §1. Its pre-experiment compares three ways of adapting SAM to 3D: a 3D adapter on a frozen SAM, fine-tuning with SAM weights, and training a 3D model from scratch. Each is scored on seen and unseen targets, and training from scratch gives the best average Dice tab. 2 §4.1. The approach therefore depends on having enough labelled 3D data to train from scratch, which is the role of the dataset described under Method.

Method

Architecture. All three parts are 3D §4.1 fig. 3.

• Image encoder. A 3D ViT with a 3D patch embedding, a learned absolute position embedding, and transformer blocks. Most blocks use window attention and four blocks use global attention. A convolutional neck maps the tokens to the prompt embedding width image_encoder3D.py.
• Prompt encoder. Each point is encoded with a random-Fourier 3D position encoding, plus a learned embedding for positive points, negative points, and "no point" padding prompt_encoder3D.py. A previous low-resolution mask, when given, is embedded with three 3D convolutions prompt_encoder3D.py.
• Mask decoder. A two-layer two-way transformer over the image tokens and the prompt tokens. Output tokens are upsampled by two stride-2 transposed convolutions, and each mask token is turned into a dot-product mask with the upsampled features. A small MLP predicts a mask-quality score §4.1 mask_decoder3D.py.

Prompting at evaluation §4.3. The first click is a random foreground voxel. Each later click is a random voxel of the error region, positive where the prediction misses the target and negative where it marks extra tissue. The released inference code follows the same rule infer_utils.py.

Large-volume inference §4.3. For volumes larger than 128 voxels per axis, the paper crops a 128-voxel patch around the first point. If the prediction reaches the patch edge, it extends the prediction with a sliding window that overlaps by 50% and keeps the prompt visible at the edge. The release does not implement this (see Paper vs code).

Data §3. SA-Med3D-140K is built from 70 public and 24 private 3D datasets, which the authors describe as 22K images with 143K masks, 245 categories and 28 modalities (CT, ultrasound and 26 MR sequences) §1 §3.1. The paper lists four cleaning steps: remove masks below 1 cm³ or with any dimension under 1.5 cm; drop masks where background covers over 99% of the volume; keep only the five largest connected components; split symmetric organs into left and right §3.2. Resampling uses 1.5 mm isotropic spacing §3.2. The training split has 131K masks (58 public datasets plus all private data) and the validation split has 12K masks from 16 public datasets §3.3. The authors say 58% of the validation cases come from data sources the training set does not contain §3.3.

Training. Stage 1 uses all 131K training masks for 800 epochs §4.2. Stage 2 fine-tunes on about 75K high-quality masks chosen from the training set §4.2. The paper does not describe the click simulation used during training, and it defers training details to a supplement that is not in the arXiv source. The released training code is described under Hyperparameters and Pitfalls.

The sketches below follow the code. They are not the paper's algorithms, which it does not give. Each step cites the code it follows.

# Interactive inference on one 128^3 ROI  [§4.3]
# Input:  roi_image (1, 1, 128, 128, 128), z-scored
# Input:  roi_gt (1, 1, 128, 128, 128), binary; num_clicks (int)
# Output: binary mask (128, 128, 128)

def segment_roi(model, roi_image, roi_gt, num_clicks):
    # Run the 3D encoder once per volume
    # [code: utils/infer_utils.py#L112]
    image_embeddings = model.image_encoder(roi_image)
    # Start with no clicks and an empty previous mask
    # [code: utils/infer_utils.py#L114-L124]
    points = empty((1, 0, 3))
    labels = empty((1, 0))
    prev_low_res = zeros((1, 1, 32, 32, 32))
    prev_mask = zeros((128, 128, 128))
    for click_index in range(num_clicks):
        # Foreground voxel first, then the error region
        # [code: utils/infer_utils.py#L39-L54]
        new_point, new_label = sample_click(prev_mask, roi_gt)
        # Keep every click so far
        # [code: utils/infer_utils.py#L145-L146]
        points = concat(points, new_point, dim=1)
        labels = concat(labels, new_label, dim=1)
        # Embed the points and the previous low-res logits
        # [code: utils/infer_utils.py#L148-L152]
        sparse, dense = model.prompt_encoder(points, labels, prev_low_res)
        low_res = model.mask_decoder(image_embeddings, sparse, dense)
        prev_low_res = low_res.detach()
        # Probability above 0.5 feeds the next error region
        # [code: utils/infer_utils.py#L164-L168]
        prev_mask = sigmoid(upsample(low_res, 128)) > 0.5
    # Final mask: probability above 0.5
    # [code: utils/infer_utils.py#L172-L179]
    return sigmoid(upsample(low_res, 128)) > 0.5

# One training volume: 11 simulated clicks, loss summed  [§4.2]
# Input:  image (1, 1, 128, 128, 128), z-scored; gt (1, 1, 128, 128, 128)
# Output: scalar loss, summed over the clicks

def training_step(model, image, gt):
    # Encode the volume once, under autocast
    # [code: train.py#L348-L349]
    image_embedding = model.image_encoder(image)
    prev_low_res = zeros((1, 1, 32, 32, 32))
    prev_logits = zeros_like(gt)
    # Mask-only step at a random index in 2..8  [code: train.py#L273]
    insert_at = random_int(2, 8)
    total_loss = 0
    # Eleven clicks, hard-coded  [code: train.py#L359]
    for click_index in range(11):
        # Click from the error region of the raw logits
        # [code: train.py#L247-L248]
        new_point, new_label = sample_error_click(prev_logits, gt)
        if click_index in (insert_at, 10):
            # Mask-only step: the new click is not passed
            # [code: train.py#L277-L282]
            prompt = None
        else:
            # Only the newest click, since multi_click is off
            # [code: train.py#L262-L264]
            prompt = (new_point, new_label)
        sparse, dense = model.prompt_encoder(prompt, prev_low_res)
        low_res = model.mask_decoder(image_embedding, sparse, dense)
        # Raw logits, not detached  [code: train.py#L241-L245]
        prev_logits = upsample(low_res, gt.shape[-3:])
        # Dice plus cross-entropy on logits, summed
        # [code: train.py#L289-L290]
        total_loss += dice_ce_loss(prev_logits, gt)
        prev_low_res = low_res
    return total_loss

# Training click: one random voxel of the error region
# [code: utils/click_method.py#L191-L223]
# Input:  prev_logits (1, 1, D, H, W) raw logits; gt (1, 1, D, H, W)
# Output: one (point, label) pair per volume

def sample_error_click(prev_logits, gt):
    # Predicted foreground: logit above 0.5
    # [code: utils/click_method.py#L199]
    pred = prev_logits > 0.5
    truth = gt > 0
    false_negative = truth and not pred
    false_positive = not truth and pred
    error = false_negative or false_positive
    # Uniform over the error voxels
    # [code: utils/click_method.py#L204-L209]
    point = random_choice(argwhere(error))
    # Positive label for a false negative
    # [code: utils/click_method.py#L211-L214]
    label = 1 if false_negative[point] else 0
    return point, label

Hyperparameters

Rows that cite §... or tab. ... come from the paper. Every other row comes from the released code: its defaults, its shell scripts, or its functions. The released checkpoint's architecture is not stated in the repository (see Pitfalls).

Results

Author-reported throughout. The main evaluation uses the 16 public validation datasets §3.3. Table 3 gives N, the number of slices containing the target, as 10 ≤ N ≤ 200 tab. 3. The paper says the 2D baselines SAM and SAM-Med2D are compared on the same validation set, and that all compared methods use a single model §4.3. Which weights were used for SAM and SAM-Med2D is not stated there. For the task-specific comparison, the authors trained nnU-Net and UNETR on each benchmark and used official weights for the general-purpose models §5.1.2.

Pre-experiment: three ways to adapt SAM to 3D (Dice %, averaged over 1 and 10 points, trained on AMOS and tested on seen and unseen targets) tab. 2:

Overall comparison on 16 public validation datasets (Dice %, N = slices containing the target) tab. 3:

The paper's headline gain is 60.12 points over SAM with one point per volume §5.1.1 §6. τ is the user's interaction time, which the paper says is usually more than one second tab. 3 caption.

Task-specific and general-purpose models on six public benchmarks (Dice %;

• marks a data source or modality not in the training set) tab. 4:

The paper says the 10-point model is best on four of the six benchmarks §5.1.2.

By anatomy and lesion, on the 16-dataset validation set (Dice %) tab. 5:

Transfer as a pre-trained encoder for semantic segmentation (Dice %, UNETR with and without a SAM-Med3D encoder) tab. 6 §5.1.6:

The paper reports a maximum gain of 5.63 Dice points, on AMOS §5.1.6.

Fine-tuning data for promptable segmentation (Dice %). The caption says the table excludes FeTA22 for fairness, so its validation set is not the one used in Table 5 tab. 7:

The last row is the model fine-tuned on high-quality data from 44 datasets tab. 7 caption. Table 7 splits bone into two columns, and its Brain value (61.91) differs from the 10-point Brain value in Table 5 (61.27). Because the validation sets differ, do not read the two tables as one series.

Qualitative results are shown in fig. 5. The text cites fig. 1 and fig. 5 for fewer prompts needed and better consistency between slices §5.2. These are visual claims, not measurements.

Paper vs code

Checked against the code at f3de1fa10da98e46f49f176773d2b1e306ba131f.

• Cleaning rules. The paper removes masks below 1 cm³ or with any dimension under 1.5 cm §3.2. The preparation script drops a mask when its voxel count times the header spacing in mm³ is below 10, which is 0.01 cm³. It has no dimension rule prepare_data_from_nnUNet.py. Its only dataset list is AMOS_val prepare_data_from_nnUNet.py. The other cleaning steps (background over 99%, five largest components, left and right split) are not in the repository. The connected-component code that does exist is SAM's 2D mask generator and is not used for 3D data amg.py.
• Large-volume inference. The paper uses 128-voxel patches with a 50% overlapping sliding window for volumes larger than 128 §4.3. The release crops one 128³ ROI around the ground-truth mask and runs it once infer_utils.py infer_utils.py. No sliding-window code exists in the 3D path.
• Stage 1 length. The paper trains for 800 epochs §4.2. The training default is 200 train.py. The repository has no stage-1 config. The shell scripts fine-tune from the released checkpoint train.sh.
• Stage 2 data. The paper filters about 75K high-quality masks §4.2. The repository has no filter or list of masks. train.py trains on whatever the configured directories contain data_paths.py.
• Training clicks. The paper does not describe the click simulation. The code uses 11 clicks, mask-only steps and a loss summed over clicks train.py.
• Released checkpoint. The README says the released checkpoint is SAM-Med3D-turbo, fine-tuned on 44 datasets readme.md readme.md. The paper's reported models are not stated to be that checkpoint. The inference scripts load it through the external medim package medim_val_single.py.
• Positional encoding. The paper names a 3D positional encoding without a formula §4.1. The code builds the three blocks as sin, cos and sin prompt_encoder3D.py.
• Evaluation clicks. The paper reports 1, 3, 5 and 10 points tab. 3. The released scripts run one click medim_val_single.py medim_val_dataset.py.
• Consistent. The first-click rule §4.3 matches infer_utils.py. The 1.5 mm spacing §3.2 matches prepare_data_from_nnUNet.py and infer_utils.py.

Pitfalls

• Click prompts differ between training and evaluation. Training sends only the newest click unless --multi_click is set, and that flag defaults to off train.py train.py. Evaluation keeps every click so far infer_utils.py. The paper does not say which rule its numbers use.
• Two thresholds for the same mask. Training decides the next click and the Dice on raw logits above 0.5, a probability of about 0.62 click_method.py train.py train.py. Evaluation uses a probability above 0.5, which is a logit above 0 infer_utils.py. Copying one rule into the other gives different click placements.
• The last click never reaches the decoder during training. At the final click, and at one earlier random step, the decoder runs with no point prompt train.py. The point sampled at that step is discarded.
• The sampled click count is ignored. train() draws a count from 1 to 20 each epoch and passes it to train_epoch, which hard-codes 11 clicks train.py train.py.
• Empty error regions crash training. The training click sampler indexes an empty set when prediction equals ground truth click_method.py. Evaluation falls back to a foreground voxel infer_utils.py. Add a guard before porting.
• Gradients flow through all clicks. The previous low-res logits are passed on without detaching train.py. The graph therefore spans every decoder call in one volume. Memory grows with the click count, even though the encoder runs once train.py.
• The IoU head is never trained. batch_forward discards its output train.py, and the loss ignores it. Do not use its scores without training it.
• Two different input normalisations. Sam3D.forward subtracts SAM's pixel mean and divides by SAM's pixel std sam3D.py sam3D.py build_sam3D.py. Those constants have three values, so by reading they cannot broadcast against a one-channel 128³ volume. I did not run it. Training and evaluation skip this step and call the encoder directly on z-scored data train.py infer_utils.py infer_utils.py.
• Z-scoring includes the padding. Statistics come from voxels above zero, but the transform is applied to every voxel. Background and padding become (0 − mean) / std, not zero train.py train.py.
• Window attention attends to zeros. Blocks pad each window with zeros and the attention has no padding mask image_encoder3D.py image_encoder3D.py. For the 128 setting, one 14³ window holds 512 real tokens and 2,232 padding tokens.
• The ROI size is fixed. The absolute position embedding is created for 8³ tokens and cannot be used with another ROI size without resizing image_encoder3D.py.
• The registry default is not the training default. The registry's default entry builds the ViT-H-sized encoder (width 1280, depth 32) build_sam3D.py build_sam3D.py. Training defaults to vit_b_ori train.py. The released checkpoint's architecture is not stated in the repository. Read it from the tensor shapes.
• The shipped inference function may not run as written. sam_model_infer calls the mask decoder without multimask_output infer_utils.py. This repository's MaskDecoder3D.forward requires that argument mask_decoder3D.py. The medim model may use a different class. I did not run the code, so this is from reading only.
• Objects larger than the ROI are cut off. One 128-voxel ROI covers 192 mm at 1.5 mm spacing. The paper's sliding window is not in the release infer_utils.py.
• Inference needs ground truth. Clicks are sampled from the label mask, so the README says a fake label must be supplied for unlabelled images readme.md.
• CT clamping happens only in training. The training loader clamps /ct_ paths to [−1000, 1000] data_loader.py. The inference preprocessing does not clamp infer_utils.py.
• Training skips small labels. A volume with 1000 or fewer label voxels is replaced by another randomly chosen case. The paper does not mention this data_loader.py.
• Multi-class labels overwrite each other. Classes are processed one at a time, and a later class overwrites an earlier one where they overlap infer_utils.py.
• Coordinates are array indices. Prompts are indices in the (D, H, W) order of the canonical ROI array, not NIfTI (x, y, z) click_method.py infer_utils.py. Match the axis order of the volume.
• Empty masks have conventions. Both empty gives NaN, and an empty ground truth with a non-empty prediction gives 0 metric_utils.py. Per-class averages use nanmean metric_utils.py.
• Costs the abstract hides. The timing formulas in tab. 3 add a constant of 2 to 6 seconds on top of τ per click. The main text gives no training hardware or memory figures. Training decodes eleven steps per volume with the graph kept (see above), so memory is higher than one decoder pass.
• The 24 private datasets are not in the release. The paper uses them in training §3.1. The README links a separate SA-Med3D-140K download readme.md. Its licence is not checked here, so check it before reusing any masks.

Lineage

• Builds on: SAM, arXiv:2304.02643, whose promptable design this code extends. The repository's header carries SAM's Meta copyright sam3D.py. The ViT block code is adapted from ViTDet, as its comment says image_encoder3D.py.
• Compared against: SAM-Med2D, arXiv:2308.16184 (fine-tuned SAM for 2D medical images, used as the main general-purpose baseline) §4.3. Also MedLSAM, arXiv:2306.14752; SAM3D, arXiv:2309.03493; MA-SAM, arXiv:2309.08842; MSA, arXiv:2304.12620; 3DSAM-Adapter, arXiv:2306.13465 tab. 1. SegVol, arXiv:2311.13385 tab. 4. Task-specific baselines are UNETR and nnU-Net tab. 4.
• Dataset used in the pre-experiment and as a benchmark: AMOS, arXiv:2206.08023 tab. 2 tab. 4.
• Superseded by: none recorded in this wiki.