A Multi-Armed Bandit to Smartly Select a Training Set from Big Medical Data

TL;DR

This paper introduces a multi-armed bandit approach using Thompson sampling to efficiently select training data from large medical image datasets, improving prediction accuracy while reducing data requirements.

Contribution

It presents a novel method for training set selection in medical imaging using a bandit model that relies solely on meta information, not image features.

Findings

Achieved higher age prediction accuracy from brain MRI data.

Required only a fraction of the total data for training.

Validated on 7,250 subjects across 10 datasets.

Abstract

With the availability of big medical image data, the selection of an adequate training set is becoming more important to address the heterogeneity of different datasets. Simply including all the data does not only incur high processing costs but can even harm the prediction. We formulate the smart and efficient selection of a training dataset from big medical image data as a multi-armed bandit problem, solved by Thompson sampling. Our method assumes that image features are not available at the time of the selection of the samples, and therefore relies only on meta information associated with the images. Our strategy simultaneously exploits data sources with high chances of yielding useful samples and explores new data regions. For our evaluation, we focus on the application of estimating the age from a brain MRI. Our results on 7,250 subjects from 10 datasets show that our approach leads to higher accuracy while only requiring a fraction of the training data.

Figures

Equations

$$P({\pi }_i|{\alpha }_i,{\beta }_i)=\frac{\Gamma ({\alpha }_i+{\beta }_i)}{\Gamma ({\alpha }_i)\Gamma ({\beta }_i)}(1-{\pi }_i{)}^{{\beta }_i-1}{\pi }_i^{{\alpha }_i-1}.$$

Full text

1 Introduction

Machine learning has been one of the driving forces for the huge progress in medical imaging analysis over the last years. Of key importance for learning-based techniques is the training dataset that is used for estimating the model parameters. Traditionally, medical data has been scarce so usually all available data for a particular task was used for training. Nowadays, many initiatives make data publicly available so that huge amounts of data can potentially be used to estimate more accurate models. However, just including all the data in the training set is becoming increasingly impractical, since processing the data to create training models can be very time consuming on huge datasets. In addition, most processing may be unnecessary because it does not help the model estimation for a given task. In this work, we propose a method to select a subset of the data for training that is most relevant for a specific task. Foreshadowing some of our results, such a guided selection of a subset of the data for training can lead to a higher performance than using all the available data while requiring only a fraction of the processing time.

The task of selecting a subset of data for training is challenging because at the time of making the decision, we do not yet have processed the data and we do therefore not know how the inclusion of the sample would affect the prediction. However, in many scenarios each image is assigned metadata about the patient (sex, diagnosis,age etc.) or the image acquisition (dataset of origin, location, imaging device, etc.). We hypothetize that some of this information can be useful to guide the selection of samples but it is a priori not clear which information is most relevant and also how it should be distributed. To address this problem, we formulate the selection of the samples to be included in a training set as reinforcement learning, where a trade-off must be reached between the exploration of new sources of data and the exploitation of sources that have been shown to lead to informative data points in the past. More specifically, we model this as a multi-armed bandit problem solved with Thompson sampling, where each arm of the bandit corresponds to a cluster of samples generated using meta information.

In this paper, we apply our sample selection method to brain age estimation [7] from MR1 T1 images. The estimated age serves as a proxy for biological age, whose difference to the chronological age can be used as indicator of disease [6, 8]. The age estimation is a well-suited application for testing our algorithm as it allows us to work with a large number of datasets since the subject’s age is one of the few variables that is included in every neuroimaging dataset.

1.1 Related Work

Our work is mostly related to active learning approaches, whose aim is to select samples to be labeled out of a pool of unlabeled data. Examples of active learning approaches applied to medical imaging tasks include the work by Hoi et al. [10], where a batch mode active learning approach was presented for selecting medical images for manually labeling the image category. Another active learning approach was proposed for the selection of histopathological slices for manual annotation in [21]. The problem was formulated as constrained submodular optimization problem and solved with a greedy algorithm. To select a diverse set of slices, the patient identity was used as meta information. From a methodological point of view, our work relates to the work of Bouneffouf et al. [1], where an active learning strategy based on contextual multi-armed bandits is proposed. The main difference between all these active learning approaches and our method is that image features are not available a priori in our application, and therefore can not be used in the sample selection process. Our work also relates to domain adaptation [16]. In instance weighting, the training samples are assigned weights according to the distribution of the labels (class imbalance) [11] and the distribution of the observations (covariate shift) [17]. Again these methods are not directly applicable in our scenario because not all the distribution of the metadata is defined on the target dataset.

2 Method

2.1 Incremental Sample Selection

In supervised learning, we model a predictive function ${f}:({\bf x},{\mathbf{p}})\mapsto y$ depending on a parameter vector ${\mathbf{p}}$, relating an observation $\mathbf{x}$ to its label $y$. In our application, ${\bf x}\in\mathbb{R}^{m}$ is a vector with $m$ quantitative brain measurements from the image and $y\in\mathbb{R}$ is the age of the subject. The parameters ${\mathbf{p}}$ are estimated by using a training set ${S^{T}}=\{{s}_{1},{s}_{2},\ldots,{s}_{N_{train}}\}$, where each sample ${s}=({\mathbf{x}},{y})$ is a pair of a feature vector and its associated true label. Once the parameters are estimated, we can predict the label $\tilde{y}$ for a new observation $\tilde{{\bf x}}$ with $\tilde{y}=f(\tilde{{\bf x}},{\mathbf{p}}^{*})$, where the prediction depends on the estimated parameters and therefore the training dataset.

In our scenario, the samples to be included in the training set ${S^{T}}$ are selected from a large source set ${S}=\{h_{1},h_{2},..,h_{N_{total}}\}$ containing hidden samples of the form $h=\{\hat{{\mathbf{x}}},\hat{{y}},\mathbf{m}\}$. Each $h$ contains hidden features $\hat{{\mathbf{x}}}$ and label $\hat{{y}}$ that can only be revealed after processing the sample. In addition, each hidden sample possesses a $d$-dimensional vector of metadata $\mathbf{m}\in\mathbb{Z}^{d}$ that encodes characteristics of the patient or the image such as sex, diagnosis, and dataset of origin. In contrast to $\hat{{\mathbf{x}}}$ and $\hat{{y}}$, $\mathbf{m}$ is known a priori and can be observed at no cost. To include a sample $h$ from set ${S}$ into ${S^{T}}$, first its features and labels have to be revealed, which comes at a high cost. Consequently, we would like to find a sampling strategy that minimizes the cost by selecting only the most relevant samples according to the metadata $\mathbf{m}$.

2.2 Multiple Partitions of the Source Data

In order to guide our sample selection algorithm, we create multiple partitions of the source dataset, where each one considers different information from the metadata $\mathbf{m}$. Considering the $j$-th meta information ($1\leq j\leq d$), we create the $j$-th partition ${S}=\cup_{i=1}^{\eta_{j}}C_{i}^{j}$ with $\eta_{j}$ a predefined number of bins for $\mathbf{m}[j]$. As a concrete example, sex could be used for partitioning the data, so ${S}=C_{\text{female}}^{\text{sex}}\cup C_{\text{male}}^{\text{sex}}$ and $\eta_{\text{sex}}=2$. All the clusters generated using different meta information are merged into a set of clusters $\mathcal{C}=\{C_{\iota}^{j}\}$. We hypothesize that given this partitioning, there exist clusters $C_{i}\in\mathcal{C}$ that contain more relevant samples than others for a specific task. Intuitively, we would like to draw samples $h$ from clusters with a higher probability of returning a relevant sample. However, since the relationship between the metadata and the task is uncertain, the utility of each cluster for a specific task is unknown beforehand. We will now describe a strategy that simultaneously explores the clusters to find out which ones contain more relevant information and exploits them by extracting as many samples from relevant clusters as possible.

2.3 Sample selection as a multi-armed bandit problem

We model the task of sequential sample selection as a multi-armed bandit problem. At each iteration ${t}$, a new sample is added to the training dataset $S^{T}$. For adding a sample, the algorithm decides which cluster $C_{i}\in\mathcal{C}$ to exploit and randomly draws a training sample ${s}_{t}$ from cluster $C_{i}$. The corresponding feature vector ${\mathbf{x}}_{t}$ and label ${y}_{t}$ are revealed and the usefulness of the sample ${s}_{t}$ for the given task is evaluated, yielding a reward ${r}_{t}\in\mathbb{R}$. A reward $r_{t}=1$ is given if adding the sample improves the prediction accuracy of the model and $r_{t}=-1$ otherwise.

At $t=0$, we do not possess knowledge about the utility of any cluster. However, this knowledge is incrementally built as more and more samples are drawn and their rewards are revealed. To this end, each cluster is assigned a distribution of rewards ${\Pi}_{i}$. With every sample the distribution better approximates the true expected reward of the cluster, but every new sample also incurs a cost. Therefore, a strategy needs to be designed that explores the distribution for each of the clusters, while at the same time exploiting as often as possible the most rewarding sources.

To solve the problem of selecting from which $C_{i}$ to sample at every iteration $t$, we follow a strategy based on Thompson sampling [18] with binary rewards. In this setting, the expected rewards are modeled using a probability $P_{i}$ following a Bernoulli distribution with parameter $\pi_{i}\in[0,1]$. We maintain an estimate of the likelihood of each $\pi_{i}$ given the number of successes $\alpha_{i}$ and failures $\beta_{i}$ observed for the cluster ${C}_{i}$ so far. Successes ($r=1$) and failures ($r=-1$) are defined based on the reward of the current iteration. It can be shown that this likelihood follows the conjugate distribution of a Bernoulli law, i.e., a Beta distribution $Beta(\alpha_{i},\beta_{i})$ so that

[TABLE]

with the gamma function $\Gamma$. At each iteration, $\hat{\pi}_{i}$ is drawn from each cluster distribution $P_{i}$ and the cluster with the maximum $\hat{\pi}_{i}$ is chosen. The procedure is summarized in Algorithm 1.

3 Results

In order to showcase the advantages of our multi-armed bandit sampling algorithm (MABS), we present an evaluation of our method in the task of estimating the biological age of a subject given a set of volume and thickness features of the brain. We choose this task in particular because of the big number of available brain images in public databases and, the relevance of age estimation as a tool for diagnostic of neuro degenerative diseases [8, 19]. For predicting the age, we reconstruct brain scans with FreeSurfer [5] and extract volume and thickness measurements to create our feature vectors ${\mathbf{x}}$. Based on these features, we train a regression model for predicting the age of previously unseen subjects.

3.1 Data

We work on MRI T1 brain scans from 10 large-scale public datasets: ABIDE [3], ADHD200 [15], AIBL [4], COBRE [14], IXI111http://brain-development.org/ixi-dataset/, GSP [2], HCP [20], MCIC [9], PPMI [13] and OASIS [12]. From all these datasets we obtain a total number of 7,250 images, which is to the best of our knowledge, the largest dataset collected for the task of age prediction. Since each one of these datasets is targeted towards different applications, the selected population is heterogeneous in terms of age, sex, and health status. Images are processed with FreeSurfer [5] and thickness and volume measurements extracted. Even though this is a fully automatic tool, the extraction of the feature is a computationally intensive task which is by far the bottleneck of our age prediction regression model.

3.2 Age estimation

We perform age estimation on two different testing scenarios. In the first, we create a testing dataset by randomly selecting subsets from all the datasets. The aim of this experiment is to show that our method is capable of selecting samples that will create a model that can generalize well to a heterogeneous population. In the second scenario, the testing dataset corresponds to a single dataset. In this scenario, we show that the sample selection permits tailoring the training dataset to a specific target dataset.

Experiment 1. For the first experiment we take all the images in the dataset and we divide them randomly into three sets: 1) a small validation set of 2% of all samples to compute the rewards given to MABS , 2) a large testing set of 48% to measure the performance of our age regression task, and 3) a large hidden training set of 50%, from which samples are taken sequentially using MABS. We perform the sequential sample selection described in Algorithm 1 using the following metadata to construct the clusters $\mathcal{C}$: age, dataset, diagnosis, and sex. We experiment with considering all of the metadata separately, to investigate the importance of each one, and the joint modeling. Since we require to evaluate a regression model every time a sample is included, we opted to use ridge regression as our learning algorithm. Rewards ${r}$ are given to each bandit estimating and observing if the $r^{2}$ score of the prediction in the validation set increases. It is important to emphasize that the testing set is not observed by the bandits in the process of giving rewards. Every experiment is repeated 20 times using different random splits and the mean results are shown. We compare with two baselines: the first one (RANDOM) consists of obtaining samples at random from the hidden set and adding them sequentially to the training set. As a second baseline (AGE PRIOR), we add samples sequentially by following the age distribution of the testing set . The results of this first experiment are shown in figure 1 (top left). In almost all the cases, using MABS as a selection strategy performed better than the baselines. It is important to observe that an increase in performance is obtained not only when the relationships between the metadata and the task are direct, like in the case of the clusters constructed by age, but also when this relationship is not clear like in the case of clustering the images using only dataset or diagnostic information. Another important aspect is that even when the meta information is not informative, like in the case of the clusters generated by sex, the prediction using MABS is not affected.

Experiment 2. For our second experiment, we perform age estimation with the test data being a specific dataset. This experiment follows the same methodology as the previous one with the important difference of how the datasets are split. This time the split is done by choosing: 1) a small validation set, taken only from the target dataset, 2) a testing set, which corresponds to the remaining samples in the target dataset not included in the validation set, and 3) a hidden dataset containing all the samples from the remaining datasets. The goal of this experiment is to show that our methodology can be deployed when samples have to be selected according to a specific population and prediction task. From the results in figure 1, we observe that depending on the dataset, using bandits with only one specific source of meta information can actually improve the sample selection algorithm. However, the best meta information for a particular task is different in every case. We also observe that in general the MABS using all available meta information extracts informative samples more efficiently than our baselines and always close to the best performing single meta information MABS. This reinforces our hypothesis that it is hard to define an a priori relationship between the meta data and the task, and is therefore a better strategy to let MABS select from multiple sources of meta information at once.

4 Conclusion

We have proposed a method for efficiently and intelligently sampling a training dataset from a large pool of data. The problem was formulated as reinforcement learning, where the training dataset was sequentially built after evaluating a reward function at every step. Concretely, we used a multi-armed bandit model that was solved with Thompson sampling. The intelligent selection considered metadata of the scan to construct a distribution about the expected reward of a training sample. Our results showed that the selective sampling approach leads to higher accuracy than using all the data, while requiring less time for processing the data. We demonstrated that our technique can either be used to build a general model or to adapt to a specific target dataset, depending on the composition of the test dataset. Since our method does not require to observe the information contained in the images, it can also be applied to predict useful samples even before the images are acquired, guiding the recruitment of subjects.

5 Acknowledgements

This work was partly funded by SAP SE, the Faculty of Medicine at the University Munich (Förderprogramm für Forschung & Lehre), and the Bavarian State Ministry of Education, Science and the Arts in the framework of the Centre Digitisation.Bavaria (ZD.B).