Improved Outcome Models with Denoising Diffusion

Abstract

Purpose:

Radiotherapy outcome modelling often suffers from class imbalance in the modelled endpoints. One of the main options to address this issue is by introducing new synthetically generated datapoints, using generative models, such as Denoising Diffusion Probabilistic Models (DDPM). In this study, we implemented DDPM to improve performance of a tumor local control model, trained on imbalanced dataset, and compare this approach with other common techniques.

Methods:

A dataset of 535 NSCLC patients treated with SBRT (50 Gy/5 fractions) was used to train a deep learning outcome model for tumor local control prediction. The dataset included complete treatment planning data (planning CT images, 3D planning dose distribution and patient demographics) with sparsely distributed endpoints (6–7 % experiencing local failure). Consequently, we trained a novel conditional 3D DDPM model to generate synthetic treatment planning data. Synthetically generated treatment planning datapoints were used to supplement the real training dataset and the improvement in the model’s performance was studied. Obtained results were also compared to other common techniques for class imbalanced training, such as Oversampling, Undersampling, Augmentation, Class Weights, SMOTE and ADASYN.

Results:

Synthetic DDPM-generated data were visually trustworthy, with Fréchet inception distance (FID) below 50. Extending the training dataset with the synthetic data improved the model’s performance by more than 10%, while other techniques exhibited only about 4% improvement.

Conclusions:

DDPM introduces a novel approach to class-imbalanced outcome modelling problems. The model generates realistic synthetic radiotherapy planning data, with a strong potential to increase performance and robustness of outcome models.

1. Introduction

Over the last few years, machine learning models have dominated in the area of treatment outcome modeling thanks to their superior performance, interpretation, and possibility to combine a variety of input data (e.g., imaging data, treatment planning data, multi-omics data, patient demographics) [1–7]. Unfortunately, many modeled endpoints (e.g., local control, regional/distant recurrence or radiation toxicities) are often distributed sparsely, with resultant class imbalance in the dataset [8–12].

There are various techniques to overcome the class imbalance problem. Traditional approaches may include Undersampling, Oversampling, Augmentation, Class Weights, Synthetic Minority Oversampling Technique (SMOTE), Adaptive Synthetic Sampling Approach (ADASYN) and others [13–23]. More advanced techniques utilize a combination of traditional approaches with certain deep learning architectures, such as deepSMOTE [24] or SMOTified-GAN[25].

In recent years, there has been a rapid growth of denoising diffusion probabilistic models (DDPM). Since its introduction by Ho, et al. in 2020 [26], it has been applied in various applications, including generation of synthetic 2D medical images [27–31] and 3D ones [32]. Such synthetic medical images can be used to enhance classification models. As demonstrated in various studies [33–35], DDPM can effectively diminish problems arising from class imbalance or general lack of sufficient training data.

This study introduces a novel conditional 3D DDPM for multi-modal radiotherapy planning data, including patient demographics. We specifically investigate its potential to improve the performance of deep learning outcome models and compare it with other techniques to solve the class imbalance. Having a generative model of reliable and complete radiotherapy planning data might considerably increase performance and robustness of outcome models, and thus, help translating them into clinics, where they can contribute to more personalized treatment.

2. Material and Methods

2.1. Local control outcome model

A dataset of 535 non-small cell lung cancer (NSCLC) patients undergoing stereotactic body radiotherapy (SBRT) was used to develop a deep learning model for time-to-local recurrence prediction (DL-surv) [36]. The patients were treated between 2009 and 2017. Most of the patients were stage I (82 %), 16 % were stage II and 2 % were stage III. More dataset details are provided in Table 1. The maximum follow-up time was 5 years, after which tumors were considered locally controlled. The mean follow-up time was 28 months.

Table 1.

Clinical characteristics of the dataset

	Total number of patients (n = 535)	Number of LR patients (n = 31)
Gender
Male	279 (52.1 %)	20 (64.5 %)
Female	256 (47.9 %)	11 (35.5 %)
Age [years]
< 60	34 (6.40 %)	0 (0 %)
60 – 70	146 (27.3 %)	8 (25.8 %)
70 – 80	222 (41.5 %)	16 (51.6 %)
> 80	133 (24.9 %)	7 (22.6 %)
PTV volume [cm3]
< 25	185 (34.6 %)	6 (19.4 %)
25 – 50	187 (35.0 %)	9 (29.0 %)
50 – 100	115 (21.5 %)	10 (32.3 %)
> 100	48 (9.0 %)	6 (19.4 %)
Clinical maximum tumor diameter [mm]
≤ 10	53 (34.6 %)	1 (3.2 %)
11–20	240 (35.0 %)	12 (38.7 %)
21 – 30	149 (21.5 %)	11 (35.5 %)
≥ 31	93 (9.0 %)	7 (22.6 %)
Lobe
LLL	80 (15.0 %)	6 (19.4 %)
LUL	150 (28.0 %)	5 (16.1 %)
RLL	110 (20.6 %)	8 (25.8 %)
RML	5 (0.9 %)	0 (0 %)
RUL	190 (35.5 %)	12 (38.7 %)

The prescribed dose was 50 Gy in 5 fractions, i.e., 100 Gy of biologically effective dose (BED [α/β=10]). Planning requirements were 95 % of PTV and 100 % of IGTV to be covered by the prescribed dose. Treatments were delivered on Varian TrueBeam or Trilogy linear accelerators, using IMRT (32% of patients) or VMAT (68% of patients).

The architecture of the DL algorithm is depicted in Figure 1. The model predicts conditional probability of local control in discrete time intervals [6,37], using multi-modal planning data (CT images, 3D planning dose) and patient demographics (PTV size, clinical maximum tumor diameter, gender, age, and lung lobe tumor location). There are 3 parallel neural networks (NN). One 3D convolutional neural network (3D CNN) for feature extraction from planning dose distribution (Dose-CNN), one 3D CNN for feature extraction from planning CT images, and one variational autoencoder (VAE) for encoding patient demographic features (Demo-VAE). Extracted features from all 3 NN are then concatenated and fed into the discrete-time survival NN (Surv-net) [6,37], which predicts probability of local control in specified time intervals. In this study, the time intervals were 0–1, 1–2, 2–3 and 3–5 years. The optimized hyper parameters were learning rate=5e-4, batch size=32, weight decay=0.05, dropout=0.2.

Figure 1:

Diagram of DL-surv model including architecture details. Abbreviations: CNN = convolutional neural network; VAE = variational autoencoder; LR = local recurrence.

The model was cross-validated (CV) on 80% of the data, using stratified 5-fold approach with 50 iterations in total. The remaining 20% of data was withheld for independent testing. The data split was done according to TRIPOD criteria 2b [38]. The CV Harrell’s concordance index (c-index) was 0.72 with 95% confidence interval (CI) 0.68–0.75, and the testing c-index was 0.69. One of the main drawbacks, holding back the model’s performance, was severe class imbalance. There were only 31 patients (5.8 %) experiencing local recurrence (LR), while the rest was either LR-free or censored before the maximum follow-up point. Therefore, we implemented and compared various techniques to solve the class imbalance problem.

2.2. Traditional approaches to solve the class imbalance

Multiple traditional techniques [39] and their combinations were implemented and evaluated. Among the most basic traditional techniques were: Undersampling, Oversampling, Augmentation (rotation, resizing, flipping) and Class Weights. With respect to the Class Weights technique, the most promising ratio of class weights was found to be 1:5 (w_major=1, w_minor=5). We then applied more advanced techniques, such as SMOTE and ADASYN.

2.3. Denoising Diffusion Probabilistic Model

The conditional DDPM for generation of synthetic NSCLC SBRT data (DDPM_NSCLC) was developed in accordance with the original paper by Ho, et al. [26] and its improved version published by Nichol, et al. [40]. Figure 2 shows diagram of the model. It was trained on the same dataset as the DL-surv model, keeping the withheld testing subset untouched. The input of the model consists of 3D treatment planning data (CT images, planning dose and PTV mask) and 3 embedded conditions - sex, class and time-to-event bin. Consequently, it is possible to generate synthetic samples specifically according to the class and time-to-event, which is necessary for generation of new samples to balance the original dataset. The remaining demographic details were calculated from generated samples (PTV volume, clinical maximum tumor diameter), and the lung lobe was imputed according to the distribution in the original dataset.

Figure 2:

Diagram of the DDPM model used for the generation of synthetic radiotherapy planning data.

Synthetic samples were evaluated utilizing Fréchet Inception Distance (FID) [41], which is a common technique to evaluate the quality of images from generative neural networks, such as VAE, GAN or DDPM. It is a measure that can compare feature representations of real and generated images. The features are extracted from a pre-trained InceptionV3 NN and their distributions are compared using Fréchet Distance. Lower FID values indicate higher feature similarity between real and synthetic images, thus, more realistic image generation.

The model was built in PyTorch 1.13.1 [42]. The optimization algorithm used in this study was Adam [43], and the loss function was mean squared error. The model was optimized for batch size (bs), learning rate (lr), weight decay (w_d) and variance schedule in the diffusion process (β_min, β_max, n_steps). The optimized values were bs=16, lr=0.0005, w_d=0, β_min=0.0001, β_max=0.02 and n_steps=1000.

Subsequently, 300 synthetically generated NSCLC SBRT data supplemented the real training data, which improved the class imbalance ratio from about 1:16 to 2:3.

We also tried building conditional Generative Adversarial Network (GAN) for the same purpose. None of the implemented GAN models (DCGAN, WGAN, WGAN-GP) provided acceptable results, as all of them failed exhibiting a mode collapse. This was presumably caused by the limited number of local failures datapoints, i.e., limited feature space for conditional GAN training.

3. Results

The comparison of real and DDPM generated data is provided in Figure 3. It shows examples of 4 real and 4 synthetic data samples, including CT images, planning dose distributions and PTV masks, which qualitatively appear to be satisfactory. The FID values of synthetic CT, Dose and PTV samples were 41.7, 27.0 and 59.4, respectively. The FIDs were calculated as mean values from all inceptionV3 feature layers. The reason for not using only the last feature layer is the low number of samples in the dataset, which limits the use of FID, as it strongly depends on the number of samples.

Figure 3:

Examples of synthetic and real treatment planning data (CT images, planning dose distribution and PTV mask). Abbreviation: FID = Fréchet Inception Distance.

Table 2 provides a comparison of DL-surv performance for 8 different class imbalance solutions. The DDPM approach outperformed the other competing methods, with a testing c-index of 0.75. We used this approach to generate 300 additional synthetic patients to reduce the class imbalance for training the DL-surv model.

Table 2:

Cross-validation and testing performance of DL-surv model for different techniques to solve the class imbalance.

Class-balancing technique	CV c-index (95% CI)	Testing c-index
Original dataset	0.72 (0.68–0.75)	0.66
Undersampling	0.70 (0.66–0.74)	0.66
Oversampling	0.66 (0.62–0.70)	0.67
Augmentation	0.70 (0.67–0.74)	0.67
Class weights (1:5)	0.72 (0.69–0.74)	0.69
Augmentation + class weights	0.68 (0.64–0.72)	0.67
SMOTE	0.66 (0.63–0.69)	0.69
ADASYN	0.67 (0.64–0.70)	0.69
DDPMNSCLC	0.74 (0.71–0.77)	0.75

4. Discussion

This work has implemented a conditional DDPM for simultaneous generation of realistic CT images, 3D dose distributions, and PTV masks, along with selected patient demographic details. As can be seen from Figure 3 the images are trustworthy and the FID values for CT, Dose and PTV mask are within acceptable range and are consistent with previous publications [27,44].

Traditional approaches used to solve the class imbalance were mostly unsuccessful (see Table 2). Undersampling and Oversampling reached testing c-indexes of 0.66 and 0.67 respectively, which is about the same performance compared to the original dataset. The reason is data scarcity in the minority class. The whole dataset contains only 31 patients with LR; thus, there are only 25 (80%) of them in the CV subset. While Undersampling and Oversampling help to balance training data batches, they do not necessarily help with the model performance, as the training data might still be inadequately sparse for reasonable training [45]. One must confront this problem in highly imbalanced datasets with a low number of minority class samples, such as this. Similarly, the Class Weights technique exhibited only slightly better testing c-index (0.69) than the original dataset. Class Weights aims to adapt the loss function according to the severity of the class imbalance. It can reduce class imbalance but does not tackle the problem of insufficient training datapoints. On the other hand, augmentation is promising not only for the class imbalance issue, but also for the lack of minority class samples, as it generates new data points. However, it did not improve the testing c-index (0.69) by much. The reason for the poor performance here is presumably due to the homogeneity of the data. All patients were positioned identically (headfirst - supine). Therefore, rotating or flipping data would not improve the model’s performance, since it only introduces features that are not normally present in the data.

ADASYN and SMOTE showed similar results. The testing c-index was 0.69, which is only a slight improvement. This is probably due to the inappropriateness of these techniques for complex imaging multi-modal data. Even though there were previously successful attempts to apply these techniques in imaging data [46–48], they were never applied in more complex data, such as radiotherapy planning data. Originally, SMOTE/ADASYN were designed for tabular data [13,14]. Hence, they might underperform when using complex multi-modal imaging data.

As shown in Table 2, DL-surv performance was considerably enhanced after supplementing the original training dataset by synthetic DDPM_NSCLC generated data. The testing c-index increased from 0.66 to 0.75, which is more than 10% improvement. DDPM is generating synthetic data, which are visually entirely new, although they have the same distribution of underlying features. Consequently, DDPM generated data do not necessarily extend the feature space covered by the training samples but do significantly increase its sampling. This approach is helpful in this particular class imbalance problem and makes the DDPM approach superior to other implemented class imbalance techniques. However, the reason(s) for limited performance in any outcome modelling should always be carefully assessed. Usually, there is an interplay of multiple factors, and the main limiting factor can often be different than the class imbalance issue.

In conclusion, denoising diffusion probabilistic models are highly promising generative models for complex multi-modal imaging data. This work presents DDPM_NCSLC, which is a novel generative model that can be used to generate complex synthetic radiotherapy planning data, including all necessary components. The generated data look very realistic and show potential to increase performance of multi-modality outcome models trained on data with a high event sparsity. In the presented tumor local control model, with the class imbalance ratio of about 1:16, the performance improved by more than 10%. Other implemented techniques (Undersampling, Oversampling, Augmentation, Class Weights, SMOTE, ADASYN) did not provide more than ~4% improvement.

• DDPM generating treatment planning data was developed and verified (FID<50)

• Synthetic data helped to improve performance of SBRT outcome model by >10%

• Other techniques for imbalanced training improved performance by less than 4%

• DDPM model can diminish the negative effect of event imbalance in outcome models.

• Introduced approach can help in clinical translation of radiotherapy outcome models