Unsupervised deep clustering and reinforcement learning can accurately segment MRI brain tumors with very small training sets

Abstract

Purpose Lesion segmentation in medical imaging is key to evaluating treatment response. We have recently shown that reinforcement learning can be applied to radiological images for lesion localization. Furthermore, we demonstrated that reinforcement learning addresses important limitations of supervised deep learning; namely, it can eliminate the requirement for large amounts of annotated training data and can provide valuable intuition lacking in supervised approaches. However, we did not address the fundamental task of lesion/structure-of-interest segmentation. Here we introduce a method combining unsupervised deep learning clustering with reinforcement learning to segment brain lesions on MRI.Materials and Methods We initially clustered images using unsupervised deep learning clustering to generate candidate lesion masks for each MRI image. The user then selected the best mask for each of 10 training images. We then trained a reinforcement learning algorithm to select the masks. We tested the corresponding trained deep Q network on a separate testing set of 10 images. For comparison, we also trained and tested a Unet supervised deep learning network on the same set of training/testing images.Results Whereas the supervised approach quickly overfit the training data and predictably performed poorly on the testing set (16% average Dice score), the unsupervised deep clustering and reinforcement learning achieved an average Dice score of 83%.Conclusion We have demonstrated a proof-of-principle application of unsupervised deep clustering and reinforcement learning to segment brain tumors. The approach represents human-allied AI that requires minimal input from the radiologist without the need for hand-traced annotation.

Deep reinforcement learning with automated label extraction from clinical reports accurately classifies 3D MRI brain volumes

“…In order to predict the actions our agent will take, we use the Deep-Q Network (DQN), as we have in prior work [9,10,11,12]. The architecture of our DQN is illustrated in Figure 1.…”

“…We should note at this point that selecting argmax a (Q) is an "on-policy" action selection. As we will see, and have described in prior work [9,10,11,12], we also need to experiment with random action selections to explore and learn about the environment. This is called off-policy behavior.…”

“…As the agent learns about the environment and the best policy to follow, the degree of randomness decreases and the agent increasingly chooses the optimal action predicted by the DQN, i.e., the argmax a {Q(a = 0), Q(a = 1)}. This is called the epsilon-greedy algorithm, and it is described in more detail in our prior RL works [9,10,11,12]. We use similar parameters regarding this part of the approach:…”

“…We can store all of these values in a tuple called a transition. As described in our prior RL work [9,10,11,12], we store these up to a maximum memory size called the replay memory buffer. Here, in order to conserve space and be able to save more transitions in the replay memory buffer, we store state information solely based on the training image number and correctness of the prior and just-made predictions, pred_corr t−1 and pred_corr t , respectively.…”

Deep Neuroevolution Squeezes More out of Small Neural Networks and Small Training Sets: Sample Application to MRI Brain Sequence Classification

“…As such, DRL may help to lead us past the above-described impasse in Radiology AI. In fact, early applications of DRL have shown remarkable ability to generalize based on small training sets [16,20,19,17,18].…”