Mask3D: Mask Transformer for 3D Semantic Instance Segmentation

In plain words

Instance segmentation in a 3D scan means finding each object (a chair, a window) and marking which points belong to it. Many earlier methods first guess object centres or boxes, then group points with hand-tuned rules.

Mask3D skips the guessing. It keeps a fixed set of 100 slots, one per possible object, and places them at points sampled from the scene. A Transformer lets each slot look at the scene's point features and sharpen into one object: a class label and the set of points that belong to it. Every slot yields a mask over the whole scene, and all masks come out of one pass. During training, each slot is paired with one ground-truth object, so the model is told directly what a good mask looks like.

The point of the approach is that it uses generic Transformer parts and needs fewer hand-picked rules, such as voting targets or clustering radii.

TL;DR

• Transformer-based 3D instance segmentation that predicts binary masks directly, without voting or geometric grouping §1 §3.
• A sparse U-Net gives multi-scale point features. A Transformer decoder refines 100 queries, initialised from farthest-point-sampled positions with zero features, by cross-attention restricted to each query's current mask §3 §3.1.
• Trained with Hungarian matching and a BCE plus Dice mask loss §3.1.
• Author-reported best scores on ScanNet, ScanNet200, S3DIS and STPLS3D tab. 1 tab. 2 tab. 3.
• The released code departs from the paper in several places (STPLS3D backbone and query counts, segment-based evaluation, sampling, top-k selection). See Paper vs code before reimplementing.

Problem

Bottom-up methods learn per-point features and group them into instances. Top-down methods detect boxes first and refine masks inside them. Many current state-of-the-art methods combine centre voting with geometric clustering that depends on hand-selected parameters, and they supervise votes rather than masks §1.

Directly predicting and supervising instance masks, as DETR and Mask2Former do in 2D, had received little attention in 3D §1. The difficulty is that a predicted mask must be paired with a ground-truth mask before a mask loss can be computed. Nearest-neighbour pairing is not optimal and leaves unmatched targets without loss, so the paper uses bipartite graph matching §1 §3.1.

Method

The model has four parts §3:

1. Sparse feature backbone. A sparse convolutional U-Net (Minkowski Res16UNet34C) produces point features at several resolutions. The full-resolution map F_0 gives per-voxel masks. Coarser maps are the attention memory. Each scale is linearly projected to a common width D §3 app. §1. The paper describes the initial voxel feature as the average RGB of the points in the voxel §3; the code keeps one point's values per voxel (see Paper vs code).
2. Mask module. An MLP maps each query X to an instance feature. A dot product with F_0, then a sigmoid and a 0.5 threshold, gives the binary mask B eq. 1 §3. A linear layer with softmax over C+1 classes (the extra one is "no object") gives the class §3. The module runs after every decoder layer, so each layer also yields an auxiliary mask.
3. Query refinement. Each decoder layer does three things in this order: cross-attention to voxel features eq. 2, self-attention between queries, and a feed-forward block. This is the Mask2Former order with cross- and self-attention swapped app. §1 §3. Cross-attention is masked: a query may only attend to voxels inside its mask from the previous layer, implemented by adding minus infinity to the other voxels' logits eq. 3 §3. The paper notes that Mask2Former found this masking helpful §3. Attention runs over the 4 coarsest of the 5 scales, from coarse to fine, and the whole pass repeats 3 times. The decoder weights are shared across those passes, giving L = 12 refinement steps §3 app. §1.
4. Sampled cross-attention. Point clouds differ in size, so voxel features are padded and masked to a fixed size. Above a threshold they are sampled, and the sample is redrawn in each decoder layer during training. All voxels are used at inference §3. The paper calls this a form of dropout and says it is essential for memory and accuracy §3.

Non-parametric queries. Instead of learned query vectors, the positions come from farthest point sampling of the input points, encoded with Fourier features computed from voxel positions scaled to [-1, 1] §3.1 app. §1. The query features start at zero, and only the position is used. The paper also tested sampling point features as the initial query features; that variant was worse §3.1 §4.2. A benefit of this design: the number of queries can change at inference time without retraining §3.1.

Training. The matching cost for predicted instance k and target k̂ is a weighted sum of Dice, BCE and class terms, with weights 2, 5 and 2 eq. 4 §3.1. Matching uses the Hungarian method §3.1. The mask loss on matched pairs is 5 times BCE plus 2 times Dice eq. 5 §3.1. The class loss is the standard cross-entropy with weight 2. Unmatched queries are pushed toward "no object" with an extra weight of 0.1 §3.1. The total loss sums the per-layer losses over all L layers eq. 6 §3.1.

Confidence score. Each instance gets its class probability times the mean mask probability over the voxels inside its binarised mask eq. 7 §3.1.

ScanNet segments. On ScanNet, the dot product is computed on graph-based segments rather than on every voxel, to save memory §3.1. Segments come from Felzenszwalb-Huttenlocher segmentation §3.1.

Post-processing (optional). DBSCAN can split instances that were merged across distant regions. It is applied to each predicted mask on its own, and points marked as noise are dropped app. §1.1 §4.3.

# Mask3D inference on one scene  [§3]  [code: models/mask3d.py#L234-L266]
# Input:  points (N, 6) with xyz and RGB
# Output: up to K instances, each a class label and a binary voxel mask

def predict_instances(points, num_queries):
    # Quantize points to 2 cm voxels for ScanNet  [§3]
    # Default voxel size in the config  [code: conf/data/indoor.yaml#L30]
    voxels = quantize(points, voxel_size=0.02)
    # Run the sparse U-Net to get five feature scales, fine to coarse  [app. §1]
    scales = sparse_unet(voxels)
    # Project every scale to the decoder width D = 128  [app. §1]
    scale_features = [project(scale, dim=128) for scale in scales]
    # Full-resolution features give the per-voxel mask dot products  [§3]
    point_features = scale_features[0]              # (M0, 128)
    # Non-parametric queries: farthest-point-sampled positions  [§3.1]
    centers = farthest_point_sample(voxels, num_queries)
    # Query features start at zero; only the position is encoded  [§3.1]
    queries = zeros(num_queries, 128)
    query_positions = fourier_encode(centers)       # fixed across layers
    # Refine queries: three shared passes over four coarse-to-fine scales  [§3]
    for pass_index in range(3):
        for scale in four_coarsest_to_fine(scale_features):
            # Mask module: dot product of instance features with voxels  [eq. 1]
            masks = mask_logits(queries, point_features)
            # Voxels outside the current mask are excluded from attention  [eq. 3]
            allowed = sigmoid(masks_at_scale(masks, scale)) >= 0.5
            queries = cross_attend(queries, scale, allowed, query_positions)
            # Queries attend to each other  [§3]
            queries = self_attend(queries, query_positions)
            # Feed-forward block  [app. §1]
            queries = feed_forward(queries)
    # Final class and mask predictions from the refined queries  [§3]
    classes = class_head(queries)                   # softmax over C + 1
    masks = mask_logits(queries, point_features) >= 0.5
    return classes, masks

# Training loss with Hungarian matching  [§3.1]
# Input:  predictions for every decoder layer, ground-truth instances
# Output: scalar loss

def training_loss(layer_predictions, gt_masks, gt_classes):
    # Accumulate the loss over all twelve refinement layers  [eq. 6]
    total = 0.0
    for predictions in layer_predictions:
        # One-to-one assignment with cost weights Dice 2, BCE 5, class 2  [eq. 4]
        pairs = hungarian_match(predictions, gt_masks, gt_classes)
        # Mask loss on matched pairs only: 5 * BCE + 2 * Dice  [eq. 5]
        total += 5.0 * bce_loss(predictions.masks, gt_masks, pairs)
        total += 2.0 * dice_loss(predictions.masks, gt_masks, pairs)
        # Class loss: cross-entropy, weight 2; no-object weight 0.1  [§3.1]
        total += 2.0 * class_loss(predictions.classes, gt_classes, pairs)
    return total

Hyperparameters

Results

Author-reported. Setup: ScanNet v2 test scores were accessed on 13 Sept. 2022; validation uses 18 classes; S3DIS uses Area 5 and 6-fold cross-validation; STPLS3D uses its hidden test set; runtime is averaged over the ScanNet validation set on a TITAN X and excludes post-processing tab. 1. Table numbers follow their order of appearance (the IEEE class prints them as Roman numerals).

Analysis experiments, ScanNet validation at 5 cm unless noted §4.2:

Backbones and size, ScanNet validation at 5 cm [supp. §1]:

Model size: the full model has 39.6 M parameters, of which 37.9 M are in the backbone and 1.8 M in the instance part (including the decoder). HAIS and SoftGroup have about 30.9 M in total tab. 6.

DBSCAN post-processing, 2 cm voxels [supp. §1.1]:

Paper vs code

Code: https://github.com/JonasSchult/Mask3D at 11bd5ff (MIT). The repository is linked from the paper and its owner is the first author. I searched only the paper's own links.

• Voxel features. Paper: each voxel gets the average RGB of its points §3. Code: the features come from unique_map, one representative point per voxel, not an average utils.py. I read the indexing only; I did not run ME's quantiser.
• Matching class cost. Paper: class term written as cross-entropy §3.1 eq. 4. Code: the class term is the negative softmax probability of the target class matcher.py. The mask and Dice terms match the paper.
• Loss weights. Paper weights are 2, 5, 2 §3.1. Code: the loss weights are the matcher's cost values, applied to each of the 12 intermediate heads and the final head trainer.py. Consistent with the paper.
• Queries. Paper: 100 queries in all experiments app. §1. Release scripts: 150 for ScanNet validation and ScanNet200, 160 for STPLS3D, 100 for S3DIS scannet_val.sh scannet200_val.sh stpls3d_val.sh s3dis_from_scratch.sh.
• STPLS3D backbone and voxels. Paper: Res16UNet34C for all experiments §3.1 app. §1. Release script: Res16UNet18B at 0.333 voxels, with 54 m crops stpls3d_val.sh stpls3d_val.sh. The paper does not state a voxel size for STPLS3D.
• DBSCAN ε for ScanNet. Paper: 0.9 app. §1.1. Code default and ScanNet validation script: 0.95 config_base_instance_segmentation.yaml scannet_val.sh.
• Top-k selection. Paper: the confidence is computed for queries whose dominant class is chosen §3.1. Code: the top-k is taken over all query-class pairs (query × class), so one query can yield several instances; with k = -1 it keeps num_queries pairs trainer.py. The paper describes no top-k step.
• Segment-based evaluation. Paper: segments are used during training §3.1. Code: when eval_on_segments is set, masks are averaged over segments and thresholded at 0.5 trainer.py. The release scripts enable this scannet_val.sh scannet_val.sh.
• Sampling thresholds. The paper gives no thresholds. Code uses per-scale sizes 200 to 51200 and samples only in training mask3d.yaml mask3d.py. At evaluation, is_eval=True disables sampling trainer.py.
• Empty attention masks. Not in the paper. If every voxel is masked out for a query, the code unmasks all voxels for that query mask3d.py.
• S3DIS fine-tuning. The paper reports the pretrain-then-fine-tune gain (1.2 mAP on Area 5) §4.1 but no fine-tuning learning rate. The release script uses 1e-5 s3dis_pretrained.sh.
• Epochs. Paper: 600 §3.1. Config: 601 trainer600.yaml. S3DIS: 1000 in the paper app. §1, 1001 in the script s3dis_from_scratch.sh.
• Result numbers that disagree. The same ScanNet validation setting appears with 40.6 in Table 4 and 40.9 in Table 5 tab. 4 tab. 5. The main validation mAP is 55.2 at 2 cm tab. 1, while the DBSCAN baseline is 54.3 AP at 2 cm tab. 7. The paper does not explain either gap.

Pitfalls

• Mixed baselines in the headline numbers. The abstract's +10.1 mAP on S3DIS 6-fold compares non-pretrained Mask3D (64.5) with pretrained SoftGroup (54.4) tab. 2. Pretrained against pretrained the gap is 61.8 − 54.4 = 7.4. The text gives +6.2 on S3DIS for Area 5, which compares pretrained Mask3D (57.8) with pretrained SoftGroup (51.6).
• ScanNet200 gain. The abstract says +12.4 mAP and the text +10.8 §4.1. Table 3 has no overall mAP. The 10.8 equals the head-class gap to LGround (38.3 − 27.5) tab. 3. I could not reproduce the abstract figure from the tables.
• STPLS3D gain. The text gives +11.2 mAP §4.1; the table gives 57.3 − 46.2 = 11.1 tab. 3.
• Runtime excludes post-processing. The runtime in Table 1 leaves out post-processing tab. 1. The code runs DBSCAN per query on the CPU when enabled trainer.py. Its cost is not reported in the paper.
• Sampling is stochastic and unquantified. The paper calls sampled cross-attention crucial, but gives no ablation §3. Training samples differ from layer to layer, and inference uses all voxels, so the train-test distributions differ.
• Queries must be retuned with care. Changing the query count at test time is free in the non-parametric design. The paper reports that fewer queries lower accuracy and more queries add less than 1 mAP §4.2; the numbers are in a figure it does not tabulate.
• Memory. One 2 cm training run takes about 78 hours on an A40 §3.1. The voxel resolution is limited by memory: the final query count is 100 because of it §4.2.
• Train-test context. S3DIS is trained on random 6 m blocks and tested on full scenes app. §1. STPLS3D is evaluated on 50 m blocks, fed as 54 m blocks app. §1.
• Merged distant instances. Far-apart objects with similar appearance can merge; the paper names this as a systematic failure §4.3 and proposes DBSCAN as a fix app. §1.1.
• Single-point scenes fail. A sampled set with a single point raises an error in the attention step mask3d.py.
• Thresholds in two places. The 0.5 threshold sets both the output masks and the attention masks eq. 1 eq. 3. Changing it changes training, not only output.

Lineage

• Builds on: DETR for learned queries and bipartite matching §1 §3.1. Mask2Former (Cheng et al., CVPR 2022) for the mask module and the masked cross-attention §3 app. §1; its arXiv id is not in this wiki, so it is not linked. 3DETR (Misra et al., ICCV 2021) for non-parametric queries §3.1; not in this wiki.
• Compared against: SoftGroup (Vu et al., CVPR 2022), HAIS (Chen et al., ICCV 2021), 3D-MPA (Engelmann et al., CVPR 2020), PointGroup, SSTNet, DyCo3D, LGround and others tab. 1 tab. 2 tab. 3. None are linked in this wiki.
• Superseded by: not known from content in this wiki.