How to use learning curves to evaluate the sample size for malaria prediction models developed using machine learning algorithms

Learning curves: construction

To generate a learning curve the following information and data needs to be available at the design stage of a study: a proposed prediction model or prediction model development process, and data from either a previous study or simulation containing outcomes and candidate predictors that resemble those that will be obtained from the target population of the study.

The following notation and concepts are required to set up a toy scenario to illustrate the basic ideas behind learning curves. Let $x$ be a continuous predictor (e.g. total IgG levels to an antigen of interest) and $y$ be a continuous outcome variable (e.g. parasitaemia) and say data $\{\left({x}_{i},{y}_{i}\right)\}$ has been collected from $i=1,\dots ,N\left(=100\right)$ individuals. In supervised learning, it is assumed that there exists a true model ($f$) that perfectly describes the relationship between the outcome and the candidate predictors [16]. In practice this model is typically unknown, and we estimate this unknown relationship with our prediction model ($\widehat{f}$) [16].

Some examples of prediction models ($\widehat{f}$ s) for continuous outcome variables that could be selected to estimate the true $f$ are linear regression models and random forests. Lastly, a measure of prediction error, or of how well predictions from $\widehat{f}$ match the observed outcome $y$, needs to be selected. For continuous outcomes, a commonly used measure of prediction error is the mean squared error ($MSE$):

$$MSE=\frac{1}{n}{\sum }_{i=1}^{N}{\left({y}_{i}-\widehat{f}\left({x}_{i}\right)\right)}^{2}$$

(1)

The $MSE$ will be small if the predictions $\widehat{f}$ are close to the observed outcomes $y$ and will be large if the predicted and observed outcomes differ greatly.

Figure 1 illustrates how learning curves would be constructed for this scenario. For the purposes of the toy example assume a linear regression model is an appropriate choice for the prediction model, $\widehat{f}$. To construct a learning curve, the data are split into a training dataset and test dataset. For example, in Fig. 1 data from 20 individuals are set aside as the test dataset; the prediction model is never trained on (i.e., fit to) the test dataset. Data from the remaining 80 individuals are used to construct training datasets, which the prediction model is trained on (i.e., fit to). In the training dataset column of Fig. 1, the size of the training datasets is increasing, while the size of the test dataset stays the same (i.e. 20 individuals in this example). In the first row, the training dataset size is one individual and the prediction model fits the outcome for that individual perfectly ($MSE$ = 0). However, predictions from that same model fit the test dataset of 20 different individuals poorly, resulting in a much higher $MSE$ on the test dataset. As the training dataset size increases, the prediction model cannot fit the training dataset perfectly anymore and the $MSE$ for the training dataset becomes larger. However, as the model is trained on more data, it starts to fit the test dataset better and the test $MSE$ decreases. Learning curves compare the $MSE$ (or another selected measure of prediction error) for the training datasets and the test dataset, as the training dataset size increases (as shown in Fig. 2 for the toy example).

Learning curves: interpretation

Learning curves can help us examine whether a prediction model’s poor performance (high error on the test dataset) is due to the model underfitting or overfitting the training data, which, in turn, can help us determine whether increasing the training dataset sample size will improve a prediction model’s accuracy on unseen data. Figure 3 illustrates how learning curves can be used to examine whether a prediction model is underfitting or overfitting the training data (see chapter “Measuring the Quality of Fit” of [16] and chapter “2.9 Model Selection and the Bias-Variance Tradeoff” of [18]).

A high error on the test dataset indicates that the trained model predicts the test data poorly (both left and right panels of Fig. 3). The magnitude/size of the errors on the training dataset can help identify if the high error on the test dataset arose because the model (1) fits the training data poorly because it is too simplistic to capture the relationship between the predictors and outcome (underfitting the training data) or (2) fits the training data too well as it captures random fluctuations/noise in the training dataset, rather than the relationship between the predictors and outcome (overfitting the training data). In both cases the trained model will have difficulty making accurate predictions on new/unseen data and exhibit high error on the test dataset.

In the scenario where the model does not fit the training data well (i.e., underfitting the training data), the learning curve for the training dataset and test dataset will show high errors as the training dataset size increases and there will be a narrow gap between the training and test curves (left panel of Fig. 3). In the scenario where the model fits the training data too well (i.e., overfitting the training data), the model will exhibit low errors on the training dataset and high errors on the test dataset and there will be a wide gap between the curves as the training dataset size increases (right panel of Fig. 3).

Identifying whether a model is underfitting or overfitting the training data can help us determine whether increasing the training dataset sample size will improve a model’s predictive performance. In the case where the model is underfitting the training data adding more training data is unlikely to help, as the two curves have either converged or appear to be converging to a high error as the training dataset size increases (left panel of Fig. 3).

In the case where the model is overfitting the training data adding more training data is likely to help, i.e., the test curve could converge toward the training curve if more training data were added (right panel of Fig. 3) due to the larger datasets providing more information for the model to learn from and reducing the chance of it modelling noise rather than underlying patterns in the data. This causes the training error to increase slightly and the test error to decrease and the curves to converge towards each other or the gap between the curves to reduce. The ideal scenario is where both the training and test curves have converged to a low error. Note the concept of irreducible error prevents these curves from converging to zero [16, 18].

In the case of underfitting the training data, where adding more training data is unlikely to help, some strategies to overcome this scenario are to change to a more complex model (e.g. random forests or support vector machines (SVM)), include interaction terms between predictors or additional transformations of the predictors to capture nonlinear relationships or consider whether more predictor variables can be collected (i.e., additional predictors included in the model).

Classification problems

Many of the concepts encountered in the toy example for continuous outcomes, such as underfitting, overfitting and the interpretation of learning curves, apply to the classification setting [16]. As the outcome $y$ is no longer continuous in classification problems, but binary or categorical, the only modification required in the classification setting is to select a measure of prediction error suitable for binary/categorical outcomes [16]. In the classification setting, measures of prediction error are derived from the confusion matrix – a cross tabulation of the observed binary/categorical outcome and the predicted binary/categorical outcome. The elements of a confusion matrix are defined in Table 1 for a binary outcome with categories positive and negative (e.g. malaria case yes/no).

Table 1 Elements of the confusion matrix for a binary outcome

A common measure of prediction error for classification problems is the balanced error rate (BER), which is the average of the false negative rate (FNR) and false positive rate (FPR) and indicates how likely it is that an individual in a particular category will be classified incorrectly. The FNR (1—sensitivity) is the proportion of observed positive outcomes that are incorrectly predicted by the model and is calculated from the confusion matrix as FN / (TP + FN). FPR (1—specificity) is the proportion of observed negative outcomes that are incorrectly predicted by the model and is calculated from the confusion matrix as FP/(TN + FP). The confusion matrix and balanced error rate can also be derived for outcomes with more than two categories/classes.

Mock study: motivating data for a malaria prediction model

The remainder of this tutorial aims to demonstrate how to generate and interpret learning curves for malaria prediction models developed using machine learning algorithms. To illustrate the approach, training dataset sizes will be evaluated to inform the design of a “mock” prediction modelling study aimed to predict the artemisinin resistance status of P. falciparum malaria isolates utilizing in vivo transcription data.

The motivating data for the mock prediction modelling study is from a large-scale genome-wide association study published by Mok et al. [19]. The objective of this study was to characterize the gene expression signatures associated with in vivo artemisinin-resistance. In brief, this study employed microarray technology to determine the global gene expression profiles of isolates sampled from 1043 patients with acute falciparum malaria enrolled into the Tracking Resistance to Artemisinin Collaboration (TRAC) study in 2011–2012 [20]. The samples were taken prior to artemisinin-based combination therapy (ACT) treatment. After treatment with ACT, these patients displayed differential rates of parasite clearance. The samples originated from 14 field sites across South East Asia (Pailin, Pursat, Preah Vihear, Rattanakiri in Cambodia; Mae Sot, Srisakhet, Khun Han, Ranong in Thailand; Shwe Kyin in Myanmar; Binh Phuoc in Vietnam; Attapeu in Laos), Bangladesh and Democratic Republic of Congo. The overall abundance of mRNA transcripts for 5061 of the ~ 5591 genes in the P. falciparum genome was measured.

Microarray data and a parasite clearance outcome were simulated based on the TRAC transcriptomics data and parasite clearance half-lives (defined in section “Outcome simulation”) published in Mok et al. [19] and the Supplementary Material of Mok et al. [19], respectively.

Microarray data simulation

Microarray data for 1043 participants/samples were simulated from a multivariate normal distribution with mean vector set to the sample means of the TRAC transcriptomics measurements and covariance matrix set to the sample covariance matrix of the TRAC transcriptomics measurements (Additional file 1).

Figure 4 compares the mean, minimum and maximum of the 5061 observed transcript abundances from the 1043 patients in the TRAC study to the mean, min and max of the 5061 transcript abundances simulated from a multivariate normal distribution for 1043 participants/samples (Additional file 1). The density plot of the means (panel “Mean”), minimums (panel “Min”) and maximums (panel “Max”) of the simulated transcript abundances overlap with the density plot of the means, minimums and maximums of the observed transcript abundances, indicating that the simulated transcript abundances resemble the observed transcript abundances.

Outcome simulation

The outcome in the mock prediction modelling study is whether a participant had a slow or fast clearing P. falciparum infection. Slow clearing is defined as a parasite clearance half-life ($P{C}_{1/2}$) > 5 h; fast clearing is a $P{C}_{1/2}$ $\le$ 5 h. In the TRAC study, 29% (299/1043) of participants had slow clearing infections, 70% (735/1043) had fast clearing infections, and for 1% (9/1043) the $P{C}_{1/2}$ could not be calculated. Typically, in sample size calculations we want to assume some of the predictors are associated with the outcome. To incorporate some association between the simulated microarray data (see section “Microarray data simulation”) and simulated outcome (slow/fast clearing infection), five genes were randomly selected to have transcript abundances associated with the outcome. A logistic regression model was used to simulate the binary outcome (slow/fast clearing infection) based on the five selected simulated predictors (further details provided in Additional file 1). The remaining 5056 transcript abundances are considered noise or not associated. 1043 outcome values were simulated and the percentage of slow clearing infections simulated was 27% (282/1043), which is similar to that observed in the TRAC study.

Conclusions regarding the sample size of omics studies of artemisinin resistant parasites should not be made based on the results of this mock study/dataset. While the simulated transcriptomics data resembles the observed transcriptomics data (Fig. 4), the simulated outcome is generated from five randomly selected (simulated) transcripts and the magnitude of the association between these transcripts and the outcome was selected to generate an outcome proportion similar to that in the observed dataset (i.e., does not reflect the actual relationship between these transcripts and slow/fast clearing infections). The purpose of the simulated data is to illustrate how to derive learning curves for data of this nature and to interpret the curves.

Statistical methods

Learning curves were constructed using two approaches. The first was a single train/test split approach where the simulated data (based on TRAC data) were randomly split into training and test datasets using an 80:20 ratio. This results in 835 individuals in the training dataset and 208 individuals in the test dataset. The training dataset of 835 individuals was then split into six smaller training datasets by randomly sampling without replacement 20, 40, 80, 160, 320, 640, and 835 individuals from the 835 individuals. The prediction models (described next in section “Prediction models”) were trained on each of the seven training sets. The trained model for each training dataset was used to make predictions on the test dataset of 208 individuals, which stays the same size across the seven training datasets. The seven training dataset BERs and corresponding seven test dataset BERs were both plotted against training dataset size to generate the learning curves for each prediction model.

The second approach is K-fold cross-validation. K-fold cross-validation generates K estimates of a model’s prediction error, that are averaged, to derive a more accurate estimate of a model’s prediction error on unseen data. This approach involves randomly dividing the dataset (all 1043 individuals) into K groups, or folds, of roughly equal size [16, 18]. K = 5 was selected for the mock study, resulting in 5 folds of size 208 or 209. For each training dataset size, the following was performed:

1.

One of the 5 folds was selected as the test dataset (either 208 or 209 individuals), and the remaining data of either 835 or 834 individuals was used as the training dataset. For the smaller training dataset sizes (n = 20, 40, 80, 160, 320 or 640), n individuals were sampled without replacement from the training dataset of 835 or 834 individuals.
2.

The BER is then computed on the training dataset and the test dataset.
3.

This procedure is repeated 5 times; each time, a different fold is treated as a test dataset.
4.

This process results in 5 estimates of the training BER and test BER, as opposed to the first approach, which produces a single estimate of the training BER and test BER for each training dataset size.

For each training dataset, the average of the training BERs and test BERs for a particular training dataset size, along with the min and max BER values, were plotted against training dataset size to generate the learning curves for each prediction model. For the last/largest training dataset size (835), the training BERs and test BERs are averaged across all 5 folds regardless of whether there were 835 or 834 in the training dataset or 208 or 209 in the test dataset.

Prediction models

Learning curves were constructed for two prediction models developed using machine learning algorithms. Note, throughout this section the training dataset refers to one of the seven training datasets of different sample sizes described in the previous section. The first machine learning method was Breiman’s random forest algorithm [21] implemented using the defaults of the randomForest function from R’s randomForest package [22].

The second was sparse Partial Least Squares-Discriminant Analysis plus Support Vector Machines (sPLSDA + SVM). In this combined approach, sPLSDA projects the predictors onto a domain of latent variables which are a linear combination of the predictors. By selecting a subset of the latent variables that best discriminate between the outcome (slow/fast parasite clearance), the number of predictors can be reduced, e.g., from 5061 predictors to 5 latent variables. SVM classifies observations by deriving non-linear boundaries (of dimension equal to the number of predictors) that separate the samples/individuals within each outcome category. The latent variables derived by sPLSDA were then used as predictors in SVM.

sPLSDA + SVM was implemented using splsda from R’s mixOmics package [23] and svm from R’s e1071 package [24, 25]. Further details on the prediction models are provided in Additional file 1.

All analyses were implemented in R version 4.4.1 [26]. The simulated data and R code to produce the learning curves presented in the Results section are available from: https://gitlab.unimelb.edu.au/sophiez/learning-curves-tutorial.

Source link