Using Distributionally Robust Optimization to improve robustness in cancer pathology
Abstract
Computer vision (CV) approaches applied to digital pathology have informed biological discovery and clinical decision-making.
However, batch effects in images represent a major challenge to effective analysis.
A CV model trained using
Empirical Risk Minimization (ERM) risks learning batch-effects when they may align with the labels and serve as spurious correlates.
The standard methods to circumvent learning such confounders include (i) application of image augmentation techniques and (ii) examination of the learning process by evaluating through external validation (e.g., unseen data coming from a comparable dataset collected at another hospital).
The latter approach is data-hungry and the former, risks occluding biological signal.
Here, we suggest two solutions from the Distributionally
Robust Optimization (DRO) families.
Our contributions are i) a DRO algorithm using abstention which is a slight variation over existing abstention-based DRO algorithms and ii) a group-DRO method where groups are deﬁned as hospitals from which data are collected.
We ﬁnd that the model trained using abstention-based
DRO outperforms a model trained using ERM by 9.9% F1 in identifying tumor vs. normal tissue in lung adenocarcinoma (LUAD) at the expense of coverage.
Further, by examining the areas abstained by the model with a pathologist, we ﬁnd that the model trained using a DRO method is more robust to heterogeneity and artifacts in the tissue.
Together, we propose selecting models that are more robust to spurious features for translational discovery and clinical decision support.
Introduction
Computer vision (CV) approaches applied to cancer histopathology image data have demonstrated potential for biological discovery, precision diagnostics and as predictive biomarkers [1–5].
Previous work has shown that models trained on one hospital and tested on another show varying levels of performance [6].
This outcome could potentially result from the model learning spurious correlates in the data, such as batch effects, which are artifacts introduced as a result of the Whole Slide Image (WSI) preparation process, and induce a signal that is readily learnable, but not biologically relevant.
Mitigating batch effects parametrically incurs challenges, as they may arise from different parts of the tissue pre-processing pipeline [7].
Further, large models are likely to learn spurious correlates when trained to near-zero training error [8], resulting in poor test performance in sub-populations of the data, especially those that are under-represented in the training set [9].
This problem is especially exacerbated in the histopathology domain owing to the giga-pixel nature of the images, which warrant the image to be processed in smaller patches whist attributing the label given to the whole gigapixel image to all its patches - a phenomenon called weak labels [10].
A model risks learning spurious correlates on patches that don’t present the phenotype given to the giga-pixel image it was taken from.
Here, we propose circumventing this problem using DRO methods [9,11,12] in a histopathology task with clinical relevance.
Experimental Setup 2.1 Network Architecture 2.1.1 ERM
We use a pretrained ResNet-50 convolutional neural network (CNN) [13] pre-trained on the ImageNet dataset [14] to train an image classiﬁcation model, herefore referred to as ‘the ERM model’.
We used a cross-entropy loss function where the loss is computed and aggregated over all samples of a batch.
2.1.2 group-DRO
In our implementation of a group-DRO method, we deﬁned the groups as hospitals from which the
WSIs were taken.
We trained an algorithm by backpropagating the loss over the tiles from the worst performing hospital, measured by cumulative loss in a batch.
However, the reported statistics, such as F1, are reported over the whole validation / testing dataset, and not the worst performing hospital.
2.1.3 Model
trained using Abstention
Input Initialize ; for ; end
Algorithm Models were trained using an abstention algorithm
(Algorithm 1) whereby we only backpropagated the losses on images for which the predicted softmax scores were greater a threshold p for the predicted class.
We ﬁrst normalized the softmax scores using temperature scaling [15].
We ablated p and measured the coverage (deﬁned by [12] as the number of samples not abstained on) as well as the F1 performance of the model on the samples not abstained on.
2.2 Lung Adenocarcinoma (LUAD)
We evaluated a DRO method on the task of detecting tumor tissue in LUAD WSIs from the Cancer
Genome Atlas (TCGA) ( n = 522 ).
We trained a binary classiﬁer using slide-level labels to classify tiles into tumor or normal tissue.
LUAD is one of the two major histologic subtypes of Non-Small Cell Lung Cancers (NSCLC).
Identiﬁcation of the tumor in a WSI can help guide pathologic assessment and guide treatment decisions [3,16,17].
However, identiﬁcation of tumor may be confounded by scarring tissue from the effects of smoking on lung tissue, amongst other features.
In one set of experiments, we compared ERM against DRO methods trained on data from one hospital and validated on an external validation set consisting of unseen data coming from a comparable dataset collected at another hospital.
We also compared the use of image augmentation via recoloring to help reduce the batch effects.
In another set of experiments [Table] we combined data from multiple hospitals in our training set and ablated the number of hospitals contributing the external validation dataset to measure the robustness of the ERM and DRO methods.
Results
[Table]
First, we evaluated models trained on a single source site and validated on either the same or different single source site on a task of LUAD identiﬁcation.
Overall, we found signiﬁcant heterogeneity in model performance based on the hospital whose data were used to train and validate the model (Figure 1).
We also found that image augmentation produced only up to 0.15 improvement in F1, and using our abstention model in addition to image augmentation produced up to 0.24 improvement in
F1.
When trained on data from multiple hospitals, we found that the DRO model outperformed a conventional convolutional neural network trained using ERM for the task of detecting tissue with
LUAD under all numbers of hospitals held out, at the expense of coverage.
(Table 1)
We subsequently conducted statistical testing on performances of models that were thresholded post-hoc at the same thresholds that the models trained with abstention were trained with.
We found no statistically signiﬁcant advantage to using models trained with abstention.
However, as we will discuss below, similarly performing models can learn drastically different features, particularly spurious ones in the case of ERM models.
Upon investigation of the tiles that the DRO model with abstention abstained on, we noted that a DRO model abstained on tiles that an ERM model predicted incorrectly (Figure 2, top).
DRO models could have presented better performance by abstaining from learning potentially conﬂicting features, since these tiles might present contradictory features to their ground truth label.
Further, the ERM methods incorrectly predicted regions of a WSI as tumorous even at higher conﬁdences, that were conﬁrmed by a pathologist as non-tumorous (Figure 2, bottom)).
On the other hand the abstention method abstained on these tiles.
This could be because these tiles presented spuriously correlated features of that were exclusively correlated with the label of tumor.
Also, ERM models predicted tiles covered by slide-preparation artifacts that leaked through the QC pipeline such as air bubbles (Figure 3) as healthy surrounding tissue in three different slides, implying that these air bubbles might have introduced a spurious correlate. DRO models trained at high conﬁdence thresholds abstained from making predictions on these regions of the WSI.
By its abstention from artifacts, the abstention method is less likely to learn spurious correlates.
However, the above results might have been produced at the expense of coverage. Even at full coverage, abstention-trained models show utility in the content of the signal used by each model, in two examples in a task of distinguishing between the subtypes of Lung cancer (Figure 4). Models trained with abstention, avoid making predictions on pen marks, while ERM allocates high conﬁdence (darker color) to these artifacts (top row). In a second example taken from a brain biopsy of a metastatic lung cancer, we observed i) the ERM model placed importance on surrounding brain tissue which was conﬁrmed by a pathologist to not bear any tumor, and thus has learned spurious signal; and ii) the model trained with stringent abstention in contrast completely disregards the brain tissue, while placing modest conﬁdence in the veriﬁed lung tumor tissue.
We would like to highlight that no coverage was lost to thresholding in visualizing these examples and all tiles of Whole Slide Image (WSI) are shown.
However, owing to the different training processes, the models learned different features.
To further demonstrate the differences in features learned in each model type, we used each model to separately produce “pruned” datasets: datasets consisting of tiles with maximum softmax value above a certain threshold.
We subsequently used these datasets to train a further set of ERM models to distinguish lung subtypes.
At higher conﬁdence levels (0.8 and 0.9), models trained on DRO-pruned data offered better performance than those trained on ERM-pruned data (0.61 ± 0.12 vs 0.42 ± 0.
F1 [ p = 0 . 10] at threshold 0.8; 0.81 ± 0.12 vs 0.53 ± 0.12 F1 [ p = 0 . 047] at threshold 0.9).
Conclusion
Here, we evaluated the impact of batch effects and developed approaches to mitigate these fundamental challenges to digital pathology.
We suggested potential causes of heterogeneity in model performance that can impact downstream analyses and proposed models that are robust to the distributional shifts between training and held-out test sets.
Prospectively, consideration of batch effects in
CV histopathology analysis will guide successful biological investigations.  Discussion
Ultimately, we found that DRO methods that aim to either optimize the model’s performance on a previously deﬁned subgroup or a learned subgroup, deﬁned in our case by the training samples that the model performed well on, were able to provide better performances on an external validation set.
The latter approach is aligned with potential clinical support use cases, whereby a model can be allowed to abstain if it is not at least p % conﬁdent that the data are not sampled from the same distribution it has been trained on.
We make the assumption that examples that a model predicts with low conﬁdence are OOD.
However, this assumption needs further validation studies to conﬁrm.
A Appendix A.1 Measurement details of task training on one hospital and validating on another
In one set of experiments done on TCGA-LUAD, we trained the model on data taken from one hospital and validated it on data taken from another, to mimic a real-world setting where data is private and cannot be shared between institutions.
In a resource scarce setting, a model trained on one hospital, cannot be re-trained on data from another.
In order to study the effect of the preprocessing steps employed by a singular hospital, we were limited in our analysis to data from hospitals that have both tumor samples and surrounding normal tissue.
In this set of experiments, we only performed three cross validation trials, owing to resource limitations.
A.2 Training details
We train our models to minimize error and stop training if the error does not improve on the validation set over ﬁve consecutive measurements [18].
The validation performance was measured four times per epoch.
We used image augmentation via jittering the RGB pixel values in the RGB space to prevent overﬁtting to the color distribution by inducing random changes in the brightness, saturation, and other properties of an image, also known as color jitter [19].
We used a random-crop size ofpixels within the 512 pixel patch during our training process as a method to prevent overﬁtting.
We performed 5-fold cross validation on all of our experiments.
However, each fold of the cross-validation was not forced to be non-overlapping, owing to data availablility constraints.
Color Jitter We used image augmentation via jittering the RGB pixel values in the RGB space to prevent overﬁtting to the color distribution by inducing random changes in the brightness, saturation, and other properties of an image, also known as color jitter [19].
To discretize the color jitter, we deﬁned a light version of the color jitter that allowed the brightness to be chosen uniformly at random between [0.875, 1.125], the contrast to be chosen uniformly at random between [0.5, 1.5], the saturation to be chosen uniformly at random between [0.5, 1.5] and the hue to be chosen between [-0.1, 0.1].
We similarly deﬁned a heavy version of the color jitter to be four times proportionally higher.
We allowed the brightness to be chosen uniformly at random between [0.5, 1.5], the contrast to be chosen uniformly at random between [0, 3], the saturation to be chosen uniformly at random between [0, 3] and the hue to be chosen between [-0.4, 0.4].
The limit on the color jitter we could introduce was placed by the hue, which was forced to be between [-0.5, 0.5].
We chose the heavy color jitter to be less than the maximum allowed by the hue, to preserve some of the color signal presented by the tile.