DSVM-UNet : Enhancing VM-UNet with Dual Self-distillation for Medical Image Segmentation

AI summary

Overview

Research area: Medical image segmentation, specifically CNN/Transformer/Mamba-based U-shaped (UNet) architectures for skin lesion and abdominal organ segmentation.

Technical level: Advanced. The paper assumes familiarity with state space models (Mamba), UNet encoder-decoder designs, knowledge distillation, and standard segmentation metrics.

Scope: The paper introduces DSVM-UNet, a dual self-distillation training framework layered on top of a Vision Mamba UNet (VM-UNet) that improves segmentation accuracy on the ISIC2017, ISIC2018, and Synapse benchmarks without adding complex architectural modules.

What This Paper Is About

Recent medical segmentation models built on Vision Mamba (VM-UNet and its variants) improve accuracy by designing ever more complicated network architectures to better perceive semantic features. This paper asks whether comparable gains can be obtained from a simpler route: instead of redesigning the network, force the features already produced inside the UNet-shaped Vision Mamba to agree with each other at global and local levels through two self-distillation losses. The goal is state-of-the-art segmentation performance while keeping parameters and computation low.

Key Contributions

1. DSVM-UNet framework: A UNet-style segmentation framework that applies dual self-distillation so that knowledge flows from high-level features to low-level features within the network, with no added architectural complexity.

2. Two self-distillation methods: Projection self-distillation, which removes global inconsistency between encoder and decoder features by projecting all levels into a common space and aligning them to the final decoder feature; and progressive self-distillation, which removes local inconsistency by aligning features between adjacent encoder layers and adjacent decoder layers.

3. Empirical validation across benchmarks: Extensive experiments on ISIC2017, ISIC2018, and Synapse showing state-of-the-art performance with reduced computational cost compared with VM-UNet, VM-UNet++, and VM-UNetV2.

Main Findings

ISIC2017 (skin lesion): DSVM-UNet reaches 82.57% mIoU, 90.62% DSC, 96.69% Acc, 98.34% Spe, and 92.08% Sen, compared with 82.34% mIoU, 90.31% DSC, 96.70% Acc, 97.67% Spe, and 91.89% Sen for VM-UNetV2. Relative to VM-UNetV2, the paper reports a 0.23% increase in mIoU, 0.31% in DSC, 0.67% in Spe, and 0.19% in Sen. The Acc value of 96.69% is marginally below VM-UNetV2's 96.70%.

ISIC2018 (skin lesion): DSVM-UNet reaches 81.51% mIoU, 90.45% DSC, 95.43% Acc, 97.66% Spe, and 91.26% Sen, versus 81.37% mIoU, 89.73% DSC, 95.06% Acc, 97.13% Spe, and 88.64% Sen for VM-UNetV2. The paper describes these as similar improvements to the ISIC2017 results.

Synapse (multi-organ CT): DSVM-UNet achieves 81.68% DSC and 19.32 HD95. The paper reports that this outperforms VM-UNetV2 by 0.47% in DSC and reduces HD95 by 1.32%, and that the method achieves the best performance across five categories of organ segmentation. Among the baselines listed, MEW-UNet holds the lowest HD95 at 16.44 and VM-UNet holds 19.21.

Ablation on ISIC2018: Starting from the BceDice baseline at 81.31 mIoU, 89.63 F1, 94.76 Acc, 96.11 Spec, 89.65 Sens (average 90.29), adding projection self-distillation gives an average improvement of 0.5%, adding progressive self-distillation alone gives 0.75%, and combining both gives 0.97%, reaching 81.51 mIoU, 90.45 F1, 95.43 Acc, 97.66 Spec, 91.26 Sens (average 91.26).

Model complexity: On ISIC2017, DSVM-UNet uses 3.65 G FLOPs and 22.63 M parameters, compared with 4.11 G / 27.42 M for VM-UNet, 4.33 G / 24.81 M for VM-UNet++, and 4.40 G / 22.77 M for VM-UNetV2. The method is therefore both lighter and more accurate than VM-UNetV2 on this measurement.

Methodology in Plain English

The base network is VM-UNet, a UNet whose encoder and decoder blocks are Vision Mamba (VSS) blocks derived from VMamba. An input image produces features at multiple levels, and because those levels have different channel counts and spatial resolutions, the features are not directly comparable.

The authors add two training-time losses rather than new layers:

• Projection self-distillation. Every encoder and every decoder feature map is linearly projected and passed through a 1x1 1D convolution so that all levels share the same channel count and spatial size. The final decoder feature, which carries the richest high-level semantics, acts as the teacher, and all other levels are pulled toward it using a mean squared error distillation loss.
• Progressive self-distillation. Adjacent levels are aligned instead of all levels at once. A 1x1 2D convolution matches channel counts and bilinear interpolation matches spatial sizes, following the rule that the lower-dimensional feature is always aligned to the higher-dimensional one. The later (more semantic) feature is treated as the teacher, and alignment uses an L1 loss.

The total objective adds these two terms to the baseline segmentation loss with weights alpha and beta. Training uses 256x256 inputs, random flipping and rotation, batch size 32, AdamW with an initial learning rate of 1e-3 and CosineAnnealingLR down to 1e-5, 300 epochs, alpha = 1 and beta = 0.5, and encoder/decoder weights initialized from VMamba-S pretrained on ImageNet-1k. All experiments run on NVIDIA RTX A40 GPUs. BceDice loss is used for ISIC2017 and ISIC2018, and CeDice for Synapse.

Why This Matters

The result suggests that in medical segmentation, careful training-time feature alignment can substitute for the architecture engineering that currently dominates VM-UNet research, and can do so while lowering FLOPs and parameter count. That is a useful signal for a field where new variants typically gain accuracy by growing in size.

Real-world applications:

• Dermatology screening tools that delineate skin lesions from dermoscopy images, where ISIC-style pipelines are the standard benchmark.
• Abdominal organ segmentation from CT for surgical planning, radiotherapy targeting, and organ volumetry, corresponding to the eight-class Synapse setup (Aorta, Gallbladder, Kidney(L), Kidney(R), Liver, Pancreas, Spleen, Stomach).
• Deployment on modest clinical hardware, since the method runs at 3.65 G FLOPs and 22.63 M parameters, below the VM-UNet and VM-UNetV2 baselines.
• Research baselines: the two distillation losses can be added to other multi-level encoder-decoder segmentation networks without redesigning them.

Industry relevance: Medical imaging software vendors and device makers care about the accuracy-per-FLOP trade-off, and the paper's claim is precisely that accuracy improves while both FLOPs and parameters fall. The approach is a training change, not a deployment change, so it can be adopted without new inference infrastructure.

Future Directions

• The paper sets alpha = 1 and beta = 0.5 but reports no sensitivity analysis over those values; how the balance between projection and progressive distillation affects performance remains an open question.
• The authors demonstrate gains on ISIC2017, ISIC2018, and Synapse only. Whether the same losses transfer to 3D volumes, other modalities (for example MRI), or additional organ and lesion datasets is not reported.
• The paper reports HD95 only for Synapse and does not report it for the ISIC datasets, so boundary-quality improvements on skin lesion tasks are not quantified.
• Because the method is architecture-agnostic in principle, testing the same dual self-distillation losses on CNN- or Transformer-based UNets, rather than only on Vision Mamba UNets, would clarify how general the effect is.

Target Audience

Graduate students and researchers working on medical image segmentation, particularly those using Mamba/state space models or UNet variants; practitioners looking for accuracy gains without larger models; and readers interested in knowledge distillation and feature-alignment techniques applied to dense prediction. Readers need a working knowledge of segmentation metrics (mIoU, DSC, HD95, Acc, Spe, Sen) and of distillation loss design to follow the methodology.

Authors’ abstract

Vision Mamba models have been extensively researched in various fields, which address the limitations of previous models by effectively managing long-range dependencies with a linear-time overhead. Several prospective studies have further designed Vision Mamba based on UNet(VM-UNet) for medical image segmentation. These approaches primarily focus on optimizing architectural designs by creating more complex structures to enhance the model's ability to perceive semantic features. In this paper, we propose a simple yet effective approach to improve the model by Dual Self-distillation for VM-UNet (DSVM-UNet) without any complex architectural designs. To achieve this goal, we develop double self-distillation methods to align the features at both the global and local levels. Extensive experiments conducted on the ISIC2017, ISIC2018, and Synapse benchmarks demonstrate that our approach achieves state-of-the-art performance while maintaining computational efficiency. Code is available at https://github.com/RoryShao/DSVM-UNet.git.