DEEP REINFORCED ACTIVE LEARNING FOR MULTI-CLASS IMAGE CLASSIFICATION

Abstract

High accuracy medical image classification can be limited by the costs of acquiring more data as well as the time and expertise needed to label existing images. In this paper, we apply active learning to medical image classification, a method which aims to maximise model performance on a minimal subset from a larger pool of data. We present a new active learning framework, based on deep reinforcement learning, to learn an active learning query strategy to label images based on predictions from a convolutional neural network. Our framework modifies the deep-Q network formulation, allowing us to pick data based additionally on geometric arguments in the latent space of the classifier, allowing for high accuracy multi-class classification in a batch-based active learning setting, enabling the agent to label datapoints that are both diverse and about which it is most uncertain. We apply our framework to two medical imaging datasets and compare with standard query strategies as well as the most recent reinforcement learning based active learning approach for image classification.

1. INTRODUCTION

Modern methods in machine learning (ML), including deep learning (DL) frameworks, require large amounts of labelled data to train sufficiently well to obtain high performance. Depending on the training task, these data can be very expensive to obtain or annotate, to the extent that traditional approaches become prohibitively costly. Active learning (AL) aims to alleviate this problem by adaptively selecting training samples with the highest value to construct a minimal training dataset with the most information for the ML model. In order to select training samples with the most information, different strategies are used in different AL cycles which can be either constructed based on knowledge of the specific problem one is aiming to learn, or using theoretical criteria to approximate mathematical bounds on information contained in the data. Standard query strategies in AL include the uncertainty-based approach (Lewis and Gale, 1994; Lewis and Catlett, 1994; Shannon, 1948; Scheffer et al., 2001; Esuli and Sebastiani, 2009; Seung et al., 1992; Dagan and Engelson, 1995) , which aim to quantify the model uncertainty about the samples to be selected using different hand-crafted heuristics. Other approaches aim to estimate the expected model change (Roy and Mccallum, 2001; Freytag et al., 2014) , or employ diversity-based approaches to promote diversity in sampling (Bilgic and Getoor, 2009; Gal et al., 2017; Nguyen and Smeulders, 2004) . Some approaches combine different techniques in hybrid-based query strategies, to take into account the uncertainty and diversity of query samples (Ash et al., 2019; Zhdanov, 2019; Shui et al., 2019; Beluch et al., 2018) . Other methods leverage the exploration-exploitation trade-off and reformulate the AL framework as a bandit problem (Hsu and Lin, 2015; Chu and Lin, 2016) or a reinforcement learning problem (Ebert et al., 2012; Long and Hua, 2015; Konyushkova et al., 2017a) , which are however still limited by their reliance on hand-crafted strategies, as opposed to learning a new one. The move towards combining deep learning methods with active learning, to combine the learning capability of the former in the context of high-dimensional data, with the data efficiency of the latter, have led to further methods development. However, combining the two is non-trivial; traditional active learning query strategies label samples one-by-one, and so batch-model deep active learning aims to use batch-based sample querying (Gal and Ghahramani, 2015; Gal et al., 2017; Kirsch et al., 2019; Cardoso et al., 2017) to ensure efficiency in sampling the data. 1

