U-Mamba2: 3D CBCT Dental Segmentation

• U-Mamba2 is a 3D medical image segmentation network designed for dental CBCT, integrating CNN with Mamba2 state-space modeling for global context.
• Its U-Net encoder–decoder architecture combines interactive click prompts and self-supervised pretraining to enhance segmentation of complex dental structures.
• The system achieves superior Dice and HD95 metrics, and its semi-supervised extension refines tooth and pulp segmentation using minimal labeled data.

U-Mamba2 is a 3D medical image segmentation network for cone-beam computed tomography (CBCT) dental anatomy, designed specifically for the ToothFairy3 challenge. It is a hybrid CNN + Mamba2 state space model architecture that keeps the familiar U-Net encoder–decoder structure but replaces the bottleneck’s global-context modeling with Mamba2 and structured state-space duality (SSD) to improve efficiency while retaining strong performance. In the reported system, the architecture is combined with interactive click prompts for interactive nerve segmentation, self-supervised pretraining, and dental domain knowledge; a later framework, U-Mamba2-SSL, uses U-Mamba2 as the backbone for multi-stage semi-supervised tooth and pulp segmentation in CBCT (Tan et al., 15 Sep 2025, Tan et al., 24 Sep 2025).

1. Problem setting and application domain

U-Mamba2 is developed for dental CBCT segmentation, where accurate delineation of teeth, jaw, sinuses, nerves, pulps, and foramina is clinically important for diagnosis, surgical planning, endodontics, implant planning, and orthodontic assessment. The motivating difficulty is that manual segmentation is slow and requires expertise, especially for small, thin, low-contrast structures such as incisive nerves and lingual foramen (Tan et al., 15 Sep 2025).

The primary benchmark context is ToothFairy3. Task 1 is multi-anatomy CBCT segmentation and extends ToothFairy2 by including the existing 42 anatomy classes plus pulps, incisive nerves, and lingual foramen, for a total of 46 anatomy classes. The dataset contains 532 CBCT scans, with scan shapes ranging from $(170, 272, 345)$ to $(298, 512, 512)$. Task 2 is interactive segmentation of the inferior alveolar nerves, where user clicks are provided as prompts (Tan et al., 15 Sep 2025).

A related but distinct setting appears in U-Mamba2-SSL, which addresses 3D tooth and pulp segmentation in CBCT under semi-supervised learning. There, the task is framed around a small labeled set and a much larger unlabeled set, reflecting the fact that dense manual annotation of CBCT volumes is expensive, slow, and requires specialized dental expertise (Tan et al., 24 Sep 2025).

2. Core architecture and the U-Mamba2 block

U-Mamba2 follows a standard encoder–decoder U-Net-like architecture with residual connections across scales. The encoder uses repeated residual blocks and strided downsampling convolutions; the decoder uses residual blocks and transposed convolutions for upsampling; skip connections link encoder and decoder features at matching scales; and the bottleneck contains the U-Mamba2 block with the Mamba2 state-space module. An optional interactive branch is added for Task 2 (Tan et al., 15 Sep 2025).

The defining architectural choice is the placement of Mamba2 in the bottleneck. The paper states that applying the U-Mamba2 block only at the bottleneck gave the best empirical performance for 3D CT, matching prior observations. A common shorthand is to treat U-Mamba2 as an all-SSM network; the reported design is instead a hybrid CNN + Mamba2 system in which convolution remains responsible for multiscale local feature extraction while Mamba2 is introduced at the lowest-resolution global-context stage. This suggests a division of labor in which local structure is handled convolutionally and long-range dependency modeling is concentrated where sequence length is smallest (Tan et al., 15 Sep 2025).

The bottleneck transformation is described explicitly. A feature map of shape

$(B, C, H, W, D)$

is reshaped and transposed into a sequence

$(B, T, C), \qquad T = H \times W \times D.$

Layer Normalization is applied, the sequence is passed through Mamba2, and the output is reshaped back to spatial form. The paper emphasizes that Mamba2 scales linearly with sequence length, uses the SSD framework, imposes stronger structural constraints on hidden-state dynamics, replaces selective scan with matrix multiplication for improved parallelism, and gives speedups without losing performance (Tan et al., 15 Sep 2025).

In the semi-supervised extension, U-Mamba2 remains the segmentation backbone and is again described as a U-Net-like encoder–decoder architecture where Mamba2 state space models are integrated into the bottleneck region. The paper explicitly states that this design improves the ability to capture long-range dependencies while remaining efficient, and positions it against transformer-based alternatives as offering better efficiency without sacrificing performance (Tan et al., 24 Sep 2025).

3. Interactive prompting, self-supervised pretraining, and dental priors

For ToothFairy3 Task 2, U-Mamba2 supports interactive segmentation of the inferior alveolar nerves. Click prompts are represented by $X$, $Y$, $Z$ coordinates, class labels, and a variable number of clicks $N$. These click points are encoded with a learnable position embedding based on spatial position and class label, fused with image features through two-way cross-attention blocks, repeated twice, and followed by Layer Normalization. The paper states that this design follows the SAM2 framework. In effect, click prompts are fused with the Mamba2 output features so that segmentation can be refined by user guidance (Tan et al., 15 Sep 2025).

Because labeled CBCT data are limited, the model is pretrained with self-supervised learning using the Disruptive Autoencoder (DAE) framework on 371 unlabeled CBCT scans from STS-3D-Tooth. The input volume is corrupted using local masking, downsampling, and Gaussian noise, and the network learns to reconstruct the original volume using

$\mathcal{L}_{SSL} = \|x - \hat{x}\|_1.$

The pretrained weights initialize U-Mamba2 except for the optional interactive branch and the final segmentation layer (Tan et al., 15 Sep 2025).

A distinctive component of the reported system is the incorporation of dental domain knowledge. First, the model uses label smoothing for related anatomies: for a voxel of class $k$, the target probability for class $(298, 512, 512)$0 is set to $(298, 512, 512)$1, and the remaining $(298, 512, 512)$2 is distributed across related classes. Second, very small classes such as the left incisive nerve, right incisive nerve, and lingual foramen are assigned a loss weight of 10 so that they are not overwhelmed by large structures. Third, left-right mirroring augmentation is made anatomically safe by swapping labels of left/right counterpart classes when an image is mirrored and also swapping the corresponding output logits during test-time augmentation; the paper states that this increases possible mirroring axis combinations from 3 to 7. Fourth, post-processing removes small false-positive blobs by computing connected components and selecting a threshold based on the 0.5th percentile of connected-component volumes from ground truth (Tan et al., 15 Sep 2025).

These additions are not peripheral. The paper presents them as part of the model design rather than as external heuristics, and the ablation study shows that each dental prior helps a little, with the combined configuration giving the best overall segmentation and the strongest gains on tiny structures (Tan et al., 15 Sep 2025).

4. Training configuration, losses, and evaluation protocol

U-Mamba2 is implemented in nnU-Net, with a 9:1 stratified train-validation split that preserves source proportions across different scanners and fields of view. The common training settings are 7 encoder-decoder stages, input patch size $(298, 512, 512)$3, batch size 1, self-supervised pretraining with DAE, and RTX 4090 hardware. Inference uses sliding-window inference, with tile size set to 0.5 initially for validation comparison and left-right mirroring disabled in test-time augmentation for fair benchmarking on validation. Later inference is optimized with tile size 0.9 and mirror axes anterior/posterior plus left/right. For the final submission, the model is trained for 1500 epochs, uses all available data, increases batch size to 2, increases input patch size to $(298, 512, 512)$4, and performs inference with tile size 0.9 and mirroring in axes 1 and 2 (Tan et al., 15 Sep 2025).

The supervised segmentation objective is the combination of cross-entropy loss and Dice loss,

$(298, 512, 512)$5

Evaluation uses Dice coefficient, HD95, and average inference time in seconds (Tan et al., 15 Sep 2025).

The semi-supervised extension formalizes the data setting as

$(298, 512, 512)$6

with $(298, 512, 512)$7. It uses nnU-Net, patch-size training, and sliding-window inference, and preprocesses the images by resizing to the median voxel spacing $(298, 512, 512)$8, giving a median input size of approximately $(298, 512, 512)$9. The data are clipped to the 0.5th and 99.5th percentiles and normalized by mean and standard deviation. Training runs for 500 epochs with SGD, momentum $(B, C, H, W, D)$0, initial learning rate $(B, C, H, W, D)$1, polynomial learning-rate decay, batch size 2, patch size $(B, C, H, W, D)$2, model size 156M parameters, and compute 6.22T FLOPs. The 30 labeled samples are split into 20 training and 10 internal validation cases (Tan et al., 24 Sep 2025).

5. Quantitative performance and ablation results

On ToothFairy3 validation, U-Mamba2 reports the best raw Dice and HD95 on both tasks among the compared models. For Task 1, it achieves Dice 0.873 and HD95 41.08, improving after post-processing to Dice 0.908 and HD95 21.35, with inference time 6.81 s. For Task 2, it achieves Dice 0.905 and HD95 1.63, improving after post-processing to Dice 0.913 and HD95 1.57, with inference time 5.70 s. The corresponding U-Mamba results are 0.865/42.06 and 0.896/25.88 with 6.98 s for Task 1, and 0.903/1.65 and 0.913/1.58 with 5.88 s for Task 2. The paper therefore states that U-Mamba2 gives the best raw Dice and HD95 on both tasks and is slightly faster than U-Mamba (Tan et al., 15 Sep 2025).

The comparison set also includes SwinUNETR and nnU-Net ResE. SwinUNETR reports Task 1 Dice 0.858 and HD95 48.86, improving after post-processing to 0.874 and 40.09 with time 7.23 s. nnU-Net ResE reports Task 1 Dice 0.861 and HD95 45.28, improving after post-processing to 0.887 and 32.05 with time 6.20 s, and Task 2 Dice 0.901 and HD95 1.98, improving after post-processing to 0.905 and 1.71 with time 5.06 s (Tan et al., 15 Sep 2025).

The ablation study on dental domain knowledge excludes post-processing and shows incremental gains from each prior. The baseline configuration reports Dice 0.867, HD95 42.36, ILN Dice 0.617, and ILN HD95 38.41. Label smoothing alone gives 0.872, 40.74, 0.628, and 38.15. Weighted loss alone gives 0.870, 41.31, 0.635, and 37.99. Left-right mirroring alone gives 0.871, 41.20, 0.642, and 36.48. The full combination of label smoothing, weighted loss, and left-right mirroring gives the best overall result: Dice 0.873, HD95 41.08, ILN Dice 0.646, and ILN HD95 35.21. The paper interprets the largest benefits for the tiny structures as coming from weighted loss and left-right mirroring (Tan et al., 15 Sep 2025).

The speed/accuracy trade-off study reports that increasing tile size to 0.9 reduces inference time by 12.9%, while Dice drops only by 0.002. The best TTA mirror-axes combination is axes 1 and 2, corresponding to anterior/posterior and left/right (Tan et al., 15 Sep 2025).

For the final hidden test submission, the reported results are mean Dice 0.792 and HD95 93.19 on Task 1, and mean Dice 0.852 and HD95 7.39 on Task 2, securing top 3 places in both tasks. The abstract also mentions an average inference time of XX, marked as TBC during the ODIN workshop, for Task 1 (Tan et al., 15 Sep 2025).

6. U-Mamba2-SSL: semi-supervised extension, reported gains, and limitations

U-Mamba2-SSL is a multi-stage semi-supervised learning framework for 3D tooth and pulp segmentation in CBCT that builds on the U-Mamba2 model and employs a staged training strategy. The framework first pre-trains U-Mamba2 in a self-supervised manner using a disruptive autoencoder, then leverages unlabeled data through consistency regularization, and finally applies pseudo-labeling with a reduced loss weighting. U-Mamba2 is the central backbone used in all stages, and each stage’s final checkpoint initializes the next stage (Tan et al., 24 Sep 2025).

In stage 1, U-Mamba2 is pretrained on all available data $(B, C, H, W, D)$3 using a disruptive autoencoder with three corruption/reconstruction tasks: denoising by adding random Gaussian noise, super-resolution by downsampling with linear interpolation and reconstructing the high-resolution version, and mask recovery by masking random cubical regions and reconstructing the missing content. The model is trained to reconstruct the original input using an $(B, C, H, W, D)$4 loss (Tan et al., 24 Sep 2025).

In stage 2, the framework uses both labeled and unlabeled data. For labeled samples, the supervised segmentation objective $(B, C, H, W, D)$5 is described as a combination of Dice loss and cross-entropy loss. For unlabeled data $(B, C, H, W, D)$6, the model produces an unperturbed prediction $(B, C, H, W, D)$7 and a perturbed prediction $(B, C, H, W, D)$8, and defines the consistency loss as

$(B, C, H, W, D)$9

The overall stage-2 loss is

$(B, T, C), \qquad T = H \times W \times D.$0

where $(B, T, C), \qquad T = H \times W \times D.$1 ramps up exponentially from zero to a fixed value $(B, T, C), \qquad T = H \times W \times D.$2 by epoch $(B, T, C), \qquad T = H \times W \times D.$3. Input perturbations include median filter, Gaussian blur, Gaussian noise, random brightness, random contrast, low-resolution simulation, and sharpening filter. Feature perturbations include Random Spatial Dropout with probability 0.5, Random Activation Dropout with

$(B, T, C), \qquad T = H \times W \times D.$4

and noise injection with

$(B, T, C), \qquad T = H \times W \times D.$5

The proportion of unlabeled data in each epoch is increased from 10% to 50% by epoch $(B, T, C), \qquad T = H \times W \times D.$6 (Tan et al., 24 Sep 2025).

In stage 3, pseudo-labeling is added. If the predicted voxel confidence exceeds $(B, T, C), \qquad T = H \times W \times D.$7, the predicted class is treated as ground truth; otherwise the voxel is assigned to background and ignored in the pseudo-label loss. The loss becomes

$(B, T, C), \qquad T = H \times W \times D.$8

and the unlabeled ratio is linearly increased from 30% to 50% by epoch $(B, T, C), \qquad T = H \times W \times D.$9 (Tan et al., 24 Sep 2025).

The reported validation results show that U-Mamba2 already outperforms nnU-Net, especially in Identification Accuracy. nnU-Net reports DSC 0.963, NSD 0.997, mIoU 0.928, IA 0.286, and average 0.794. U-Mamba2 reports DSC 0.965, NSD 0.998, mIoU 0.930, IA 0.464, and average 0.839. U-Mamba2-SSL with pretraining only reports DSC 0.967, NSD 0.998, mIoU 0.937, IA 0.731, and average 0.908; with pretraining + CR it reports DSC 0.967, NSD 0.999, mIoU 0.935, IA 0.736, and average 0.910; and with pretraining + CR + PL it reports DSC 0.967, NSD 0.999, mIoU 0.935, IA 0.738, and average 0.910. The paper highlights that pretraining gives the biggest gain, especially in IA, which rises from 0.464 to 0.731, while consistency regularization and pseudo-labeling further improve it slightly (Tan et al., 24 Sep 2025).

The reporting of final performance includes several numerically distinct summaries. The abstract reports an average score of 0.872 and DSC 0.969. The results table reports 0.910 average on the validation split used in the experiment section. The conclusion says the method achieved 1st place with an average score of 0.928 on the STSR 2025 validation set. In the final challenge submission, using larger patches and longer training, the reported validation scores are DSC 0.969, NSD 0.998, mIoU 0.940, and IA 0.806. The paper also notes qualitative failure cases, including inaccurate estimation of pulp thickness or length and more false positives near image edges in limited field-of-view CBCT scans. Its stated limitations are that the dataset contains both full-FOV and LFOV CBCTs, which have different properties, and that the foreground occupies only a small ROI, suggesting future work on tailored augmentation, processing, and reduction of wasted computation on background-only volumes (Tan et al., 24 Sep 2025).