LENAS: Learning-based Neural Architecture Search and Ensemble for 3D Radiotherapy Dose Prediction

TL;DR

LENAS addresses the challenge of accurate, efficient 3D radiotherapy dose prediction by integrating neural architecture search with ensemble learning and knowledge distillation. It introduces U-NAS to automatically generate a diverse teacher ensemble and KDA-Net to distill this knowledge into a lightweight student, balancing model accuracy and inference cost. Empirical results on OpenKBP and AIMIS demonstrate state-of-the-art performance and strong generalization, along with insights into how diversity and ensemble strategy impact performance. The work provides practical guidelines for constructing and selecting diverse base learners in KBP, with potential to enhance clinical treatment planning.

Abstract

Radiation therapy treatment planning requires balancing the delivery of the target dose while sparing normal tissues, making it a complex process. To streamline the planning process and enhance its quality, there is a growing demand for knowledge-based planning (KBP). Ensemble learning has shown impressive power in various deep learning tasks, and it has great potential to improve the performance of KBP. However, the effectiveness of ensemble learning heavily depends on the diversity and individual accuracy of the base learners. Moreover, the complexity of model ensembles is a major concern, as it requires maintaining multiple models during inference, leading to increased computational cost and storage overhead. In this study, we propose a novel learning-based ensemble approach named LENAS, which integrates neural architecture search with knowledge distillation for 3D radiotherapy dose prediction. Our approach starts by exhaustively searching each block from an enormous architecture space to identify multiple architectures that exhibit promising performance and significant diversity. To mitigate the complexity introduced by the model ensemble, we adopt the teacher-student paradigm, leveraging the diverse outputs from multiple learned networks as supervisory signals to guide the training of the student network. Furthermore, to preserve high-level semantic information, we design a hybrid-loss to optimize the student network, enabling it to recover the knowledge embedded within the teacher networks. The proposed method has been evaluated on two public datasets, OpenKBP and AIMIS. Extensive experimental results demonstrate the effectiveness of our method and its superior performance to the state-of-the-art methods.

Figures (7)

• Figure 1: Overview of the proposed LENAS. $O_i$ in the hybrid module denotes the operation and $\alpha_i$ denotes its weight. $Dis$ in KDA-Net denotes the discriminator.
• Figure 2: (a) DVHs of the dose distribution of ground-truth plan (solid curves) and predictions by a single U-Net with and without the proposed KDA method, illustrated by dashed and dotted lines, respectively. The PTV70, PTV63, and PTV56 are shown in red, orange, and yellow lines, respectively. (b) An example of dose distributions of the clinical plan and predicted plans of the single U-Net with and without the KDA method.
• Figure 3: The dose score (MAE) of (a) 20 NAS models; (b) the ensembles.
• Figure 4: The dose score (MAE) of the ensembles with different numbers of models. The yellow bars indicate the models selected based on performance; the blue and green bars indicate the models selected based on diversity from the top 20 models and ten models, respectively.
• Figure 5: (a) Diversity and (b) dose score (MAE) of individual models in different ensemble strategies: NAS, bagging, random initializations and iterations, and different off-the-shelf architectures (denotes public).