---
title: 'MTMed3D: Unified 3D Medical Imaging'
url: https://www.emergentmind.com/topics/mtmed3d
type: topic
---

# MTMed3D: Unified 3D Medical Imaging

MTMed3D is an end-to-end multi-task Transformer-based framework developed for 3D medical imaging, designed to jointly perform object detection, semantic segmentation, and classification from volumetric data. Addressing the inefficiencies of single-task models, MTMed3D introduces a unified architecture where all three core tasks share a 3D Swin-Transformer encoder, with subsequent specialized CNN decoders. The model is evaluated on multi-modal brain tumor MRIs (BraTS 2018/2019), delivering state-of-the-art detection performance while maintaining competitive segmentation and classification, with substantial reductions in computational cost, parameter count, and inference latency compared to single-task models [2511.12373].

## 1. Model Architecture

### 1.1. Shared 3D Swin-Transformer Encoder

MTMed3D inputs a multi-modal MRI volume ($X \in \mathbb{R}^{4 \times 240 \times 240 \times 155}$), which is center-cropped and normalized to $\tilde X \in \mathbb{R}^{4 \times 96 \times 96 \times 96}$. The volume is divided into non-overlapping 3D patches, then linearly embedded:
$$
Z^0 = \operatorname{LinearEmbed}\bigl(\operatorname{PatchPartition}(\tilde X)\bigr) \in \mathbb{R}^{M \times C},
$$
where $M$ is the number of patches and $C$ the feature dimension. Four hierarchical stages of Swin-Transformer blocks and patch-merging layers generate a set of multi-scale feature maps $\{Z^i\}_{i=1}^4$.

Each Swin-Transformer block alternates Windowed Multi-Head Self-Attention (W-MSA) with Shifted-Window Multi-Head Self-Attention (SW-MSA), ensuring localized and global context aggregation:
$$
\begin{aligned}
Z'_l &= \mathrm{W\text{-}MSA}(\mathrm{LN}(Z_{l-1})) + Z_{l-1},\\
Z_l  &= \mathrm{MLP}(\mathrm{LN}(Z'_l)) + Z'_l,\\
Z'_{l+1}&=\mathrm{SW\text{-}MSA}(\mathrm{LN}(Z_l)) + Z_l,\\
Z_{l+1}&=\mathrm{MLP}(\mathrm{LN}(Z'_{l+1})) + Z'_{l+1}.
\end{aligned}
$$

This design enables hard sharing of contextual, multi-scale features—eliminating redundancy across tasks.

### 1.2. Task-Specific CNN Decoders

- **Detection Decoder:** Features $\{F_2, F_3, F_4, F_5\}$ are unified in channel space using $1 \times 1 \times 1$ convolutions, normalized (GroupNorm), and activated (ReLU). The decoder employs a Path Aggregation Network (PANet) to propagate low-level features and facilitates multi-scale object prediction, producing classification scores and 3D bounding box offsets.
- **Segmentation Decoder:** Adopts a U-Net-style architecture as in Swin UNETR. Encoder features are routed via long skip connections into residual + deconvolutional layers, outputting region probability maps for Whole Tumor (WT), Tumor Core (TC), and Enhancing Tumor (ET).
- **Classification Decoder:** Refined segmentation features are input to a DenseNet-121, where feature maps are concatenated at every layer ($x_{l} = H_l([x_0, x_1, \dots, x_{l-1}])$). After global pooling, a fully connected layer predicts high- vs. low-grade glioma status.

## 2. Multi-Task Learning and Optimization

### 2.1. Composite Training Objective

MTMed3D is trained with a joint objective:
$$
\mathcal{L}_{\mathrm{total}} = w_1\,\mathcal{L}_{\mathrm{seg}} + w_2\,\mathcal{L}_{\mathrm{cls}} + w_3\,\mathcal{L}_{\mathrm{det}}
$$
where $\mathcal{L}_{\mathrm{seg}}$ (segmentation) is the Dice loss over three regions, $\mathcal{L}_{\mathrm{cls}}$ (classification) is Focal Loss, and $\mathcal{L}_{\mathrm{det}}$ (detection) combines Smooth L1 loss for bounding box regression and Focal Loss for object classification.

### 2.2. Gradient-Normalization for Dynamic Task Weighting

Task-balancing is achieved using GradNorm, adaptively adjusting $\{w_i\}$ during training. Let $G_w^{(i)}(t) = \|\nabla_w[w_i(t)\,L_i(t)]\|_2$ be the gradient norm for task $i$ at iteration $t$, with average $\overline G_w(t) = \frac{1}{T}\sum_i G_w^{(i)}(t)$, and relative inverse training rate
$$
r_i(t) = \frac{L_i(t)/L_i(0)} {\frac{1}{T}\sum_j [L_j(t)/L_j(0)]}.
$$
GradNorm minimizes $|G_w^{(i)}(t)-\overline G_w(t)\,r_i(t)|$ with respect to $\{w_i\}$. This ensures balanced gradient flow and prevents any single task from monopolizing encoder learning capacity. Synergistic feature learning is promoted as decoders back-propagate into the shared Swin backbone.

## 3. Data Sources, Preprocessing, and Training

### 3.1. Datasets and Labels

- **Datasets:** BraTS 2018 (285 cases: 210 high-grade glioma, 75 low-grade glioma) and BraTS 2019 (335 cases: 259 HGG, 76 LGG).
- **Modalities:** T1, T1ce, T2, FLAIR (stacked into 4-channel volumes).
- **Labels:** Segmentation masks for WT/TC/ET; categorical classification label (HGG/LGG); detection bounding boxes autogenerated from segmentation.

### 3.2. Augmentation and Training Protocols

- **Spatial augmentations:** random 3-axis flips ($p=0.5$), random 3D rotations ($p=0.75$).
- **Photometric augmentations:** per-channel intensity shifts and scalings (range $(-0.1, 0.1)$).
- **Cropping:** All volumes cropped to $96^3$ to accommodate GPU memory constraints.
- **Optimization:** AdamW optimizer with cosine annealing learning rate. Branch-specific learning rates: $10^{-4}$ (segmentation), $10^{-5}$ (detection/classification).
- **Batch size:** 1.
- **Validation:** Five-fold cross-validation; typical training runs are $200-300$ epochs.

### 3.3. Hardware

Experiments executed on Intel i7-13700F CPU and NVIDIA RTX 4070 Ti (12 GB VRAM).

## 4. Quantitative Results and Empirical Analysis

### 4.1. Overall Task Performance

| Task         | MTMed3D (5-fold Avg/Best)           | Leading Baseline(s)          |
|--------------|-------------------------------------|------------------------------|
| Detection    | mAP@[0.1:0.5]=0.9711/0.9082         | RetinaNet: 0.8664 / 0.6768   |
|              | mAP@0.5=0.8650                      | nnDetection/MedYOLO: ~0.8360/0.8610 |
| Segmentation | Dice WT=0.8793, TC=0.8019, ET=0.7755| Top single-task methods: WT=0.9082, TC=0.8183, ET=0.8038|
| Classification| Acc=0.9193, Sens=0.9761, Spec=0.7634| Mask-RCNN(2D)=0.9630, CNN(2D)=0.9715, SE-ResNeXt=0.9745 |

MTMed3D outperforms previous 3D detection methods and matches or exceeds segmentation accuracy of leading single-task architectures.

### 4.2. Efficiency Metrics

- **Task-averaged accuracy**: MTMed3D: Seg Dice avg $=0.8189$; Acc $=0.9193$; Det mAP@[0.1:0.5] $=0.9082$ vs. Single-task averages: Seg Dice $=0.8052$; Acc $=0.9017$; Det mAP@[0.1:0.5] $=0.8217$.
- **Resource reductions** vs. sum of three single-task models:
  - MACs: $2 \times 10^{11}$ vs $4 \times 10^{11}$ ($-46.6$\%)
  - FLOPs: $4 \times 10^{11}$ vs $7 \times 10^{11}$ ($-42.9$\%)
  - Parameters: $8 \times 10^7$ vs $1 \times 10^8$ ($-47.8$\%)
  - Inference time: $0.10$ s vs $0.22$ s ($-55.7$\%)
  - Model size: 307 MB vs 588 MB

### 4.3. Ablation Studies

- **PANet vs. FPN**: PANet head in detection, combined with GradNorm, yields superior ET Dice, lower Hausdorff distance (HD), and improved mAP/mAR.
- **Balancing strategy**: GradNorm produces more balanced multi-task results and notably higher classification specificity and accuracy than MGDA.

## 5. Design Rationale, Limitations, and Prospective Developments

Hard parameter sharing via a 3D Swin-Transformer backbone leverages both long-range dependencies and hierarchical feature context, supporting robust multi-task performance, particularly for detection due to transfer of segmentation cues. A single forward pass through MTMed3D suffices for all three tasks, resulting in halved computational, memory, and latency footprints relative to triplicate single-task systems.

Classification accuracy, however, remains below leading dedicated 2D/3D classifiers. This is attributed to the lack of a learnable $[\textrm{CLS}]$ token, imbalanced 3D data, and lack of slice-level pre-processing. Future directions proposed include introducing a learnable classification token or cross-attention mechanisms, pre-training the Swin encoder on larger 3D medical datasets, expanding the framework to additional anatomical sites and tasks such as registration or prognosis, and evaluating dynamic task weighting/auxiliary losses beyond GradNorm.

## 6. Context and Significance in Multi-Task Medical Imaging

MTMed3D constitutes the first end-to-end, Swin-Transformer-based multi-task network for 3D medical imaging that simultaneously addresses detection, segmentation, and classification. This demonstrates both technical feasibility and substantial system-level efficiency for multi-purpose clinical AI pipelines, suggesting a shift toward unifying architectures in volumetric diagnostic applications [2511.12373].

Source: https://www.emergentmind.com/topics/mtmed3d