PaperlayerView exact revision on arXiv

Source-linked AI summary

MedSegDiff-V2: Diffusion based Medical Image Segmentation with Transformer

Junde Wu, Wei Ji, Huazhu Fu, Min Xu, Yueming Jin, Yanwu Xu

arXiv:2301.11798v2eess.IVcs.CV

TL;DR

Medical image segmentation lacks an effective straightforward way to combine transformer representations with diffusion models because their features differ and transformer outputs are variance-sensitive. MedSegDiff-V2 introduces Anchor and Semantic Conditions, using U-SA, SS-Former, and NBP-Filter to integrate these representations, and reports superior performance across 20 segmentation tasks and five image modalities.

  • Problem

    Straightforwardly combining transformers with diffusion performs poorly because clean-image semantic features differ from noisy-mask diffusion features, while transformers are more sensitive to input variance.

  • Method

    MedSegDiff-V2 integrates transformer mechanisms into diffusion segmentation through Anchor Condition with U-SA and Semantic Condition with SS-Former and NBP-Filter.

  • Results

    MedSegDiff-V2 outperformed previous state-of-the-art methods on 20 medical image segmentation tasks with various image modalities.

  • Takeaways & Limitations

    MedSegDiff-V2 is presented as a benchmark for future research as the first transformer-based diffusion model for medical image segmentation.

Abstract

from arXiv · show

The Diffusion Probabilistic Model (DPM) has recently gained popularity in the field of computer vision, thanks to its image generation applications, such as Imagen, Latent Diffusion Models, and Stable Diffusion, which have demonstrated impressive capabilities and sparked much discussion within the community. Recent investigations have further unveiled the utility of DPM in the domain of medical image analysis, as underscored by the commendable performance exhibited by the medical image segmentation model across various tasks. Although these models were originally underpinned by a UNet architecture, there exists a potential avenue for enhancing their performance through the integration of vision transformer mechanisms. However, we discovered that simply combining these two models resulted in subpar performance. To effectively integrate these two cutting-edge techniques for the Medical image segmentation, we propose a novel Transformer-based Diffusion framework, called MedSegDiff-V2. We verify its effectiveness on 20 medical image segmentation tasks with different image modalities. Through comprehensive evaluation, our approach demonstrates superiority over prior state-of-the-art (SOTA) methodologies. Code is released at https://github.com/KidsWithTokens/MedSegDiff

Introduction

Medical image segmentation supports diagnosis and image-guided surgery, while diffusion models offer stochastic, uncertainty-aware predictions. MedSegDiff-V2 addresses the difficulty of combining transformers with diffusion through Anchor and Semantic Conditions.

  • Medical image segmentation divides images into regions of interest for applications including diagnosis and image-guided surgery.
  • Diffusion models generate diverse segmentation predictions through stochastic sampling, directly capturing prediction uncertainty.
  • Straightforward transformer-diffusion combinations perform poorly because their features arise from clean images versus corrupted masks, creating difficult fusion and variance challenges.
  • MedSegDiff-V2 integrates transformer mechanisms into diffusion-based medical segmentation using Anchor Condition and Semantic Condition strategies.
  • Anchor Condition uses Uncertain Spatial Attention to relax conditional features, while Semantic Condition uses SS-Former and a timestep-adaptive NBP-Filter to align semantic and noise features.
  • The framework reports state-of-the-art performance across 20 organ segmentation tasks spanning five image modalities.

Related Work

Prior work established transformers and diffusion models as promising approaches for medical image segmentation. Diffusion models additionally exploit stochastic sampling to form implicit segmentation ensembles, but controlling diversity remains important for convergence and accuracy.

  • Transformer-based models such as TransUNet, Swin-UNet, Swin-UNetr, and DS-TransUNet have been developed for medical image segmentation.
  • UNet-based diffusion segmentation models have emerged as strong approaches, motivating efforts to integrate transformer architectures into diffusion backbones.
  • Diffusion models use stochastic sampling to generate implicit ensembles of segmentations, improving segmentation performance across tasks including medical imaging.
  • Without effective diversity control, diffusion-generated ensembles may struggle to converge, making sample accuracy at each iteration important.

Method

MedSegDiff-V2 combines diffusion-based segmentation with transformer conditioning through an encoder-side Anchor Condition and embedding-level Semantic Condition. U-SA stabilizes anchor fusion, while SS-Former aligns noise and semantic features in frequency space across diffusion timesteps.

  • Diffusion backbone: Diffusion segmentation gradually noises the label and trains a reverse process to recover the clear segmentation step by step.The model starts from Gaussian noise and transforms the latent distribution toward the data distribution.
  • Conditioning framework: The diffusion model conditions its noise-estimation function on raw-image features extracted by a separate condition model.Transformer attention combines raw-image and current segmentation-map embeddings before UNet decoding.
  • Optimization: Training combines standard diffusion noise-prediction loss with an anchor loss supervising the condition model using soft Dice and cross-entropy losses.Condition-model supervision occurs at intervals controlled by α, with cross-entropy weighted by β; both are set to 5 and 10, respectively.
  • Conditioning framework: Direct transformer integration is difficult because raw-image semantic features differ from features abstracted from corrupted masks, while transformer outputs are sensitive to input variance.The paper identifies feature incompatibility and large per-timestep output variance as integration challenges.
  • Anchor Condition with U-SA: Anchor Condition inserts decoded condition features into diffusion encoder features, and U-SA smooths and selectively fuses them to provide a prediction range for refinement.U-SA applies a learnable Gaussian kernel, retains maximum responses, reduces anchor channels, and adds the enhanced feature across diffusion channels.
  • Semantic Condition with SS-Former: SS-Former uses bidirectional cross-attention in Fourier space with a timestep-adaptive Neural Band-pass Filter to align diffusion-noise and semantic embeddings.The filter learns a smooth frequency attention map from coordinates conditioned on timestep information, then transforms filtered features back with IFFT.

Dataset

Experiments use five medical image segmentation datasets spanning multi-organ abdominal segmentation and additional public medical imaging tasks.

  • General segmentation: AMOS2022 contains sixteen annotated anatomies and BTCV contains twelve annotated anatomies for abdominal multi-organ segmentation.Both datasets assess general segmentation performance.
  • Additional tasks: The additional public datasets are REFUGE-2, BraTs-2021, ISIC 2018, and TNMIX.

Implementation Details

Implementation uses standardized 256×256 inputs, end-to-end AdamW training on four NVIDIA A100 GPUs, and 100-step diffusion inference with ten ensemble runs.

  • Training setup: All images are resized to 256×256 pixels, and training uses AdamW with batch size 32 and initial learning rate 1 ×10−4.
  • Inference setup: Experiments run on 4 NVIDIA A100 GPUs with 100 diffusion steps for inference and 10 model runs for the ensemble.The ensemble uses fewer runs than MedSegDiff's 25 runs.

Main Results

MedSegDiff-V2 is evaluated against established CNN-, transformer-, and diffusion-based segmentation methods across abdominal multi-organ and multi-modality tasks. It achieves superior performance through its proposed conditioning modules and transformer-based design.

  • Abdominal Multi-organ Segmentation: MedSegDiff-V2 shows superior performance over SOTA methods on abdominal multi-organ segmentation benchmarks AMOS and BTCV.The comparison reports Dice scores against widely used CNN-, transformer-, and diffusion-based methods.
  • Abdominal Multi-organ Segmentation: MedSegDiff-V2 predicts more precise segmentation details than competing methods in low-contrast or ambiguous regions.The qualitative comparison is presented on BTCV.
  • Multi-modality Images: MedSegDiff-V2 surpasses all other methods in five tasks spanning different medical image modalities.The result is reported as evidence of generalization across segmentation tasks and image modalities.
  • Multi-modality Images: 2.0% on Optic-Cup, 1.9% on Brain-Tumor, and 3.9% on Thyroid Nodule are MedSegDiff-V2's Dice-score improvements over UNet-based MedSegDiff.The improvements underscore the effectiveness of the transformer-based backbone.
  • Ablation Study: Anchor Condition and Semantic Condition improve the diffusion model, while U-SA outperforms previous Spatial Attention and NBP-Filter strengthens SS-Former.SS-Former alone provides only marginal improvement, whereas combining it with NBP-Filter produces a significant improvement.

Analysis and Discussion

MedSegDiff-V2 improves diffusion-based segmentation through efficient implicit ensembling, diverse yet accurate samples, and fewer computational steps to stability. Its component combination balances uncertainty and confidence while outperforming simpler diffusion–transformer integrations.

  • Implicit Ensemble Effect: Implicit ensembles combine predictions from multiple samplings of one diffusion model, with performance typically improving rapidly before stabilizing near 50 ensembles.The evaluation uses average Dice Score on AMOS, with each configuration run 20 times.
  • Implicit Ensemble Effect: MedSegDiff-V2 surpasses MedSegDiff by 5% initially and maintains a lead exceeding 2% while requiring fewer ensembles to converge.The authors attribute this pattern to U-SA divergence control and SS-Former’s improved starting point.
  • Analysis of Uncertainty: U-SA increases sample diversity, while combining it with SS-Former significantly improves performance with still high confidence.SS-Former alone has the highest confidence but lowest diversity, whereas the combined model balances both properties.
  • Model Efficiency and Complexity: Model parameter count is not directly correlated with Gflops in diffusion segmentation because implicit ensembling contributes to computational cost.This makes convergence speed a practical efficiency measure for iterative diffusion inference.
  • Model Efficiency and Complexity: MedSegDiff-V2 reaches stability in fewer steps despite using more parameters and transformer blocks, requiring only half the Gflops of MedSegDiff.Gflops measures processing for one 256 × 256 image until implicit-ensemble stability, defined by less than 0.1% performance variance across the last ten ensembles.

Conclusion

MedSegDiff-V2 incorporates transformer mechanisms into a diffusion-based UNet for medical image segmentation. Comparative experiments report performance exceeding previous state-of-the-art methods across 20 tasks with varied image modalities.

  • MedSegDiff-V2 incorporates transformer mechanisms into the original UNet backbone for diffusion-based medical image segmentation.
  • Comparative experiments show the model outperformed previous SOTA methods on 20 medical image segmentation tasks.The tasks used various image modalities.
  • The authors present MedSegDiff-V2 as the first transformer-based diffusion model for medical image segmentation.
  • The paper proposes that MedSegDiff-V2 can serve as a benchmark for future research.
Loading 2301.11798v2…