Deep Reinforcement Learning Classification of Brain Tumors on MRI

Abstract

We have recently shown that deep reinforcement learning can achieve high accuracy for lesion localization and segmentation even with minuscule training sets. Here, we introduce reinforcement learning for image classification, specifically binary classification of normal versus tumor-containing 2D MRI brain scans. We employed multi-step image classification via Deep Q learning with TD(0) environmental sampling. We trained on a set of 30 images (15 normal and 15 tumor-containing.) We tested on a separate set of 30 images (15 normal and 15 tumor-containing.) For comparison, we also trained and tested a supervised deep learning classification network on the same set of training and testing images. Whereas the supervised approach quickly overfit the small training set and, as expected, performed poorly on the testing set ($50\%$ accuracy, equivalent to random guessing), deep reinforcement learning achieved an accuracy of $100 \%$. The difference was statistically significant, with a p-value of $6.1 \times 10^{-5}$. Class activation maps for the Deep Q networks used in deep reinforcement learning highlight the lesions. In contrast, those of supervised deep learning’s convolutional neural networks show no focus of network attention. Hence, in this proof of principle work, we have shown not only that deep reinforcement learning is able to train effectively on very small data sets, but how it learns to classify, by focusing on the regions of greatest salience.