Patient Aware Active Learning for Fine-Grained OCT Classification

Active learning is a branch of human-in-the-loop computing that assumes the model and unlabeled dataset will evolve over time as the most informative samples are selected from a dataset. This application applied to medical image analysis is practical as it reduces the costly and cumbersome human effort needed for expert manual annotation. Even though, active learning on medical imagery has demonstrated progress in overcoming the reliance on large, balanced datasets [1, 2], the practical deployment of this paradigm still has major challenges.

First, most conventional active learning paradigms applied to medical image analysis are still largely considered as “black box” algorithms. This inhibits the health care professionals from interpreting, understanding and even correcting predictions a model has made [3, 4, 5]. From a physician’s point-of-view, this creates a higher risk of patient harm. Second, a single disease can present itself in visually diverse formats across multiple patients. This is exemplified in Fig. 1 where, although each patient’s optical coherence tomography (OCT) is diagnosed with diabetic macular edema (DME), the visual characteristics of them are unalike. In other words, data diversity is captured within the medical meta data and remains un-exploited in existing active learning paradigms. Additionally, medical datasets are widely imbalanced both across classes and patients [6]. This ultimately results in existing methods training models without properly accounting for one or more patient disease manifestations from whom there is less data. For this purpose, the model could overfit to a minority patient group and fail to generalize to the broader masses. For practical deployments, this represents a major risk.

From a machine learning perspective, the root cause of the problem resides in the structure of the representation space. Conventional active learning methods rely on a well established representation of the data. If samples are conventionally chosen from an unbalanced dataset, the structure of the representation space will be underdeveloped and the model will generalize poorly. For this purpose, we propose to shape the representation space through medical meta-information. Specifically, we provide data diversity through patient identity. By imposing a patient constraint on the query strategies, we structure the representation space based on a medically grounded prior (Fig. 2). We show that our method is robust to both unbalanced data distributions and architectures and is also effective at fine-grained disease classification. Based on the original paradigm, we define our method as Patient Aware Active Learning. In summary, the contributions of this paper are as follows:

1. i

   We incorporate patient information to augment the active learning paradigm and improve medical interpretability.

2. ii

   We develop a modular plug-in method that can be applied to arbitrary active learning strategies that match or outperform existing methods in terms of generalization.

3. iii

   We test our method on a popular OCT benchmark with two different architectures and five different query strategies. This amounts to 100 experiments.

Active learning frameworks strive to select the most informative subset of samples from an unlabeled dataset in order to combat the laborious, costly and time-consuming nature of developing large annotated datasets [7, 8, 9, 10, 11]. In relation to medical image analysis, active learning has been applied to histopathological image analysis [1], skin lesion segmentation [12], heart magnetic resonance imaging and computerized tomography scan analysis [13], the diagnosis of digital mammograms [14], and for tuberculosis detection of chest radiography [15]. A significant area of active learning research is focused on defining sample informativeness. In this context, several approaches define informativeness with generalization difficulty [16, 17] while others focus on data diversity within the acquisition batch [9, 10]. Even though diversity is considered within the context of learned model features, existing approaches do not account for diversity inherent within the medical data. To the best of our knowledge, our approach represents the first to consider medical meta-information within the active learning workflow.

In a modular context, our method integrates within existing query strategies and enhances their performance with medical meta-information. Patient awareness is injected into the learning framework by partitioning the unlabeled training pool on patient identity prior to ranking sample informativeness with a query function. It is a modular method that integrates with arbitrary active learning strategies.

The dataset used in this paper, obtained from [6], has grayscale, cross-sectional, foveal OCT scans of varying sizes belonging to an unbalanced distribution of a healthy class and three types of retinal diseases: Drusen, choroidal neovascularization (CNV) and diabetic macular edema (DME). We are interested in distinguishing between the disease states and thus use only imagery from these three classes for fine-grained classification. Sample imagery from each class is shown in Fig. 4. A total of 10488 DME, 36345 CNV and 7756 Drusen images from 1852 patients were used in the training and unlabeled set. Within the test set, there were 250 images for each disease collected from 486 patients. All imagery were resized to $128\times 128$ and normalized to have zero mean and unit standard deviation. Additional implementation details are shown in Table 1. There was no overlap in patients or imagery in train or test sets. The patient and class distribution in train and test sets are shown in Fig. 5.

We use Resnet-18 [18] and Densenet-121 [19] architectures trained from scratch to classify images in an active learning framework. A learning rate of 1.5e-4 was used along with the Adam [20] optimizer. During each training round, the Resnet and Densenet models were trained for as many epochs it took to achieve 98% and 94% accuracy respectively. This is repeated with five different random seeds and an average accuracy is aggregated.

Classification performance is evaluated on the following baseline algorithms:

1. 1.

   Random: Randomly sampling $k$ instances at each round as a naive baseline.

2. 2.

   Least Confidence: Selecting the samples for which the model has the lowest predicted probability [17].

   $\displaystyle x^{*}_{LC}=\operatorname*{arg\,max}_{x}1-P_{\theta}(\hat{y}|x)$

   (2)

3. 3.

   Margin: The difference between the top two most probable predictions, $\hat{y_{1}}$ and $\hat{y_{2}}$, is used to identify the most informative sample [17].

   $\displaystyle x^{*}_{M}=\operatorname*{arg\,min}_{x}P_{\theta}(\hat{y_{1}}|x)-P_{\theta}(\hat{y_{2}}|x)$

   (3)

4. 4.

   Entropy: Uses the full distribution to determine the most uncertain prediction from a model [17].

   $\displaystyle x^{*}_{E}=\operatorname*{arg\,max}_{x}\sum_{i}P(y_{i}|x)\log{P(y_{i}|x)}$

   (4)

5. 5.

   BADGE: Samples disparate and high magnitude points from a gradient space at every round [21].

In Figs. 6 and 7 we show learning curves showing the accuracy of baseline query strategies compared to the corresponding patient aware query. In the interest of space, we visualize results for four of five query strategies on Resnet-18 and Densenet-121 architectures. In these plots, x-axis corresponds to the number of samples in the train set at that round and y-axis corresponds to the performance accuracy on the test set. Each colored curve is the average of five trials using different seeds, with standard errors being shown by the shaded regions. Our patient aware learning applied to the baselines is shown as the green curve in all sub-figures while baseline queries are shown in orange. The blue curve represents random sampling which serves as the naive baseline for all experiments. All results were achieved using 5.5% (3000 samples) of the OCT dataset (54589 samples).

The learning curves provide intuition about having medical insights guide the decisions of neural networks. For instance, most plots show patient aware learning surpassing the baselines on both architectures. This means that in majority cases, querying the dataset with patient-level insights in an active learning setup allows the model to better characterise disease states. Furthermore, we often see patient-aware learning having an edge over baseline algorithms from the early rounds of training onward like in Figs. 6(b) and 7(d). This is because the role of patient awareness plays a critical part in incorporating diversity throughout the training set. The model is thus exposed to several manifestations of a pathology from the onset. Noticeably too, patient aware learning can sometimes perform the same as baseline strategies, as shown in Fig. 7(b). This can happen when the selection of diverse patients do not contain significantly different disease manifestations within the imagery for the model to perform much different from the baselines. This is also possible when the model in some trials is poorly initialized yielding a poor set of built-in assumptions to make predictions. Overall, these findings suggest that patient aware learning is a good choice for medical applications of active learning.

In this paper we introduced a framework that incorporates clinical context in the form of patient awareness into an active learning framework that is medically interpretable. We show that our methods, in a plug-in sense, can be used in a modular context with arbitrary sampling strategies. We perform controlled experiments to validate the effectiveness of patient aware active learning for fine-gained OCT classification on different architectures using an unbalanced dataset. Based on our contributions in this paper, we will further investigate learning methods to evaluate the robustness of patient aware active learning.