SegDT: Diffusion Transformer in Medical Imaging

• The paper introduces SegDT, a diffusion transformer-based segmentation model that integrates a compact latent-space architecture and rectified flow for accelerated inference.
• It achieves state-of-the-art performance on ISIC skin lesion benchmarks with top Dice (94.76) and IoU (91.40) scores while significantly reducing computational steps.
• SegDT offers real-time inference on commodity GPUs with a minimal memory footprint, making it highly suitable for near-patient and edge-device deployments.

SegDT is a diffusion transformer-based image segmentation model designed for medical imaging tasks, particularly skin lesion segmentation. Developed to achieve state-of-the-art accuracy with high specificity and efficiency on commodity GPUs, SegDT integrates a compact latent-space architecture, a DiT backbone, and a rectified flow sampling strategy that accelerates inference without significant quality degradation. The design emphasizes adaptability to low-cost hardware, real-time throughput, and a minimal memory footprint, making it suitable for near-patient or edge-device deployment (Bekhouche et al., 21 Jul 2025).

1. Model Components and Architecture

SegDT's architecture consists of three main components operating in the latent space:

• VAE Encoder: Utilizes a pretrained VAE encoder from the TAESD system to project both input dermoscopic images ($x$) and, during training, ground-truth masks ($m$), into a shared latent space $\mathcal Z$ of size $C' \times H' \times W'$ with $H'=H/8,\;W'=W/8,\;C'=4$.
• Diffusion Transformer (DiT–XS): Receives flattened patches of the latent space as input and is conditioned on embeddings for both the diffusion timestep $t$ and the image latent $y$. The main DiT backbone consists of 12 blocks, incorporating self-attention, cross-attention to $y$, feedforward networks, adaptive layer normalization (adaLN), and DropPath regularization.
• VAE Decoder: A pretrained, frozen VAE decoder reconstructs a three-channel image from the denoised latent. Segmentation is obtained via thresholding (threshold: 0.2) the pixel values of this reconstructed image.

The inference pipeline is as follows:

1. Encode the input image $x$ to obtain the image latent $y$.
2. Initialize $m$0.
3. For $m$1, process $m$2 and $m$3 through DiT blocks, output a $m$4 tensor (denoised latent and variance), and update $m$5 via rectified flow:

$m$6

1. Decode $m$7 with the VAE decoder and binarize to yield the mask.

This design leverages approximately $m$8 million parameters and a compute footprint of $m$9 GFLOPs per $\mathcal Z$0 image.

2. Diffusion and Rectified Flow Formulation

SegDT follows a Gaussian diffusion process defined by a fixed noise schedule $\mathcal Z$1 with:

$\mathcal Z$2

and forward noise injection:

$\mathcal Z$3

The conventional reverse process samples

$\mathcal Z$4

SegDT replaces the direct parameterization of $\mathcal Z$5 with a velocity field $\mathcal Z$6, learning to predict:

$\mathcal Z$7

and updates latents via a one-step Euler integration. The primary training loss is the mean squared error between predicted and true velocity:

$\mathcal Z$8

By adopting rectified flow, SegDT reduces the typical sampling steps from 35–1000 in DDPM-based segmentation to 15, with negligible loss in mask quality.

3. Training Regime and Implementation

Datasets and Preprocessing

SegDT is trained and evaluated on the ISIC 2016, 2017, and 2018 skin lesion segmentation benchmarks:

• ISIC 2016: 900 train / 335 test images.
• ISIC 2017: 2000 train / 150 val / 600 test images.
• ISIC 2018: 2594 train / 100 val / 1000 test images.

Inputs are bilinearly resized to $\mathcal Z$9 and pixel-normalized to $C' \times H' \times W'$0.

Hyperparameters

• Diffusion timesteps: $C' \times H' \times W'$1 (rectified flow in 15 steps)
• Noise schedule: $C' \times H' \times W'$2 linearly spaced in $C' \times H' \times W'$3
• Optimizer: Adam, batch size 32, learning rate $C' \times H' \times W'$4, 100 epochs, learning rate divided by 10 at epoch 50
• DiT-XS backbone: Patch size $C' \times H' \times W'$5, embedding dim $C' \times H' \times W'$6, MLP ratio = 4, DropPath = 0.1
• Binarization threshold: 0.2

Hardware and Runtime

• Training: dual NVIDIA RTX 3090 GPUs (24 GB each)
• Inference: single RTX 3060 (12 GB), achieving approximately 7 fps with 15 rectified steps
• Memory footprint: $C' \times H' \times W'$7M parameters, $C' \times H' \times W'$8 GFLOPs per image

4. Quantitative and Qualitative Evaluation

SegDT achieves leading performance on standard skin lesion segmentation metrics, using significantly less compute than prior methods.

Selected results (ISIC 2016):

Method	Dice↑	IoU↑	ACC↑	SE↑	SP↑	GFLOPs↓	Params(M)↓
MobileUNETR	92.80	87.47	96.59	93.03	96.87	1.30	3.00
DU-Net+	92.46	85.25	96.89	95.75	95.64	54.00	39.00
SegDT	94.76	91.40	97.08	93.35	99.44	3.68	9.95

SegDT consistently yields the highest Dice and IoU across all three ISIC benchmarks, along with top specificity scores, indicating reduced false positives—an important criterion for clinical deployment.

Inference speed: IDDPM-style (35 steps) yields ∼0.25 s/image; SegDT's rectified-flow (15 steps) ∼0.12 s/image ($C' \times H' \times W'$9 2× faster).

Qualitative output: SegDT successfully delineates challenging, irregular lesions. Occasional under-segmentation of very small or highly textured lesions is observed, often attributed to DiT's patch-level receptive field.

5. Strengths and Limitations

Strengths

• Achieves state-of-the-art Dice/IoU on ISIC datasets with 7–15× fewer FLOPs compared to UNet-based diffusion models.
• Enables real-time inference (≈ 7 fps) on standard GPUs as a consequence of rectified flow acceleration.
• Compact parameter count (9.95M) and latent-space diffusion via an 8× downsampled representation.

Limitations

• Reduced sensitivity to very small or highly irregular lesions; potential improvements could involve explicit multi-scale processing or refinement stages.
• Uses a fixed binarization threshold, which may require tuning on new datasets.
• Fixed, pretrained VAE decoder imposes limitations when applying the model to out-of-domain image types due to potential domain shift.

A plausible implication is that the decoder constraint and thresholding heuristic may need to be revisited for robust generalization to non-dermoscopic modalities.

6. Deployment Considerations and Future Directions

SegDT's low compute requirement, compactness, and high specificity support its application in resource-constrained and real-time settings, such as edge devices or near-patient diagnostic tools. High specificity reduces false alarms, minimizing unnecessary follow-up and improving clinical workflow efficiency.

Future extensions suggested include adapting SegDT for 3D volumetric data (e.g., CT/MRI segmentation) using volumetric diffusion transformers or integrating clinical metadata for enhanced context conditioning. This suggests generalizability to broader medical image analysis applications, contingent on appropriate retraining and adaptation (Bekhouche et al., 21 Jul 2025).