Attention‐aware 3D U‐Net convolutional neural network for knowledge‐based planning 3D dose distribution prediction of head‐and‐neck cancer

PURPOSE: Deep learning–based knowledge‐based planning (KBP) methods have been introduced for radiotherapy dose distribution prediction to reduce the planning time and maintain consistent high‐quality plans. This paper presents a novel KBP model using an attention‐gating mechanism and a three‐dimensional (3D) U‐Net for intensity‐modulated radiation therapy (IMRT) 3D dose distribution prediction in head‐and‐neck cancer. METHODS: A total of 340 head‐and‐neck cancer plans, representing the OpenKBP—2020 AAPM Grand Challenge data set, were used in this study. All patients were treated with the IMRT technique and a dose prescription of 70 Gy. The data set was randomly divided into 64%/16%/20% as training/validation/testing cohorts. An attention‐gated 3D U‐Net architecture model was developed to predict full 3D dose distribution. The developed model was trained using the mean‐squared error loss function, Adam optimization algorithm, a learning rate of 0.001, 120 epochs, and batch size of 4. In addition, a baseline U‐Net model was also similarly trained for comparison. The model performance was evaluated on the testing data set by comparing the generated dose distributions against the ground‐truth dose distributions using dose statistics and clinical dosimetric indices. Its performance was also compared to the baseline model and the reported results of other deep learning‐based dose prediction models. RESULTS: The proposed attention‐gated 3D U‐Net model showed high capability in accurately predicting 3D dose distributions that closely replicated the ground‐truth dose distributions of 68 plans in the test set. The average value of the mean absolute dose error was 2.972 ± 1.220 Gy (vs. 2.920 ± 1.476 Gy for a baseline U‐Net) in the brainstem, 4.243 ± 1.791 Gy (vs. 4.530 ± 2.295 Gy for a baseline U‐Net) in the left parotid, 4.622 ± 1.975 Gy (vs. 4.223 ± 1.816 Gy for a baseline U‐Net) in the right parotid, 3.346 ± 1.198 Gy (vs. 2.958 ± 0.888 Gy for a baseline U‐Net) in the spinal cord, 6.582 ± 3.748 Gy (vs. 5.114 ± 2.098 Gy for a baseline U‐Net) in the esophagus, 4.756 ± 1.560 Gy (vs. 4.992 ± 2.030 Gy for a baseline U‐Net) in the mandible, 4.501 ± 1.784 Gy (vs. 4.925 ± 2.347 Gy for a baseline U‐Net) in the larynx, 2.494 ± 0.953 Gy (vs. 2.648 ± 1.247 Gy for a baseline U‐Net) in the PTV_70, and 2.432 ± 2.272 Gy (vs. 2.811 ± 2.896 Gy for a baseline U‐Net) in the body contour. The average difference in predicting the D (99) value for the targets (PTV_70, PTV_63, and PTV_56) was 2.50 ± 1.77 Gy. For the organs at risk, the average difference in predicting the [Formula: see text] (brainstem, spinal cord, and mandible) and [Formula: see text] (left parotid, right parotid, esophagus, and larynx) values was 1.43 ± 1.01 and 2.44 ± 1.73 Gy, respectively. The average value of the homogeneity index was 7.99 ± 1.45 for the predicted plans versus 5.74 ± 2.95 for the ground‐truth plans, whereas the average value of the conformity index was 0.63 ± 0.17 for the predicted plans versus 0.89 ± 0.19 for the ground‐truth plans. The proposed model needs less than 5 s to predict a full 3D dose distribution of 64 × 64 × 64 voxels for a new patient that is sufficient for real‐time applications. CONCLUSIONS: The attention‐gated 3D U‐Net model demonstrated a capability in predicting accurate 3D dose distributions for head‐and‐neck IMRT plans with consistent quality. The prediction performance of the proposed model was overall superior to a baseline standard U‐Net model, and it was also competitive to the performance of the best state‐of‐the‐art dose prediction method reported in the literature. The proposed model could be used to obtain dose distributions for decision‐making before planning, quality assurance of planning, and guiding‐automated planning for improved plan consistency, quality, and planning efficiency.