Optimization in Sparse 2D to Dense 3D Weakly Supervised Learning: Application to Multi-Label Segmentation of Large ex vivo MRI Data

Technical Summary: Optimization in Sparse 2D to Dense 3D Weakly Supervised Learning for Multi-Label Segmentation of Large ex vivo MRI Data

Problem Statement
High-resolution ex vivo MRI of the spinal cord (typically 50–100 µm isotropic) offers detailed anatomical characterization of multiple sclerosis (MS) pathology, including gray matter (GM) and white matter (WM) lesions. However, the sheer volume of data (thousands of slices per sample) renders dense, voxel-wise manual annotation prohibitively expensive. Researchers are often forced to rely on sparse 2D annotations (e.g., 428 slices across 104,000+ total slices). This creates a technical dichotomy: 2D models trained on sparse slices lack volumetric context, leading to inconsistent predictions along the Z-axis, while 3D models require dense ground truth that does not exist. Existing solutions, such as 2.5D strategies or cross-teaching, attempt to bridge this gap, but guidelines for optimizing such "Sparse-to-Dense" pipelines remain ambiguous, particularly regarding the transferability of 2D optimization strategies to 3D contexts.

Methodology
The authors propose a three-stage Sparse-to-Dense framework utilizing a 2D Teacher and a 3D Student architecture, both based on nnU-Net v2.

1. Dataset: The study utilizes 9.4T ex vivo spinal cord MRI data from MS patients (12 samples, 43,719 axial slices). Ground truth consists of 428 manually segmented slices across four classes: Healthy WM, Healthy GM, Lesion WM, and Lesion GM.
2. Stage 1 (2D Teacher): A 2D nnU-Net is trained on the sparse annotations. The authors systematically optimized this stage by testing input modalities (Magnitude vs. Phase), preprocessing techniques (CLAHE, contrast stretching), spatial augmentation, and soft-label regularization.
3. Stage 2 (Pseudo-Label Generation): The optimized 2D Teacher generates dense volumetric pseudo-labels for the entire dataset using ensembling (4-fold cross-validation) and Test Time Augmentation (TTA).
4. Stage 3 (3D Student): A 3D nnU-Net is trained on the dense pseudo-labels to learn volumetric features and longitudinal consistency. The authors applied the same systematic optimization (preprocessing, augmentation, soft labels) to the 3D Student to determine if 2D strategies transfer effectively.
5. Evaluation: Performance was assessed using Dice Score, 95th Percentile Hausdorff Distance (HD95), and a novel Geometric Consistency metric (Inter-slice Dice, DSCz) to measure Z-axis smoothness.

Key Contributions
The paper makes three primary contributions:

1. The Perception Divergence: Empirical evidence demonstrating that preprocessing techniques optimized for human visual perception (specifically CLAHE and contrast stretching) are detrimental to machine learning performance in this domain. These methods degraded Gray Matter Lesion Dice Scores by up to 25 percentage points by disrupting global intensity statistics essential for the model.
2. Divergent Regularization Needs: Identification of a critical conflict in Sparse-to-Dense pipelines. While aggressive spatial augmentation and soft-labeling significantly improved the 2D Teacher (improving Lesion WM detection by >11 Dice points), these same techniques degraded the performance of the 3D Student.
3. Open Access Model: The release of a trained model for segmenting WM, GM, and MS lesions in ex vivo human spinal cord MRI, available via the Spinal Cord Toolbox.

Results

• Input Modalities: Dual-input models (Magnitude + Phase) outperformed magnitude-only models, particularly for Lesion GM stability and boundary precision (HD95), confirming the necessity of phase data for resolving fine anatomical boundaries.
• Preprocessing: Contrary to the hypothesis that visual enhancement aids the model, applying CLAHE to magnitude images and contrast stretching to phase images caused severe performance drops. The 2D Lesion GM Dice fell from 0.665 to 0.415 with magnitude CLAHE. The 3D model showed similar degradation, though it exhibited slightly more resilience in boundary metrics (HD95) for Lesion GM, suggesting volumetric context offers some protection against texture distortion.
• Augmentation:
  • 2D: Aggressive spatial augmentation (rotations, scaling, shearing) was critical. Without it, Lesion WM Dice dropped to 0.631; with moderate augmentation (Aug2), it reached 0.743.
  • 3D: The opposite trend was observed. The 3D Student achieved its highest performance (Total Average Dice 0.772) with no spatial augmentation. Introducing augmentation degraded performance, with the most aggressive settings causing a collapse in boundary fidelity.
• Soft Labels:
  • 2D: Soft labeling (Soft Seg 2) improved performance, raising Lesion WM Dice to 0.744 and refining boundary precision.
  • 3D: Soft labeling provided no benefit and slightly worsened results. The authors suggest the pseudo-labels generated by the 2D ensemble already smooth out uncertainty, making additional soft-loss regularization counter-productive (over-smoothing).
• Volumetric Consistency: The transition from 2D to 3D yielded a significant gain in geometric continuity. The naive stacking of 2D predictions resulted in a low Inter-slice Dice (DSCz) of 0.785, whereas the 3D Student achieved 0.900. However, this smoothness came with a trade-off: the 3D Student had a slightly higher HD95 (1.00 mm) compared to the optimized 2D Teacher (0.72 mm), indicating a shift from strict in-plane contour adherence to biological plausibility.

Significance and Claims
The paper claims that the transition from 2D to 3D in weakly supervised learning is not merely an inference trick but a fundamental shift in learning dynamics. The authors argue that 3D architectures trained on dense pseudo-labels exhibit different optimization landscapes than their 2D counterparts. Specifically, the "shortcut learning" risks that necessitate aggressive augmentation in sparse 2D settings are mitigated in 3D by the density of the pseudo-labels, rendering such augmentation harmful. Similarly, the uncertainty modeling beneficial for 2D human-rater mimicry becomes redundant or destructive when the 3D model is trained on already-averaged ensemble predictions.

The study concludes that standard computer vision practices (human-centric preprocessing, aggressive augmentation, soft labeling) cannot be blindly transferred from 2D to 3D medical segmentation tasks. Instead, 3D models trained on dense pseudo-labels require a distinct, conservative regularization strategy that prioritizes volumetric consistency over the specific noise-reduction techniques required for sparse 2D training. This work challenges the assumption that "more data" (dense pseudo-labels) simply allows for the direct application of 2D optimization heuristics, highlighting a "dimension gap" in regularization needs.