Source-linked AI summary
Multimodal Deep Learning for Uncertainty-Aware Radiation Pneumonitis Risk Prediction
Jin Yang, Tian Liu, Jing Wang, Robert Samstein, Kenneth Rosenzweig, Julie Bloom, Ming Chao
TL;DR
RP risk prediction is limited by conventional representations and by the lack of individualized uncertainty estimates. MM-DiT combines self-supervised CT–dose learning, latent diffusion, and Bayesian prediction with label-noise modeling. Across two independent cohorts, it showed strong discrimination and calibration, while uncertainty decomposition and decision-curve analysis supported reliability-aware clinical assessment.
Problem
Conventional RP models inadequately capture complex spatial and patient-specific factors, while existing models generally do not distinguish sources of individual predictive uncertainty.
Method
MM-DiT integrates planning CT and 3D dose distributions through self-supervised multimodal pretraining, latent diffusion refinement, Bayesian prediction, and learnable label-noise modeling.
Results
Across two independent cohorts, MM-DiT demonstrated strong predictive discrimination and probability calibration with consistent uncertainty patterns, and decision-curve analysis showed positive net benefit across broad thresholds.
Takeaways & Limitations
MM-DiT extends RP risk stratification by pairing individualized risk probabilities with aleatoric, epistemic, and label uncertainty measures that inform prediction reliability.
Takeaways & Limitations
The retrospective study primarily used pretreatment CT and dose distributions, so additional multi-institutional validation and toxicity-rich longitudinal data are needed.
Abstract
from arXiv · showhide
Radiation pneumonitis (RP) is a common and clinically significant toxicity of thoracic radiation therapy that can cause pulmonary morbidity and impair quality of life. Although conventional dose-volume histogram-based metrics and normal tissue complication probability models are widely used for RP risk assessment, they inadequately capture the complex spatial, anatomical, and patient-specific factors underlying radiation-induced lung injury. Recent machine learning approaches have improved RP risk prediction by integrating multimodal clinical and imaging information; however, most provide a point risk estimate without quantifying the reliability of individual predictions, limiting their potential clinical utility. We propose a Multimodal Bayesian Diffusion Transformer (MM-DiT) framework that jointly estimates RP risk and characterizes the sources of predictive uncertainty. MM-DiT integrates planning computed tomography (CT) images and three-dimensional radiation dose distributions through self-supervised multimodal pre-training, reducing reliance on limited and potentially noisy toxicity labels. The resulting representations are further refined using a latent diffusion transformer and transferred to a Bayesian prediction framework for probabilistic RP risk estimation. A learnable label-noise model is incorporated to explicitly account for uncertainty arising from imperfect toxicity annotations. Therefore, it provides individualized RP risk estimates with complementary measures of aleatoric, epistemic, and label uncertainty, enabling assessment of prediction reliability at the individual-patient level. We evaluated MM-DiT in two independent cohorts using complementary assessments of predictive discrimination, calibration, and uncertainty. The results demonstrate its potential to provide accurate RP risk estimates while quantifying clinically relevant sources of predictive uncertainty.
1 Introduction
RP is a clinically significant thoracic-radiotherapy toxicity, but conventional DVH/NTCP models and most AI approaches incompletely represent its spatial, patient-specific, and uncertainty-related complexity. MM-DiT addresses this gap by combining multimodal learning with Bayesian risk estimation and uncertainty decomposition.
- Clinical motivation: RP can cause pulmonary morbidity, reduced quality of life, treatment interruption, and potentially life-threatening respiratory failure.Accurate pretreatment identification of high-risk patients could support risk stratification and individualized clinical management.
- Limitations of conventional assessment: DVH and NTCP approaches discard spatial information from 3D dose distributions and incompletely represent patient-specific and biological factors influencing RP.Relevant factors include pulmonary function, tumor location, and anatomical imaging features.
- AI-based prediction: AI methods can integrate demographics, clinical characteristics, radiation dose distributions, and CT-derived features to model heterogeneous RP-related information.Their multimodal integration is intended to capture interactions among anatomical, biological, and treatment-related factors.
- Uncertainty gap: Existing RP models generally provide risk estimates without explicitly distinguishing uncertainty from patient variability, limited model knowledge, or toxicity-assessment ambiguity.This limits information about whether an individual prediction is well supported.
- Proposed framework: MM-DiT uses self-supervised CT–dose representation learning, latent diffusion refinement, Bayesian prediction, and learnable label-noise modeling for personalized RP risk estimation.The framework produces calibrated probabilistic risk estimates with aleatoric, epistemic, and label uncertainty measures.
- Evaluation: MM-DiT was evaluated in two independent cohorts using discrimination, calibration, uncertainty quantification, and decision curve analysis.The study assesses whether uncertainty-aware estimates may provide greater net benefit across clinically relevant thresholds.
2.1 Patient cohorts
The framework was developed and evaluated using two independent patient cohorts with planning CT, corresponding 3D dose distributions, and binary RP-severity groups.
- Cohort definition: Patients were categorized as no-RP2 (RP grade ≤1) or RP2 (RP grade ≥2).Each patient record included a planning CT image set and corresponding 3D spatial dose distribution.
- Data sources: One cohort was derived from the NRG RTOG 0617 randomized clinical trial and accessed through The Cancer Imaging Archive.
- Study design: The proposed framework used data from two independent patient cohorts.
2.2 Data Preprocessing
The preprocessing pipeline standardized planning CT and dose volumes through intensity clipping, lung-mask-guided thoracic cropping, fixed-size resampling, normalization, and augmentation. Table 1 compares cohort characteristics, while Figure 1 illustrates the MM-DiT framework.
- Intensity preprocessing: CT intensities were clipped to (−1000, 400) HU and dose distributions to (0, 90) Gy to reduce outlier and noise influence.
- Spatial preprocessing: Both modalities were cropped to the thoracic foreground using their respective lung masks.
- Cohort characterization: Table 1 compares demographic, clinical, and tumor-biology characteristics across the two cohorts.Age is summarized by mean and [Min, Max], while categorical variables use absolute and relative frequencies (%).
- Standardization: The cropped volumes were resampled to 128 × 128 × 128 voxels and z-score normalized.
- Data augmentation: Augmentation included x/y mirroring, affine transformations, additive Gaussian noise, and CT brightness adjustments with specified probabilities.
- Framework overview: Figure 1 presents the MM-DiT framework for RP risk prediction and uncertainty quantification.
2.3 Model Development and Implementation
MM-DiT learns multimodal CT–dose representations without RP labels, refines them with latent diffusion, and transfers them to a Bayesian classifier for probabilistic RP prediction.
- Framework stages: The model uses self-supervised multimodal representation learning followed by supervised Bayesian fine-tuning.Pretraining uses planning CT images and 3D dose distributions without RP labels.
- Multimodal fusion: CT and dose modalities are encoded separately, fused in latent space with PoE or MoE, and processed by a Diffusion Transformer.
- Bayesian prediction: The Bayesian classifier uses MC Dropout to generate a distribution of RP-risk predictions, with the predictive mean serving as the final estimate.
- Probabilistic representation: The modality-specific posterior distributions are modeled as diagonal-covariance multivariate Gaussians and fused into a joint latent posterior.A deterministic 512-dimensional representation is obtained using the posterior mean.
- Fusion strategies: PoE combines unimodal posteriors proportionally to their product, whereas MoE forms a weighted mixture allowing independent modality contributions.
- Self-supervised objective: Self-supervised training combines reconstruction and denoising losses, with modality-specific decoders reconstructing CT and dose feature embeddings.
- Latent diffusion: The diffusion transformer operates in a compact 512-dimensional latent space, progressively denoises corrupted representations, and is trained by regressing injected noise.The forward process uses T = 1000 timesteps, while the reverse process predicts the denoising noise.
2.4 RP Risk Prediction and Model Evaluation
The framework generates probabilistic RP risk predictions from Bayesian stochastic inference and converts mean risk into binary RP classifications at a selected threshold. Evaluation emphasizes discrimination, accuracy, sensitivity, and specificity, with AUC prioritized under class imbalance.
- Probabilistic risk prediction: MC Dropout produces stochastic forward-pass predictions whose sigmoid-transformed logits yield probabilistic RP risks.The final prediction is based on the stochastic probability set and its predictive mean.
- Probabilistic risk prediction: The final RP risk prediction is defined as the predictive mean across T stochastic predictions.
- Binary classification: A decision threshold τ converts the mean probability into a binary RP-grade classification via an indicator function.
- Model evaluation: AUC is prioritized because extreme class imbalance makes its threshold-independent discrimination measure especially relevant.Accuracy, sensitivity, and specificity are also reported at the selected threshold.
2.5 Model Calibration Analysis
Calibration is assessed by comparing predicted RP probabilities with observed outcome frequencies across confidence bins. Expected Calibration Error summarizes the weighted bin-wise disagreement, with lower values indicating better calibration.
- Confidence binning: Predicted probabilities are partitioned into M confidence bins for calibration analysis.Each bin contains predictions grouped by confidence level.
- ECE computation: Expected Calibration Error quantifies agreement between predicted RP probabilities and observed outcome frequencies across bins.Its calculation uses the bin sizes and the absolute difference between bin accuracy and confidence.
- Bin statistics: For each bin, the mean predicted probability and empirical outcome frequency are calculated.The predicted probability uses p̂_i, while the observed frequency uses the corresponding ground-truth labels.
- ECE interpretation: Lower ECE values indicate better calibration, while ECE equal to zero denotes perfect agreement between predicted probabilities and actual event rates.N denotes the total number of samples in the ECE calculation.
2.6 Uncertainty Quantification and Decomposition
The framework decomposes RP predictive uncertainty into aleatoric, epistemic, and label components to distinguish intrinsic outcome variability, model limitations, and annotation inconsistency. These components support more transparent reliability assessment and can flag noisy labels or cases requiring cautious interpretation.
- Uncertainty decomposition: Total RP predictive uncertainty is decomposed into aleatoric, epistemic, and label uncertainty components.The decomposition is intended to isolate distinct underlying sources of uncertainty.
- Clinical interpretation: Separating uncertainty sources helps distinguish data noise, model limitations, and unreliable ground-truth annotations for prediction reliability assessment.The framework is described as robust to noisy clinical endpoints while flagging unreliable annotations.
- Aleatoric uncertainty: Aleatoric uncertainty captures irreducible stochasticity in the input data and identifies intrinsically ambiguous toxicity outcomes.It is quantified by averaging Bernoulli-probability variances across stochastic MC Dropout passes.
- Epistemic uncertainty: Epistemic uncertainty reflects parameter limitations caused by insufficient or non-representative training data.It is reducible with larger and more diverse training samples and may be elevated for out-of-distribution anatomical or treatment patterns.
- Label uncertainty: Label uncertainty models inconsistency and subjectivity in toxicity annotations, identifying cases where clinical grading may be ambiguous or observer-dependent.It is quantified through variance induced by the label-transition process using predicted risk and estimated noise rates.
3 Results
MM-DiT achieved strong RP prediction across two cohorts, with multimodal CNN+PoE configuration selected for subsequent calibration, clinical utility, and uncertainty analyses. Predictions were well calibrated, clinically useful across relevant thresholds, and accompanied by uncertainty patterns that identified borderline and cohort-dependent reliability.
- 3.1 Performance in Risk Outcome Prediction: CNN+PoE achieved the highest AUC across both cohorts and was selected as the final MM-DiT configuration.AUC was 85.59% in RTOG 0617 and 89.11% in MSHS.
- 3.2 Comparison with Other Methods: MM-DiT achieved the best overall predictive performance against multimodal VAE, CT-only, and dose-only baselines in both cohorts.AUC was 85.59% versus 75.27% for multimodal VAE in RTOG 0617, and 89.11% versus 76.43% in MSHS.
- 3.2 Comparison with Other Methods: Single-modality models had AUC values below 60% on both cohorts, whereas multimodal VAE and MM-DiT reached 75.27%/76.43% and 85.59%/89.11%, respectively.The paired values correspond to RTOG 0617 and MSHS.
- 3.2 Comparison with Other Methods: MM-DiT achieved accuracy of 90.87% in RTOG 0617 and 92.02% in MSHS, with high specificity and improved sensitivity versus the VAE baseline.Specificity was 96.34% and 97.74% in RTOG 0617 and MSHS, respectively.
- 3.3 Model Calibration and Reliability: ECE was 0.0330 in RTOG 0617 and 0.0198 in MSHS, with reliability curves closely following the perfect-calibration diagonal.These results indicate consistent probability calibration across the two independent cohorts.
- 3.4 Risk Stratification and Clinical Utility: MM-DiT maintained higher net benefit than treat-all and treat-none strategies across a broad range of clinically relevant threshold probabilities.Total predictive uncertainty was highest near predicted probability 0.5, while MSHS showed moderately higher uncertainty and relatively greater epistemic contribution than RTOG 0617.
4 Discussion
MM-DiT combines RP risk prediction with explicit uncertainty characterization, distinguishing predicted risk from confidence and decomposing uncertainty into aleatoric, epistemic, and label components. The framework shows clinical utility across decision thresholds, while its retrospective pretreatment inputs and analytically selected thresholds require further validation.
- Overall findings: MM-DiT demonstrated strong predictive discrimination and probability calibration across two independent cohorts while also quantifying individual-prediction reliability.The framework distinguishes what it predicts from how certain that prediction is.
- Uncertainty interpretation: Total uncertainty peaked near the decision boundary at probability 0.5 and decreased toward both probability extremes.Uncertainty was decomposed into epistemic, aleatoric, and label components.
- Uncertainty interpretation: Aleatoric uncertainty was dominant, particularly near the decision boundary, suggesting pretreatment CT and dose distributions cannot fully capture patient-specific biological variability.The discussion indicates that reducing this uncertainty may require toxicity-rich longitudinal information during and after treatment.
- Uncertainty interpretation: Greater epistemic uncertainty in the MSHS cohort may reflect differences in populations, imaging, treatment patterns, or other factors between development and evaluation cohorts.Elevated epistemic uncertainty can identify predictions with limited representation in training data that may warrant greater clinical scrutiny.
- Clinical interpretation: Patients were categorized into four risk–uncertainty groups using predicted RP probability and overall uncertainty, with 0.5 separating risk groups and the upper 30% defining high uncertainty.The framework supports different clinical interpretations for combinations of predicted risk and confidence.
- Clinical utility: Decision curve analysis showed positive net benefit across a broad range of threshold probabilities compared with treat-all and treat-none strategies.Prospective studies are still needed to determine whether uncertainty-aware assessment improves treatment selection, toxicity monitoring, or other clinical management.
- Limitations: The study was retrospective, used primarily pretreatment CT and dose distributions, and requires additional multi-institutional validation for generalizability.The four-group risk and uncertainty thresholds were selected for analytical demonstration and require clinical validation; decision-curve utility represents potential rather than prospective effectiveness.
5 Conclusion
MM-DiT integrates planning CT and 3D radiation dose distributions to predict individualized RP risk while explicitly quantifying multiple uncertainty sources. The framework extends conventional risk stratification toward uncertainty-aware clinical interpretation and supports prospective multi-institutional evaluation.
- Framework: MM-DiT integrates planning CT imaging and 3D radiation dose distributions for individualized RP risk prediction and explicit uncertainty quantification.It characterizes aleatoric, epistemic, and label uncertainty alongside RP risk.
- Implication: By combining risk estimation with prediction reliability, MM-DiT extends conventional RP risk stratification toward uncertainty-aware clinical interpretation.The authors present it as a transparent and reliable approach to AI-assisted toxicity assessment.
- Future evaluation: The findings provide a foundation for prospective multi-institutional evaluation of uncertainty-aware decision support in thoracic radiation oncology.