Domain knowledge driven 3D dose prediction using moment-based loss function

Objective. To propose a novel moment-based loss function for predicting 3D dose distribution for the challenging conventional lung intensity modulated radiation therapy plans. The moment-based loss function is convex and differentiable and can easily incorporate clinical dose volume histogram (DVH) domain knowledge in any deep learning (DL) framework without computational overhead. Approach. We used a large dataset of 360 (240 for training, 50 for validation and 70 for testing) conventional lung patients with 2 Gy × 30 fractions to train the DL model using clinically treated plans at our institution. We trained a UNet like convolutional neural network architecture using computed tomography, planning target volume and organ-at-risk contours as input to infer corresponding voxel-wise 3D dose distribution. We evaluated three different loss functions: (1) the popular mean absolute error (MAE) loss, (2) the recently developed MAE + DVH loss, and (3) the proposed MAE + moments loss. The quality of the predictions was compared using different DVH metrics as well as dose-score and DVH-score, recently introduced by the AAPM knowledge-based planning grand challenge. Main results. Model with (MAE + moment) loss function outperformed the model with MAE loss by significantly improving the DVH-score (11%, p < 0.01) while having similar computational cost. It also outperformed the model trained with (MAE + DVH) by significantly improving the computational cost (48%) and the DVH-score (8%, p < 0.01). Significance. DVH metrics are widely accepted evaluation criteria in the clinic. However, incorporating them into the 3D dose prediction model is challenging due to their non-convexity and non-differentiability. Moments provide a mathematically rigorous and computationally efficient way to incorporate DVH information in any DL architecture. The code, pretrained models, docker container, and Google Colab project along with a sample dataset are available on our DoseRTX GitHub (https://github.com/nadeemlab/DoseRTX)