LEAF: Latent Diffusion with Efficient Encoder Distillation for Aligned Features in Medical Image Segmentation

TL;DR

LEAF introduces a novel latent diffusion-based approach for medical image segmentation that improves accuracy through feature distillation and direct segmentation map prediction without increasing computational costs.

Contribution

The paper presents LEAF, a method that enhances diffusion models for segmentation by replacing noise prediction with direct segmentation and aligning features via distillation, without adding model complexity.

Findings

Improved segmentation accuracy across multiple datasets.

No additional parameters or inference cost.

Effective feature alignment with transformer encoders.

Abstract

Leveraging the powerful capabilities of diffusion models has yielded quite effective results in medical image segmentation tasks. However, existing methods typically transfer the original training process directly without specific adjustments for segmentation tasks. Furthermore, the commonly used pre-trained diffusion models still have deficiencies in feature extraction. Based on these considerations, we propose LEAF, a medical image segmentation model grounded in latent diffusion models. During the fine-tuning process, we replace the original noise prediction pattern with a direct prediction of the segmentation map, thereby reducing the variance of segmentation results. We also employ a feature distillation method to align the hidden states of the convolutional layers with the features from a transformer-based vision encoder. Experimental results demonstrate that our method enhances the performance of the original diffusion model across multiple segmentation datasets for different disease types. Notably, our approach does not alter the model architecture, nor does it increase the number of parameters or computation during the inference phase, making it highly efficient.