Medical Semantic Segmentation with Diffusion Pretrain

TL;DR

The paper tackles the challenge of learning robust voxel-level representations for 3D medical image segmentation under limited labels. It introduces an anatomically guided diffusion pretraining framework that couples an auxiliary diffusion process with a Body Part Regressor to provide universal 3D coordinates, yielding generalized feature representations. In experiments on BTCV and FLARE, the method achieves a Dice of 67.8 in non-linear probing and outperforms restorative pretraining by 7.5%, while remaining competitive with leading contrastive pretraining approaches. The work demonstrates that diffusion-based self-supervision, when guided by anatomical structure, can be a strong, non-contrastive alternative for high-resolution medical segmentation with good cross-dataset transfer.

Abstract

Recent advances in deep learning have shown that learning robust feature representations is critical for the success of many computer vision tasks, including medical image segmentation. In particular, both transformer and convolutional-based architectures have benefit from leveraging pretext tasks for pretraining. However, the adoption of pretext tasks in 3D medical imaging has been less explored and remains a challenge, especially in the context of learning generalizable feature representations. We propose a novel pretraining strategy using diffusion models with anatomical guidance, tailored to the intricacies of 3D medical image data. We introduce an auxiliary diffusion process to pretrain a model that produce generalizable feature representations, useful for a variety of downstream segmentation tasks. We employ an additional model that predicts 3D universal body-part coordinates, providing guidance during the diffusion process and improving spatial awareness in generated representations. This approach not only aids in resolving localization inaccuracies but also enriches the model's ability to understand complex anatomical structures. Empirical validation on a 13-class organ segmentation task demonstrate the effectiveness of our pretraining technique. It surpasses existing restorative pretraining methods in 3D medical image segmentation by $7.5\%$, and is competitive with the state-of-the-art contrastive pretraining approach, achieving an average Dice coefficient of 67.8 in a non-linear evaluation scenario.

Figures (2)

• Figure 1: Model schema. Image patch $x_0$ is passed through diffusion process resulting in noised patch $x_t$. Independently, it is passed through a Body Part Regressor network to produce dense coordinate map $x_{BPR}$ of the same spatial resolution. $x_t$ and $x_{BPR}$ are then concatenated and passed through a trainable backward diffusion network to produce dense voxel-wise representations $\widetilde{x_t}$. For downstream supervised segmentation task a non-linear probing model is trained on top of frozen feature representations. All patches are three-dimensional.
• Figure 2: Ablation of the number of diffusion timesteps. The results of image segmentation for objects of different sizes: Big, Medium, and Small, non-linear probing regime. Average corresponds to all organs' average. Recall how segmentation quality spikes around 10 diffusion steps and then degrades for Medium and Small objects, while plateauing for Big objects.