A Fast and Efficient Modern BERT based Text-Conditioned Diffusion Model for Medical Image Segmentation

TL;DR

This paper introduces FastTextDiff, a diffusion-based medical image segmentation model that leverages ModernBERT for efficient, text-conditioned segmentation, significantly improving accuracy and training speed by integrating clinical notes and multi-modal attention.

Contribution

The paper presents a novel, scalable diffusion model using ModernBERT for text-conditioned medical image segmentation, replacing Clinical BioBERT and enhancing efficiency and accuracy.

Findings

ModernBERT effectively encodes clinical knowledge for segmentation.

FastTextDiff outperforms traditional diffusion models in accuracy.

Training efficiency is improved using FlashAttention 2 and large-scale pretraining.

Abstract

In recent times, denoising diffusion probabilistic models (DPMs) have proven effective for medical image generation and denoising, and as representation learners for downstream segmentation. However, segmentation performance is limited by the need for dense pixel-wise labels, which are expensive, time-consuming, and require expert knowledge. We propose FastTextDiff, a label-efficient diffusion-based segmentation model that integrates medical text annotations to enhance semantic representations. Our approach uses ModernBERT, a transformer capable of processing long clinical notes, to tightly link textual annotations with semantic content in medical images. Trained on MIMIC-III and MIMIC-IV, ModernBERT encodes clinical knowledge that guides cross-modal attention between visual and textual features. This study validates ModernBERT as a fast, scalable alternative to Clinical BioBERT in diffusion-based segmentation pipelines and highlights the promise of multi-modal techniques for medical image analysis. By replacing Clinical BioBERT with ModernBERT, FastTextDiff benefits from FlashAttention 2, an alternating attention mechanism, and a 2-trillion-token corpus, improving both segmentation accuracy and training efficiency over traditional diffusion-based models.