Knowledge-based automated planning with three-dimensional generative adversarial networks

A GAN consists of two neural networks known as a generator and discriminator. [Goodfellow et al., 2014] We focus on conditional GANs which, in addition to a Gaussian input, learn to generate different outputs conditioned on known problem-specific characteristics (e.g., CT images) [Isola et al., 2017]. Specifically, let $\mathbf{z}\sim p_{\mathbf{z}}$ denote a sample from a Gaussian input. The generator network takes as input $\mathbf{z}$ and a CT image $\mathbf{c}$ and outputs a predicted dose distribution $\mathbf{x}=G(\mathbf{z},\mathbf{c})$. The discriminator network then takes a CT image and the predicted dose distribution as input and outputs a “belief” regarding whether the dose distribution is the actual clinical dose (as opposed to artificially produced by the generator). That is, $D(\mathbf{x},\mathbf{c})\in[0,1]$ where $D(\mathbf{x},\mathbf{c})=1$ suggests the discriminator is confident the dose distribution is the clinically delivered dose. Both networks are trained iteratively with a single loss function $\mathcal{L}(D,G)$. Letting $\mathbf{x}\sim p_{\text{data}}$ denote the distribution of real delivered plans, we write the training problem as:

where

The above formulation represents the objective for the most common class of conditional GANs used in problems associated with image generation [Isola et al., 2017, Zhu et al., 2017]. By minimizing the first term in $\mathcal{L}(D,G)$, the generator learns to construct dose distributions such that $D(G(\mathbf{z},\mathbf{c}),\mathbf{c})=1$. That is, the generator attempts to fool the discriminator into believing that the generated dose is from the real clinical distribution. The discriminator adversarially maximizes the second term in $\mathcal{L}(D,G)$ to output $D(G(\mathbf{z},\mathbf{c}),\mathbf{c})=0$ for $\mathbf{z}\sim p_{\mathbf{z}}$ and $D(\mathbf{x},\mathbf{c})=1$ for $\mathbf{x}\sim p_{\text{data}}$, i.e., the discriminator attempts to correctly distinguish between artificially generated versus clinically delivered plans. The final term in $\mathcal{L}(D,G)$ is an $l_{1}$ loss function which forces the generated samples to better resemble the ground truth dataset (i.e., the clinically delivered dose distribution). The hyperparameter $\lambda$ balances the tradeoff between minimizing the GAN loss (first two terms) and having images resemble deliverable plans.

We constructed generator and discriminator networks derived from the pix2pix architecture proposed in the canonical Style Transfer GAN [Isola et al., 2017]. The generator possesses a U-net architecture that passes a contoured CT image through consecutive convolution layers, a bottleneck layer, and then several deconvolution layers. The U-net employs skip connections, i.e., the output of each convolution layer is concatenated to the input of a corresponding deconvolution layer. This allows the generator to easily pass “high-dimensional” information (e.g., structural outlines) between the inputted CT image and the outputted dose. The discriminator takes as input a dose distribution and CT image and passes them through several consecutive convolution layers until outputting a single scalar value between zero and one. We refer the reader to the original Style Transfer GAN work for full details on the number and size of the convolution and deconvolution layers in the pix2pix architecture [Isola et al., 2017].

We trained three GAN variants each with slightly modified architectures; we refer to them as 2D-RGB, 2D-dose, and 3D-dose. The “2D” or “3D” designation refers to whether dose is being predicted for 2D slices individually or the full 3D image at once. The “-RGB” or “-dose” designation refers to the input/output channels of the GAN, i.e., whether the output of the generator is a 3-dimensional vector of numbers representing the red-green-blue scale, or a 1-dimensional scalar that represents dose directly (i.e., grayscale). The 2D-RGB GAN is a direct replica of the original pix2pix network and has been previously implemented and tested for KBAP [Mahmood et al., 2018]. It takes as input a single axial slice of the contoured CT image, i.e., a $128\times 128$ pixel image with three channels for color, and outputs a three channel $128\times 128$ pixel image that represents dose. Two-dimensional convolution and deconvolution filters are used for the generator and discriminator. The 2D-dose GAN has an almost identical architecture except the output of the generator is a one channel $128\times 128$ image. In both 2D GANs, a 3D dose distribution is obtained by concatenating all outputted axial slices for a single patient. In contrast, the 3D-dose GAN takes an entire contoured CT image ($128\times 128\times 128$ voxels) and generate the corresponding 3D dose distribution. Thus, it uses three-dimensional convolution and deconvolution filters [Hermoza and Sipiran, 2017] with a one channel output to represent dose. Apart from filter dimensions, the 2D-dose and 3D-dose GANs have the same architecture including the same number of layers and filter sizes.

The generator and discriminator networks were trained iteratively using gradient descent. After training was complete, the discriminator was disconnected and the generator was used on out-of-sample CT images. Figure 2 summarizes the difference between how the GANs were trained and tested. In our experiments, we used the loss function given by $\mathcal{L}(D,G)$ with $\lambda=90$, and trained the networks using the Adam optimizer [Kingma and Ba, 2014] with learning rate $\alpha=0.0002$ and $\beta_{1}=0.5$ and $\beta_{2}=0.999$. These hyperparameters are the default Adam settings and have been used in many style transfer problems [Isola et al., 2017]. We also performed a parameter sweep for $\lambda$ and found the default setting sufficient. We stopped training when the standard GAN loss function $\mathbb{E}_{\mathbf{x}\sim p_{\text{data}}}\left[\log D(\mathbf{x},\mathbf{c})\right]+\mathbb{E}_{\mathbf{z}\sim p_{\mathbf{z}}}\left[\log(1-D(G(\mathbf{z},\mathbf{c}),\mathbf{c}))\right]$ was roughly equal to the regularization term $\lambda\mathbb{E}_{\mathbf{x}\sim p_{\text{data}},\mathbf{z}\sim p_{\mathbf{z}}}\left[\|\mathbf{x}-G(\mathbf{z},\mathbf{c})\|_{1}\right]$; if the loss from the $l_{1}$ penalty fell too low, the GAN began to memorize the dataset. Therefore, the 2D-dose and 2D-RGB GANs were trained for $50$ epochs and the 3D-dose GAN was trained for $200$ epochs. It is natural for the 3D-dose GAN to require more epochs to converge as it contains significantly more parameters and generates the entire dose distribution. The code for all experiments is provided at http://github.com/rafidrm/gancer.

We obtained plans for 217 oropharyngeal cancer treatments delivered at a single institution with 6 MV, step-and-shoot, intensity-modulated radiation therapy. Plans were prescribed 70 Gy and 56 Gy in 35 fractions to the gross disease (PTV70) and elective target volumes (PTV56), respectively; in 130 plans there was also a prescription of 63 Gy to the intermediate risk target volume (PTV63). The organs-at-risk (OARs) included the brainstem, spinal cord, right parotid, left parotid, larynx, esophagus, mandible, and the limPostNeck, which is an artificial structure for sparing the posterior neck. Patient geometry was discretized into voxels of size $4$ mm $\times$ $4$ mm $\times$ $2$ mm.

The CT images and dose distributions for all 217 treatment plans were converted into a suitable format for use by the GAN architectures. The CT images were reconstructed so that each voxel had RGB channels, which were assigned values according to Table 1, and converted into 128 axial slices of $128\times 128$ voxels. The images were separated into a training set of 130 random samples (a total of 15,657 pairs of CT image slices and dose distributions) and a testing set of the remaining 87 samples for out-of-sample evaluation.

During out-of-sample testing, predictions produced by the generator were used as input into an optimization model with two stages. In the first stage, given a predicted dose distribution, the objective weights for a standard inverse planning model were estimated using a parameter estimation technique that has been previously validated in oropharynx [Babier et al., 2018b]. In the second stage, the estimated weights were used in an inverse planning optimization model to generate treatment plans. The objective minimized the sum of 65 functions: seven per OAR and three per target. In our experiments, objectives for the OARs included the mean dose, max dose, and the average dose above 0.25, 0.50, 0.75, 0.90, and 0.975 of the maximum predicted dose to the OAR. Objectives for the target included the maximum dose, average dose below prescription, and average dose above prescription. The complexity of all generated treatment plans was constrained to a sum-of-positive-gradients (SPG) value of 55 [Craft et al., 2007]. SPG was used since it is a convex surrogate for the physical deliverability of a plan and the parameter 55 was chosen as it is two standard deviations above the average SPG [Babier et al., 2018a]. The dose influence matrices required for the optimization model were derived with A Computational Environment for Radiotherapy Research [Deasy et al., 2003]. Each of the KBP-generated plans were delivered from nine equidistant coplanar beams at angles 0^∘, 40^∘, …, 320^∘. We used Gurobi 7.5 to solve the optimization model.

We also generated plans using a scaling procedure that multiplicatively increased the entire predicted dose distribution by the smallest amount to satisfy all target dose criteria. The scaled predictions were then input into the optimization model. Note that this scaling does not affect the fairness of the analysis because the final KBAP plans must satisfy the same constraints (e.g., SPG) as plans generated from the unscaled predictions. To study the full effect of the scaling step, we generated three additional populations of scaled final plans alongside the initial three unscaled plans corresponding to 2D-RGB, 2D-dose, and 3D-dose. We refer to the three scaled KBAP plans as 2D-RGB^′, 2D-dose^′, and 3D-dose^′ respectively.

We conducted two primary analyses. First, we evaluated the quality of the KBAP plans by computing the fraction of clinical planning criteria that were satisfied. We compared these results against the performance of the clinical plans. Second, we evaluated the quality of the KBP predictions by comparing the predicted dose distributions to clinical dose distributions, and the KBAP plan DVHs to their respective predicted DVHs.

KBAP Plan Quality: The quality of the final KBAP plans were measured by evaluating how often they satisfied the clinical criteria presented in Table 2. For each class of clinical criteria, i.e., OARs, targets, and all regions-of-interest (ROIs), which includes both OARs and targets, we generated confusion matrices to compare the KBAP plans with the clinical plans. We then analyzed the organ-specific criteria that was achieved by each clinical plan to determine whether the corresponding KBAP plan also passed that criterion.

KBP Prediction Quality: We measured KBP prediction quality by evaluating how similar the KBP predictions were to their corresponding KBAP plans. For every ROI and every patient, we calculated the average absolute error between the KBP predicted DVH and the KBAP plan DVH. Differences between 2D-RGB and 2D-dose, 2D-dose and 2D-dose^′, and 2D-dose^′ and 3D-dose^′ were evaluated; each comparison attempts to isolate the impact of each of the three KBAP-specific modifications. We then evaluated the difference between 2D-RGB and 3D-dose^′ to quantify the joint contribution of all three modifications. The distribution of the errors over the population of plans was visualized using a separate boxplot for each set of comparisons. Finally, we visualized the predicted dose distributions to determine whether the 3D model produced smoother predictions across the longitudinal axis than a 2D model.

In Table LABEL:clinCritConfusion, we present confusion matrices comparing the quality of the final KBAP plans with the clinical plans. The columns represent KBAP performance using each of the six KBP approaches while the rows indicate the clinical plans and the performance targets. Overall, 3D-dose^′ plans best replicated the performance of the clinical plans since they agreed most on what criteria passed and failed (i.e., Pass/Pass and Fail/Fail). Where they differed, 3D-dose^′ plans satisfied five times as many criteria (Fail/Pass) as the clinical plans (Pass/Fail). We also observed that scaling made a substantial difference as scaled plans outperformed their unscaled counterparts. For example, scaled 3D plans satisfied 99.5% of all target criteria whereas only 52.3% were satisfied for unscaled 3D plans. Finally, 2D-dose^′ and 3D-dose^′ performed the best satisfying 77.0% and 76.4% of all ROI criteria, respectively.

Table LABEL:clinCrit summarizes the performance of the KBAP plans, focusing only on the criteria that the corresponding clinical plans also passed. Both 3D-dose and 3D-dose^′ performed the best across all OAR and target criteria. In particular, 3D-dose^′ achieved the highest passing rate for all target criteria. It also achieved the highest passing rate for all OAR criteria except the larynx and mandible. However, for some regions such as the brainstem, esophagus, and PTV63, multiple KBAP approaches yielded the top result.

In Table LABEL:clinCritPlans, we recorded the proportion of KBAP plans that satisfied all of the same criteria as the corresponding clinical plans. Note that the subtle difference between these results and the results from Table LABEL:clinCrit is that here we only record a pass if the KBAP plan meets every single criteria that the clinical plan meets. Table LABEL:clinCrit, on the other hand, considers each criterion independently. We found that plans generated from the scaled predictions performed better than their unscaled counterparts in terms of satisfying all ROI criteria; we observed the biggest improvement between 3D-dose and 3D-dose^′ (34.5 percentage points). Like all scaled plans, the improvement of 3D-dose^′, which satisfied the most target criteria (98.9%), was the result of much better target coverage at the expense of less OAR sparing compared to the 3D-dose plans, which satisfied the most OAR criteria (94.3%). Dose encoded as a grayscale image led to an increase in plan quality best shown by the difference of 27.6 percentage points separating 2D-RGB^′ and 2D-dose^′. Finally, it was clear that the 3D GAN architecture resulted in large improvements in prediction quality as compared to the 2D approach. That is, 3D-dose^′ satisfied criteria more frequently (by 2.3 percentage points) than 2D-dose^′. Overall, 3D-dose^′ achieved the same criteria as the clinical plans in 78.2% of cases, which was more than any other approach.

In Figure 3, we present box plots comparing the effect of each modification on prediction error, i.e., the average absolute error between the KBP predicted DVH and the KBAP plan DVH. Each subplot represents a pairwise comparison of the prediction error between KBAP pipelines with different GAN architectures for prediction. The larger the skew towards one side, the more accurate that GAN architecture is. By comparing 2D-RGB to 2D-dose, plotted in Figure 3(a), we found that dose encoded as a scalar value resulted in predictions that were closer (median error difference of 0.19 Gy) to deliverable plans as compared to predictions where dose was encoded as a RGB color map. We also found that the prediction errors were lower among scaled plans than unscaled plans (Figure 3(b)). That is, by comparing 2D-dose and 2D-dose^′, we found that half of all predictions generated via scaling were about 0.31 Gy more accurate than those generated without; the PTV56 was the only structure where the scaled KBAP plans were less accurate (average median error of 1.00 Gy). For both modifications, although the median differences were modest, the upper limit of the differences could be upwards of 2 Gy. Additionally, we found that the median error of 3D-dose^′ was 0.07 Gy lower than 2D-dose^′ (Figure 3(c)). Finally, in Figure 3(d), we evaluated the difference between our best set of KBAP plans, 3D-dose^′, and the plans generated by the previous state-of-the-art in GAN-based KBAP, 2D-RGB. The median error of 3D-dose^′ was 0.63 Gy smaller than 2D-RGB, which suggests that the predictions made by 3D-dose^′ were closer to the final dose distributions than those generated by 2D-RGB.

Figure 4 presents predicted dose distributions corresponding to the 2D-dose and a 3D-dose KBP methods. Note that the two models predicted distinct dose distributions that differed along the vertical (i.e., longitudinal) axis. In particular, the 3D GAN better learned the vertical relationship between adjacent axial slices and thus, generated smoother dose predictions across the longitudinal axis. In contrast, the 2D GAN predictions included more “streaky” and unrealistic discontinuities in the dose distribution, particularly around axial slices that were adjacent to the target boundaries. For example, the dose falloff was impossibly steep in the plans generated using 2D predictions (Figure 4(c)), while plans generated using 3D predictions had more realistic dose gradients (Figure 4(d)).