U-Mamba2-SSL: Semi-supervised 3D Tooth Segmentation

• U-Mamba2-SSL is a semi-supervised segmentation framework that integrates disruptive autoencoder pretraining, consistency regularization, and pseudo-labeling to delineate tooth and pulp in CBCT scans.
• It employs a U-Net style encoder–decoder augmented with Mamba2 state-space modules, achieving high metrics including DSC 0.969, NSD 0.998, mIoU 0.940, and IA 0.806.
• The framework effectively utilizes abundant unlabeled data to reduce manual annotation burden, enhancing diagnostic and treatment planning in dental imaging.

U-Mamba2-SSL is a semi-supervised 3D segmentation framework for tooth and pulp delineation in Cone-Beam Computed Tomography (CBCT), introduced for settings in which voxel-level annotation is scarce but unlabeled scans are abundant (Tan et al., 24 Sep 2025). The method builds on the U-Mamba2 backbone—a U-Net style encoder–decoder augmented with Mamba2 state-space modules at the bottleneck—and organizes training into three sequential stages: self-supervised pretraining with a disruptive autoencoder, consistency regularization on unlabeled data, and pseudo-labeling with reduced loss weighting. In the reported STSR 2025 Task 1 validation setting, the framework achieved an “average score” of 0.872 and a Dice similarity coefficient (DSC) of 0.969; the final challenge submission reports DSC 0.969, NSD 0.998, mIoU 0.940, and IA 0.806 on the validation set (Tan et al., 24 Sep 2025).

1. Problem formulation and clinical context

The framework addresses semi-supervised segmentation of teeth and pulp in CBCT volumes, a task motivated by the role of CBCT in diagnosis, orthodontics, treatment planning, and surgery planning. The central difficulty is that manual voxel-level annotation is labor-intensive, requires domain expertise, and scales poorly because of large voxel counts, anatomical variability across patients, and CBCT-specific artifacts and noise (Tan et al., 24 Sep 2025).

Formally, the labeled and unlabeled datasets are $$\mathcal{D}_l=\{(x_{1}^l, y_{1}), ..., (x_{n}^l, y_{n})\}$$ and $\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$, where $x_i^l \in \mathbb{R}^{H \times W \times D}$ is a labeled input volume, $y_i \in \mathbb{R}^{C \times H \times W \times D}$ is its voxel-level label, and $x_i^u$ is an unlabeled volume. The objective is to exploit the regime $m \gg n$ to train a 3D segmentation model that benefits from unlabeled CBCT scans rather than relying only on the limited labeled subset (Tan et al., 24 Sep 2025).

This problem framing is important because U-Mamba2-SSL is not presented as a purely self-supervised pretraining method. A common simplification is to view it only as a reconstruction-based initializer, but the reported method is a full semi-supervised pipeline in which unlabeled data participate in all three stages: pretraining, consistency regularization, and pseudo-labeling (Tan et al., 24 Sep 2025).

2. Architectural substrate: U-Mamba2 and bottleneck state-space modeling

U-Mamba2-SSL inherits its core network from U-Mamba2. The backbone is a U-Net style 3D encoder–decoder with skip connections, in which the encoder progressively downsamples and extracts multi-scale feature maps and the decoder upsamples while fusing encoder features. The key architectural distinction is the insertion of Mamba2 state-space modules at the bottleneck, where they are used to enlarge the effective receptive field and capture long-range dependencies and global context without quadratic-cost attention (Tan et al., 24 Sep 2025).

In the U-Mamba2-SSL configuration, the network has seven encoder–decoder stages, 156M parameters, and 6.22T FLOPs for the configured patch size. Although the paper does not elaborate normalization layers or activation functions, implementation follows nnU-Net conventions for 3D medical segmentation (Tan et al., 24 Sep 2025).

Relative to the original U-Mamba, which used Mamba rather than Mamba2, the Mamba2 version is described as imposing stronger constraints on the hidden state-space structure, yielding higher efficiency while maintaining accuracy relative to transformer-based alternatives (Tan et al., 24 Sep 2025). The broader U-Mamba2 architecture paper describes the same design principle as a hybrid CNN–SSM system in which a single Mamba2 block is placed at the bottleneck and 3D bottleneck features are flattened to a sequence, processed by Mamba2, and reshaped back to volumetric form (Tan et al., 15 Sep 2025). This architectural lineage indicates that U-Mamba2-SSL should be understood as a semi-supervised training framework layered on top of a specific state-space segmentation backbone rather than as a standalone model family.

3. Three-stage semi-supervised learning pipeline

The training framework proceeds in three stages, each initialized from the final checkpoint of the previous stage, and is implemented in nnU-Net with patch-size training and sliding-window inference (Tan et al., 24 Sep 2025).

The first stage is self-supervised pretraining via a disruptive autoencoder. Pretraining uses all available data, $\mathcal{D}_l \cup \mathcal{D}_u$, and applies three forms of disruption to the input volume: denoising with additive Gaussian noise, super-resolution through downsampling via linear interpolation followed by reconstruction, and masking of random cubical regions that are small relative to the input dimensions. After corruption, U-Mamba2 reconstructs the original volume with an $L_1$ objective,

$L_{\mathrm{rec}} = \lVert \hat{x} - x \rVert_1 .$

The intended effect is dual: denoising and super-resolution encourage recovery of fine local details such as edges and textures, while masking forces the model to infer missing content from larger-scale context (Tan et al., 24 Sep 2025).

The second stage applies consistency regularization. Labeled patches are optimized with a supervised segmentation loss $\mathcal{L}_S$, while unlabeled patches are passed through paired unperturbed and perturbed pathways. The consistency loss is the $\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$0 distance between the two predictions,

$\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$1

and the total objective is

$\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$2

The method explicitly follows the smoothness assumption and avoids non-local spatial transforms such as flips and rotations in the consistency branch because they would violate local correspondence in dense segmentation (Tan et al., 24 Sep 2025).

The third stage adds pseudo-labeling. A consistency-trained model generates voxel-wise pseudo-labels for unlabeled scans; if the predicted class confidence exceeds $\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$3, the predicted class is accepted as pseudo ground truth, otherwise the voxel is set to background and ignored in the loss. The objective becomes

$\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$4

where $\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$5 reduces the influence of pseudo-labeled supervision in order to limit error propagation (Tan et al., 24 Sep 2025).

This staged design makes the framework structurally conservative. Rather than trusting pseudo-labels from the outset, it first learns reconstruction priors, then enforces prediction stability under perturbation, and only then incorporates confidence-gated self-training. A plausible implication is that the authors intended the three stages to address different failure modes of low-label 3D segmentation: representation quality, perturbation sensitivity, and unlabeled target exploitation.

4. Perturbation design, supervised objectives, and data protocol

The perturbation scheme in Stage 2 is unusually explicit. For the unlabeled consistency branch, input perturbations include median filtering, Gaussian blur, Gaussian noise, random brightness, random contrast, low-resolution simulation, and image sharpening. Feature perturbations are applied on encoder feature maps before skip-connection fusion and include Random Spatial Dropout with probability 0.5, Random Activation Dropout with threshold $\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$6 that zeros activations above the sampled percentile, and Noise Injection with $\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$7 applied as $\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$8 for feature map $\mathcal{D}_{u}=\{x_1^u, ..., x_m^u\}$9 (Tan et al., 24 Sep 2025).

For labeled data, supervised segmentation uses a combination of Dice loss and voxel-wise cross-entropy. The cross-entropy is

$x_i^l \in \mathbb{R}^{H \times W \times D}$0

and the soft Dice coefficient is

$x_i^l \in \mathbb{R}^{H \times W \times D}$1

In pseudo-labeling, the reduced-weight objective is written as

$x_i^l \in \mathbb{R}^{H \times W \times D}$2

with $x_i^l \in \mathbb{R}^{H \times W \times D}$3 and confidence gating at $x_i^l \in \mathbb{R}^{H \times W \times D}$4 (Tan et al., 24 Sep 2025).

The dataset and preprocessing protocol are also tightly specified. The labeled set comprises 30 CBCT scans from the STSR 2025 Task 1 Challenge, split into 20 training and 10 internal validation cases. Unlabeled data consist of all remaining provided scans and are used in pretraining and semi-/pseudo-labeled stages. Inputs are resampled to the median voxel spacing $x_i^l \in \mathbb{R}^{H \times W \times D}$5, producing a median input size of $x_i^l \in \mathbb{R}^{H \times W \times D}$6. Intensities are clipped to the 0.5th and 99.5th percentiles and normalized by mean and standard deviation. Random crops are sampled at patch size $x_i^l \in \mathbb{R}^{H \times W \times D}$7, with at least 33% of samples containing foreground (Tan et al., 24 Sep 2025).

General augmentation uses rotation, scaling, Gaussian noise and blur, brightness and contrast changes, low-resolution simulation, and mirroring. The unlabeled consistency branch is stricter: only intensity, blur, noise, and sharpening transforms are used, with no spatial transforms, in order to preserve local alignment between perturbed and unperturbed predictions (Tan et al., 24 Sep 2025).

5. Training regime, implementation, and computational profile

Training is conducted in nnU-Net 2.6.2 with PyTorch 2.7.1. The optimizer is SGD with momentum 0.99, the initial learning rate is 0.01 with polynomial decay, batch size is 2, and the standard training configuration runs for 500 epochs; the final submission is trained for 1000 epochs. The reported patch size is $x_i^l \in \mathbb{R}^{H \times W \times D}$8, while the final submission uses $x_i^l \in \mathbb{R}^{H \times W \times D}$9 (Tan et al., 24 Sep 2025).

The semi-supervised loss weights and thresholds are explicitly reported: $y_i \in \mathbb{R}^{C \times H \times W \times D}$0, $y_i \in \mathbb{R}^{C \times H \times W \times D}$1, and $y_i \in \mathbb{R}^{C \times H \times W \times D}$2. The consistency weight $y_i \in \mathbb{R}^{C \times H \times W \times D}$3 ramps up exponentially from 0 to $y_i \in \mathbb{R}^{C \times H \times W \times D}$4 by epoch $y_i \in \mathbb{R}^{C \times H \times W \times D}$5, and the fraction of unlabeled data per epoch increases linearly from 10% to 50% by epoch $y_i \in \mathbb{R}^{C \times H \times W \times D}$6 in Stage 2. In Stage 3, the unlabeled proportion increases linearly from 30% to 50% by epoch $y_i \in \mathbb{R}^{C \times H \times W \times D}$7 (Tan et al., 24 Sep 2025).

Hardware and runtime are reported for a practical deployment setting: NVIDIA RTX 5090 32 GB, Ubuntu 24.04, CUDA 12.9, and approximately 13 hours of training for the 500-epoch configuration. Inference uses sliding windows with test-time augmentation. A speed-optimized configuration with tile size 0.9 and mirroring on axes “1,2” yields an average score drop of only 0.002 and an inference time of 17.08 seconds (Tan et al., 24 Sep 2025).

The method is therefore positioned within the nnU-Net engineering ecosystem rather than as a bespoke training stack. This matters because many implementation choices that are not restated in the paper are inherited from nnU-Net conventions, while the novel contributions are concentrated in the bottleneck state-space backbone and the three-stage semi-supervised procedure.

6. Empirical results, ablations, and interpretation

On the STSR 2025 Task 1 validation set, the abstract reports an average score of 0.872 and DSC 0.969. The final challenge submission, using scaled training and a larger patch, reports DSC 0.969, NSD 0.998, mIoU 0.940, and IA 0.806. In this challenge, the “average score” is defined as the arithmetic mean of DSC, NSD, mIoU, and IA (Tan et al., 24 Sep 2025).

The reported comparison and ablation results isolate the contribution of the backbone and each semi-supervised stage. The nnU-Net baseline attains DSC 0.963, NSD 0.997, mIoU 0.928, IA 0.286, and Average 0.794. Replacing the baseline with U-Mamba2 improves the same metrics to DSC 0.965, NSD 0.998, mIoU 0.930, IA 0.464, and Average 0.839. Adding disruptive autoencoder pretraining yields DSC 0.967, NSD 0.998, mIoU 0.937, IA 0.731, and Average 0.908. Adding consistency regularization gives DSC 0.967, NSD 0.999, mIoU 0.935, IA 0.736, and Average 0.910. Adding pseudo-labeling results in DSC 0.967, NSD 0.999, mIoU 0.935, IA 0.738, and Average 0.910 (Tan et al., 24 Sep 2025).

Several observations are explicitly stated. Integrating Mamba2 into U-Net already improves IA over nnU-Net. DAE pretraining produces the largest IA jump, from 0.464 to 0.731, which the paper interprets as evidence that self-supervised reconstruction on unlabeled CBCT volumes is highly beneficial. Consistency regularization and pseudo-labeling then provide further IA gains while maintaining high DSC, NSD, and mIoU. The reported configuration ranked first on the STSR 2025 validation set, and the paper later notes an “average score of 0.928” under a scaled training setup (Tan et al., 24 Sep 2025).

These results also clarify what U-Mamba2-SSL is optimizing beyond overlap. The modest differences in DSC across later ablations, contrasted with larger IA changes, suggest that the framework was especially effective at improving the challenge-specific metric balance rather than merely maximizing Dice. This suggests that the semi-supervised stages contribute not only to voxel overlap but also to boundary or instance-related properties captured by the broader metric suite.

7. Limitations, failure modes, and relation to adjacent work

The paper reports several failure cases. Thickness and length or extent of pulp can be under-estimated or over-estimated, and limited field-of-view (LFOV) CBCT scans tend to produce more false positives near image boundaries. The authors also identify dataset heterogeneity between full and LFOV CBCT scans as a remaining issue, and propose that future work should specialize preprocessing and augmentation per CBCT type. Because only a small portion of the full volume contains tooth and pulp foreground, the paper further suggests ROI detection or cropping as a way to reduce computation and focus modeling capacity on relevant anatomy (Tan et al., 24 Sep 2025).

The broader U-Mamba2 architecture paper adds adjacent context from another dental benchmark. There, DAE-style self-supervised pretraining is also used on unlabeled CBCT, but the isolated quantitative contribution of SSL cannot be measured because all compared models are SSL-pretrained; that work instead emphasizes multi-anatomy segmentation, dental domain knowledge, and interactive prompting (Tan et al., 15 Sep 2025). This contrast is useful: in U-Mamba2-SSL for STSR 2025, the three-stage design is directly ablated, whereas in the broader U-Mamba2 study SSL is part of the standard training recipe rather than the experimental variable.

Outside dental imaging, Mamba-based self-supervised learning has also been studied for speech, where replacing Transformer blocks in HuBERT with Mamba preserves the masked-unit prediction paradigm while enabling linear-time long-sequence processing; that work does not implement Mamba2, but it provides a separate line of evidence that state-space backbones can be effective substrates for SSL in long-context settings (Lin et al., 14 Jun 2025). The connection is conceptual rather than task-specific: both lines of work treat selective state-space models as a mechanism for combining efficient sequence modeling with self-supervised or semi-supervised training.

In applied terms, the clinical implication stated for U-Mamba2-SSL is that more reliable tooth and pulp segmentation from semi-supervised learning can accelerate treatment planning, endodontic assessment, and orthodontic interventions by reducing annotation burden and enabling automated analyses across diverse CBCT acquisitions (Tan et al., 24 Sep 2025). The method’s main open questions remain robustness across acquisition types, better handling of sparse regions of interest, and further control of pseudo-label error in difficult boundary regions.