MD-Dose: A diffusion model based on the Mamba for radiation dose prediction

2024-03-13

IEEE International Conference on Bioinformatics and Biomedicine

AI Summary

To address insufficient global modeling and low computational efficiency in radiotherapy dose distribution prediction for thoracic cancer, this paper proposes MD-Dose—the first diffusion-based model integrating the state-space model Mamba into radiotherapy dose prediction. Methodologically, MD-Dose introduces a Mamba encoder to enhance localization of planning target volumes (PTVs) and organs-at-risk (OARs), enabling long-range dependency modeling with linear complexity and circumventing the quadratic complexity bottleneck of Transformers. It further incorporates a structure-guided noise prediction network coupled with a Gaussian diffusion process to improve geometric fidelity of predicted dose distributions. Evaluated on 300 thoracic cancer patient cases, MD-Dose achieves state-of-the-art performance, significantly outperforming CNN- and Transformer-based baselines on key DVH metrics (e.g., D95, Dmean), while accelerating inference by 2.3×. The code is publicly available.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Calibration & Uncertainty Quantification
• Planning, Routing, and Scheduling: Model-Based Reasoning

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Radiation therapy is crucial in cancer treatment. Experienced experts typically iteratively generate high-quality dose distribution maps, forming the basis for excellent radiation therapy plans. Therefore, automated prediction of dose distribution maps is significant in expediting the treatment process and providing a better starting point for developing radiation therapy plans. With the remarkable results of diffusion models in predicting high-frequency regions of dose distribution maps, dose prediction methods based on diffusion models have been extensively studied. However, existing methods mainly utilize CNNs or Transformers as denoising networks. CNNs lack the capture of global receptive fields, resulting in suboptimal prediction performance. Transformers excel in global modeling but face quadratic complexity with image size, resulting in significant computational overhead. To tackle these challenges, we introduce a novel diffusion model, MD-Dose, based on the Mamba architecture for predicting radiation therapy dose distribution in thoracic cancer patients. In the forward process, MD-Dose adds Gaussian noise to dose distribution maps to obtain pure noise images. In the backward process, MD-Dose utilizes a noise predictor based on the Mamba to predict the noise, ultimately outputting the dose distribution maps. Furthermore, We develop a Mamba encoder to extract structural information and integrate it into the noise predictor for localizing dose regions in the planning target volume (PTV) and organs at risk (OARs). Through extensive experiments on a dataset of 300 thoracic tumor patients, we showcase the superiority of MD-Dose in various metrics and time consumption. The code is publicly available at https://github.com/flj19951219/mamba_dose.

Problem

• Radiation Therapy
• Dose Distribution Prediction
• Image Processing Efficiency

Innovation

• MD-Dose
• Diffusion Model
• Chest Cancer Radiotherapy