Efficient Medical Image Segmentation with Rectified Flow-Based Diffusion Transformer

Shunichi Kato, Masaki Nishimura, Takaya Ueda, Yu Song, Ikuko Nishikawa

Abstract: Several diffusion models have recently been proposed for high-quality data generation, including image and video synthesis. Medical image segmentation is one of the most promising application areas, where high pixel-level accuracy is essential. Achieving such accuracy using diffusion models generally requires a large number of denoising steps, which incurs significant computational costs. To address this issue, the present work proposes a segmentation method that uses Rectified Flow to approximate the stochastic reverse diffusion process with a deterministic ordinary differential equation (ODE) trajectory. Furthermore, the Transformer-based segmentation model incorporates skip connections, similar to those in U-Net, to reduce information loss. To verify the effectiveness of the proposed model, experiments were conducted on the ISIC2018 skin lesion segmentation dataset. The results show that the model achieves superior accuracies in IoU and dice scores compared to existing diffusion models, even with only three denoising steps.

External IDs:doi:10.1007/978-981-95-4100-3_1