Unsupervised Task Design to Meta-Train Medical Image Classifiers

TL;DR

This paper introduces an unsupervised method for designing classification tasks to improve meta-training of medical image classifiers, reducing reliance on costly hand-designed tasks and achieving competitive results.

Contribution

It presents a novel unsupervised task design approach for meta-training medical image classifiers, eliminating the need for hand-crafted tasks.

Findings

Outperforms other unsupervised pre-training methods.

Achieves results comparable to meta-training with hand-designed tasks.

Improves classification accuracy after fine-tuning.

Abstract

Meta-training has been empirically demonstrated to be the most effective pre-training method for few-shot learning of medical image classifiers (i.e., classifiers modeled with small training sets). However, the effectiveness of meta-training relies on the availability of a reasonable number of hand-designed classification tasks, which are costly to obtain, and consequently rarely available. In this paper, we propose a new method to unsupervisedly design a large number of classification tasks to meta-train medical image classifiers. We evaluate our method on a breast dynamically contrast enhanced magnetic resonance imaging (DCE-MRI) data set that has been used to benchmark few-shot training methods of medical image classifiers. Our results show that the proposed unsupervised task design to meta-train medical image classifiers builds a pre-trained model that, after fine-tuning, produces better classification results than other unsupervised and supervised pre-training methods, and competitive results with respect to meta-training that relies on hand-designed classification tasks.

Tables

Table 1. Table 1: AUC results ( ± plus-or-minus \pm standard error) for breast screening baselines.

Training Method Baseline	AUC
From Scratch [19]	$0.83\pm 0.04$
MIL based fine-tuning [7]	$0.85\pm 0.04$
Multi-Task  [8]	$0.85\pm 0.04$
Variational Autoencoder + Nearest Neighbour	$0.61\pm 0.06$
Variational Autoencoder + Fine-Tune in breast screening	$0.84\pm 0.04$
Deep Clustering + Nearest Neighbour	$0.53\pm 0.06$
Deep Clustering + Fine-Tune in breast screening	$0.80\pm 0.05$

Table 2. Table 2: AUC for the breast screening task for our proposed method (U-MT) as a function of the number of image clusters K 𝐾 K and the task sampling method (random and curriculum). We also present the results of supervised meta-training  [ 4 ] (S-MT) for the cases where the breast screening is included (labelled as S) and not included (labelled as NS) in the meta-training tasks. N/A indicates that the experiment is not feasible due to the lack of extra ground truth labels.

	Random			Curriculum
	$K=3$	$K=4$	$K=5$	$K=3$	$K=4$	$K=5$
S-MT [4] (S)	$0.86\pm 0.04$	N/A	N/A	$0.90\pm 0.04$	N/A	N/A
S-MT [4] (NS)	$0.85\pm 0.04$	N/A	N/A	$0.89\pm 0.04$	N/A	N/A
U-MT (Ours)	$0.81\pm 0.05$	$0.88\pm 0.04$	$0.89\pm 0.04$	$0.87\pm 0.04$	$0.86\pm 0.04$	$0.88\pm 0.04$

Equations

$$\underset{\theta ,\omega }{\min }\frac{1}{|\mathcal{D}|}\sum _{i=1}^{|\mathcal{D}|}\ell (g_{\omega }(f_{\theta }({\mathbf{v}}_i)),{\tilde{\mathbf{y}}}_i),$$

$$\underset{\mathbf{C}}{\min }\frac{1}{|\mathcal{D}|}\sum _{i=1}^{|\mathcal{D}|}\underset{{\tilde{\mathbf{y}}}_i}{\min }\Vert f_{\theta }({\mathbf{v}}_i)-\mathbf{C}{\tilde{\mathbf{y}}}_i{\Vert }_2^2,$$

$$\kappa =\frac{1}{|\mathcal{D}|}\sum _{i=1}^{|\mathcal{D}|}\frac{b(i)-a(i)}{\max \left(a(i),b(i)\right)},$$

$$\underset{\psi }{\max }\log p({\mathcal{Y}}_{l=\mathrm{1..}T}^{(v)}|{\mathcal{Y}}_{l=\mathrm{1..}T}^{(t)},{\mathcal{V}}_{l=\mathrm{1..}T}^{(v)},{\mathcal{V}}_{l=\mathrm{1..}T}^{(t)},\psi ),$$

$$\log p({\mathcal{Y}}_{l=\mathrm{1..}T}^{(v)}|{\mathcal{Y}}_{l=\mathrm{1..}T}^{(t)},{\mathcal{V}}_{l=\mathrm{1..}T}^{(v)},{\mathcal{V}}_{l=\mathrm{1..}T}^{(t)},\psi )$$

$${\varphi }_l^{\ast }=\psi -\alpha {\nabla }_{{\varphi }_l}\left[-\log p({\mathcal{Y}}_l^{(t)}|,{\mathcal{V}}_l^{(t)},{\varphi }_l)\right],$$

Full text

1 Introduction

The accuracy and robustness of deep learning based medical image classifiers is generally positively correlated with the size of the annotated training set used during the modelling process [1]. However, large annotated training sets are expensive and not readily available for some medical image analysis applications, such as breast screening from DCE-MRI [2]. Therefore, training medical image classifiers with small annotated training sets has become a highly investigated topic, particularly after the advent of deep learning [1].

The most competitive medical image classifiers are currently based on convolutional neural networks (CNNs) [1] that need large training sets to be properly modelled. To reduce the need for such large annotated sets, pre-training approaches have been explored in medical image analysis, where the most relevant for our paper are: 1) supervised pre-training using independent data sets [5], where the model is pre-trained by solving a classification problem in a different data set; 2) unsupervised pre-training using clustering [3], where the model is pre-trained by performing clustering without any knowledge about the ground truth labels; and 3) unsupervised pre-training using input reconstruction [6], where the model is pre-trained by reconstructing the input images of the training set. Arguably, the main issue with these pre-training methods is that their objective functions are irrelevant for the medical image classifier being developed downstream. Alternatively, the need for pre-training methods can be alleviated with the use of other types of training methods, such as multiple instance learning (MIL) [7] or multi-task learning [8], but both methods still need large training sets. More recently, a pre-trained model produced by supervised meta-training (i.e., a meta-training process that depends on hand-designed classification tasks) showed superior performance compared to the previously described pre-training methods [4]. Nevertheless, these promising meta-training results are counterbalanced by an unappealing need of an expensive hand-designing process to produce the classification tasks [4]. Given the high cost of this process, the availability of a large number of hand-designed classification tasks is rare, which hampers the exploration of meta-training for medical image classifiers.

In this paper, we propose a new method to unsupervisedly produce a large number classification tasks to meta-train medical image classifiers. To this end, we use deep clustering [3] to automatically build image clusters that can be grouped in different ways to enable the design of multiple classification tasks employed in the meta-training process – see Fig. 1. We evaluate our method on the breast screening classification task from a breast DCE-MRI data set that has been used to benchmark few-shot training algorithms of medical image classifiers [4]. Results show that our proposed approach produces classification results that are significantly better than other unsupervised and supervised pre-training methods, and competitive to supervised meta-training.

2 Literature Review

DCE-MRI is a recommended image modality in breast screening programs for patients at high-risk [9]. However, DCE-MRI interpretation is time-consuming and prone to high inter-observer variability [10]. Thus, computer-aided diagnosis (CAD) systems are being developed to assist radiologists increase their diagnosis sensitivity [11] and specificity [12], and reduce analysis time. However, the development of CAD systems for breast DCE-MRI is challenging due in part to the small size of annotated data sets available for training.

Meta-training has been shown to be an effective strategy to improve the learning of classifiers using relatively small training sets [13]. For instance, Maicas et al. [4] proposed the use of hand-designed breast classification tasks to meta-train a model that was then fine-tuned to solve the breast screening task. Results showed that this method improves over other strategies to train classifiers from small data sets, such as MIL [7] and multi-task learning [8]. However, the method proposed in [4] relies on costly hand-designed classification tasks.

Similarly to our paper, Hsu et al. [14] proposed an unsupervised method to design computer vision classification tasks for meta-training. Results showed that this approach produced worse classification performance than meta-training modelled with hand-designed tasks (i.e., supervised meta-training). We believe that the reason behind this drop in performance lies in the large number of hand-designed tasks already available for supervised meta-training in computer vision applications [14], enabling a good classification performance baseline. The difficulty to obtain a large number of hand-designed tasks for medical image classification problems means that the number of these hand-designed tasks will be small, which may result in a relatively low classification performance baseline. We hypothesize that our proposed method that unsupervisedly designs a large number of classification tasks to meta-train a medical image classifier can achieve a classification performance that is at least comparable to supervised meta-training [4] trained with a small number of hand-designed tasks. Our proposed method has the advantage that it does not rely on costly hand-designed tasks.

3 Data Set and Methods

3.1 Data set

The data set is represented by ${\cal D}=\left\{\left(\mathbf{v}_{i},\mathbf{t}_{i},b_{i},y_{i}\right)\right\}_{i=1}^{|{\cal D}|}$, where $\mathbf{v}:\Omega\rightarrow\mathbb{R}$ corresponds to the first DCE-MRI subtraction volume ($\Omega$ denotes the volume lattice) [15], $\mathbf{t}:\Omega\rightarrow\mathbb{R}$ represents the T1-weighted MRI only used to separate the left and breast regions of the volume, $b\in\{\text{left},\text{right}\}$ indicates the left or right breast, and $y\in\mathcal{Y}=\{0,1\}$ indicates the classification label: no malignant findings, or malignant findings, respectively.

3.2 Deep Clustering to Unsupervisedly Design Classification Tasks

The proposed unsupervised task design method builds several binary classification tasks from image groups formed by deep clustering [3]. The training of deep clustering alternates an optimisation of two objective functions [3]. We denote the $\theta$-parameterised model that produces the unsupervised learning features by $f_{\theta}(\mathbf{v})\in\mathbb{R}^{D}$ and the $\omega$-parameterised classifier that produces a pseudo-label representing one of the unknown $K$ classes and is placed on top of $f_{\theta}(.)$ by $g_{\omega}(f_{\theta}(\mathbf{v}))\in\{0,1\}^{K}$. The first objective function is the cross-entropy loss $\ell(.)$ with respect to the pseudo-labels $\{\widetilde{\mathbf{y}}_{i}\}_{i=1}^{|\mathcal{D}|}$, with $\widetilde{\mathbf{y}}\in\widetilde{\mathcal{Y}}=\{0,1\}^{K}$,

[TABLE]

which is used to estimate the optimal $\theta^{*}$ and $\omega^{*}$. The second objective function finds the $K$ centroids, denoted by $\mathbf{C}\in\mathbb{R}^{D\times K}$, and pseudo-labels $\widetilde{\mathbf{y}}$ with

[TABLE]

where $\widetilde{\mathbf{y}}_{i}$ is a $K$-dim one-hot vector.

Each step of the optimization above will generate new values for the model parameters, centroids and pseudo-labels. We extend deep clustering [3] with a model selection process based on maximising the Silhouette coefficient that measures clustering quality [16] with

[TABLE]

where $a(i)$ represents the average $\ell_{2}$ distance between $f_{\theta}(\mathbf{v}_{i})$ and all points $f_{\theta}(\mathbf{v}_{j})$ where $i\neq j$ and $\widetilde{\mathbf{y}}_{i}=\widetilde{\mathbf{y}}_{j}$; and $b(i)$ denotes the smallest average $\ell_{2}$ distance between $f_{\theta}(\mathbf{v}_{i})$ and $f_{\theta}(\mathbf{v}_{j})$ where $i\neq j$ and $\widetilde{\mathbf{y}}_{i}\neq\widetilde{\mathbf{y}}_{j}$.

The unsupervised design of classification tasks is based on the formation of $L$ binary classification problems derived from the pseudo-labels obtained from (2). Each of these $L$ binary classification problems is built by randomly selecting 2 nonempty and disjoint subsets $\mathcal{K}_{l}^{(0)}$ and $\mathcal{K}_{l}^{(1)}$ from the pseudo label set $\{1,2,\dots,K\}$ and labelling their corresponding data points as class 0 and 1, respectively. Note that the number of classification tasks for a given $K$ is $L=\sum_{i=1}^{n-1}\sum_{k=1}^{\min(i,n-i)}\frac{\binom{n}{i}\times\binom{n-i}{k}}{1+\delta(i-k)}$, where $\binom{A}{B}$ denotes the binomial coefficient, and $\delta(.)$ represents the Dirac delta function.

3.3 Meta-training with the Unsupervised Classification Tasks

Meta-training estimates the parameters of a meta-learner, so it can be used as a pre-trained model that is efficiently fine-tuned to previously unseen classification tasks, using small annotated training sets [13]. The algorithm assumes that there exists a task distribution $\mathcal{T}$, from which each classification task $\mathcal{T}_{l}$ is drawn, where each task comprises a training set $\{\mathbf{v}_{i}^{(l,t)},\widetilde{y}_{i}^{(l,t)}\}_{i=1}^{M}$ and a testing set $\{\mathbf{v}_{i}^{(l,v)},\widetilde{y}_{i}^{(l,v)}\}_{i=1}^{N}$, with $M<<N$ and $M+N=|\mathcal{T}_{l}|$. Meta-training iteratively samples $T$ tasks from $\mathcal{T}$, and re-trains a multi-target classifier for those tasks using the training and testing sets defined above.

We use the MAML meta-training [17] that consists of a Bayesian hierarchical model, where $\psi$ denotes the classifier meta parameter, and $\phi_{l}$ represents the parameter for task $\mathcal{T}_{l}$. The meta-training objective function is defined by:

[TABLE]

where $T$ is the number of tasks per meta-training iteration, $\mathcal{Y}^{(v)}_{l}=\{\widetilde{y}_{i}^{(l,v)}\}_{i=1}^{N}$, $\mathcal{Y}^{(t)}_{l}=\{\widetilde{y}_{i}^{(l,t)}\}_{i=1}^{M}$, $\mathcal{V}^{(v)}_{l}=\{\mathbf{v}_{i}^{(l,v)}\}_{i=1}^{N}$, and $\mathcal{V}^{(t)}_{l}=\{\mathbf{v}_{i}^{(l,t)}\}_{i=1}^{M}$. In (4), we have

[TABLE]

where the lower bound is derived from Jensen’s inequality [18]. Therefore, the maximisation in (4) is approximated with the lower bound maximisation in (5), where the posterior $p(\phi_{l}|\mathcal{Y}_{l}^{(t)},\mathcal{V}^{(t)}_{l},\psi)$ is approximated with a Dirac delta function at a local optimal task-specific model parameter $\phi_{l}^{*}$, with $p(\phi_{l}|\mathcal{Y}_{l}^{(t)},\mathcal{V}^{(t)}_{l},\psi)=\delta(\phi_{l}-\phi_{l}^{*})$. The local optimal model parameter $\phi_{i}^{*}$ is obtained with truncated gradient descent initialised by the meta parameters $\psi$:

[TABLE]

where $\alpha$ is the learning rate, and the truncated gradient descent consists of a single step of (6). Maximising the lower bound of the log likelihood in (5) represents the MAML algorithm in [13], which produces a pre-trained model that can quickly learn new tasks drawn from $\mathcal{T}$.

4 Experiments and Results

4.1 Experimental Set-Up

We evaluate our proposed method on a breast DCE-MRI data set [2] (formally defined in Sec. 3.1), which has previously been used to evaluate few-shot training methods [4]. To allow a fair comparison with previous papers, we split the data set in a patient-wise manner into the same training, validation and testing sets, containing 45, 13, and 59 patients, respectively. We use the T1-weighted MRI to automatically extract the left and right breast regions from the first DCE-MRI subtraction volume [4]. Each breast region is resized into a volume of $100\times 100\times 50$ [4]. For the breast screening problem, only breasts that contain a malignant finding(s) are considered positive, while breasts with only benign findings or no findings are considered negative. There are 30, 9, and 38 positive and 60, 17, and 80 negative breasts in the training, validation and testing sets, respectively.

The model $f_{\theta}(\bf{v})$ that unsupervisedly produces the volume features is a 3D Densenet [19] composed of five dense blocks, each containing two dense layers. The features represent the input to the deep clustering algorithm, explained in Sec.3.2, with the number of clusters $K\in\{3,4,5\}$. The model that is meta-trained, and fine-tuned, has the same architecture as $f_{\theta}(.)$. During meta-training, we use a meta learning rate $\alpha=0.001$ in (6). At each meta-iteration, a meta-batch size of $T=4$ classification tasks is sampled according to a random or a curriculum learning strategy [4]. The meta-trained model is fine-tuned to the breast screening task using the entire training set, where model selection is performed using the validation set and results are reported in the test set.

The evaluation for the breast screening problem is based on the area under the ROC curve (AUC). We also measure the standard error utilising an estimate based on the Wilcoxon test [20] that estimates confidence intervals based on the testing set. In this evaluation, we study the type of task sampling for meta-training, i.e. random, or curriculum learning [4], and the influence of the number of clusters $K$ in (1), used to build the tasks. We compare our method (U-MT) with the previously proposed supervised meta-training for the case where the breast screening task is included (S-MT (S)) and not included (S-MT (NS)) in the meta-training process. We also compare our method with: a) Densenet trained from scratch on the breast screening task; b) Densenet from (a) fine-tuned with MIL [7]; c) Densenet trained with multi-tasking (using hand-designed tasks) [4]; d) Densenet pre-trained as a variational autocoder (i.e., unsupervised training) and fine-tuned for the breast screening task; and e) Densenet pre-trained with deep clustering (i.e., unsupervised training) and fine-tuned for the breast screening task. All Densenet models of these competing methods have the same architecture as the meta-trained model described above. The rationale for baselines (d) and (e) is to evaluate the effect of pre-training based on a reconstruction or a clustering scheme. With this purpose, we present results based on nearest neighbor classification and the fine-tuned classification model.

4.2 Results

We show the AUC results ($\pm$ standard error) for breast screening baselines in Tab. 1. Table 2 presents the results of meta-training, as a function of $K\in\{3,4,5\}$, with supervised and unsupervised task design using random and curriculum learning task sampling methods. Figure 2 presents examples of breast screening classification.

We measure the statistical significance of the difference in performance between our best performing approaches (Random with $K=5$ and Curriculum with $K=5$) and all baseline methods, obtaining a p-value $p\leq 0.001$ for all cases (unpaired two-tailed t-test). Also, comparing our newly proposed U-MT (Random with $K=5$) and S-MT (S) (Curriculum with $K=3$) [4], we obtain a p-value $p>0.05$.

5 Discussion and Conclusion

We have presented a new method that unsupervisedly designs classification tasks to meta-train medical image classifiers. Our method significantly outperforms several baselines consisting of traditional pre-training methods based on variational autoencoder, deep clustering, MIL, and multi-task learning (see Tab. 1). Our method also produces results comparable to the state-of-the-art set by meta-training using hand-designed tasks [4] (see Tab. 2). However, instead of using manually defined labels during meta-training, we unsupervisedly build classification tasks, allowing us to build a larger set of tasks, compared to the hand-designed ones. Also from Tab. 2, we notice that larger number of tasks, which increases with the number of clusters (Sec. 3.2), generally implies better AUC results. This confirms our initial hypothesis that, differently from computer vision problems, automatically building tasks is of great importance for medical image classification problems, where image labels that allow a large number of tasks are costly to obtain. We also observe that sampling tasks according to curriculum learning provides a good improvement of accuracy compared to random task sampling for a small number of clusters ($K=3$), but not for larger number of tasks ($K=5$). We hypothesize that meta-training with curriculum learning sampling needs a larger number of meta-iterations to learn a curriculum that is better than random task sampling. Given the large number of tasks for $K\in\{4,5\}$, the meta-training process converged before the curriculum learning algorithm – that deserves further research.