UNDERWHELMING GENERALIZATION IMPROVEMENTS FROM CONTROLLING FEATURE ATTRIBUTION
Year: 2019

Abstract:
Overfitting is a common issue in machine learning, which can arise when the model learns to predict class membership using convenient but spuriouslycorrelated image features instead of the true image features that denote a class. These are typically visualized using saliency maps. In some object classification tasks such as for medical images, one may have some images with masks, indicating a region of interest, i.e., which part of the image contains the most relevant information for the classification. We describe a simple method for taking advantage of such auxiliary labels, by training networks to ignore the distracting features which may be extracted outside of the region of interest, on the training images for which such masks are available. This mask information is only used during training and has an impact on generalization accuracy in a dataset-dependent way. We observe an underwhelming relationship between controlling saliency maps and improving generalization performance.

heading: 1 INTRODUCTION
text: Overfitting is a common problem in machine learning, particularly when one uses powerful function approximators such as deep neural networks. When training these models with backpropagation, the network will evolve from modelling simple to more complicated functions until it finds salient discriminative features in the data. Once the model has found these, the gradients of the loss do not encourage the model to find other discriminative features in the data, even if they exist (Reed & Marks, 1999). In the classification case, this can be problematic if there exists some distractor feature xd in the data that is correlated with one of the output classes. This is a common issue in industry data (e.g., medical) where datasets are typically small and there are many confounding variables.
Consider the extreme case in a binary classification problem where in the training distribution there exists a confounding distractor element xd of the input data such that for Dtrain, p(y = 1|xd) = 1, while in the validation distribution Dvalid, p(y = 0|xd) = 1 (Figure 1). In this scenario, predicting using xd is easier than predicting using the true features that denote class membership and a classifier trained on Dtrain with traditional classification loss would predict the incorrect class with 100% probability on Dvalid. This is a textbook example of overfitting (Goodfellow et al., 2016; Reed & Marks, 1999). The existence of these overfit features is the motivation behind methods seeking to learn domain-invariant representations (Ganin & Lempitsky, 2014; Fernando et al., 2014), and is a common problem with real-world data (Badgeley et al., 2019; Zhao et al., 2019; Young et al., 2019).
In this paper, we explore the utility of various methods that allow one to use a mask on the input data to guide the network to avoid predicting from the defined region and penalize the network for attributing a prediction to a distractor. We present a synthetic dataset that encourages all models tested to overfit to an easy to represent distractor instead of a more complicated counting task. We present a novel “activation difference” (actdiff ) regularizer which mitigates this behaviour directly. We also present a method where we train an autoencoder/UNet to reconstruct a masked version of the input, indirectly controlling feature representations used for classification. We compare these methods with the recently-proposed gradmask (Simpson et al., 2019b), and present an expanded analysis of this algorithm’s behaviour. All code for this paper, and this dataset, are available here: https://github.com/bigtrellis2222/activmask.
We compare the real-life performance of these methods on open medical datasets with traditional classifiers, and demonstrate the differences in their feature attributions using saliency maps. Finally, we describe a medical dataset curated from two openly-available X-ray databases, and describe how samples can be drawn from each to generate a dataset biased by a site-diagnosis correlation inspired by previous work (Zhao et al., 2019). We demonstrate that, similarly to our synthetic datasets, classifiers are likely to predict using features unrelated to the task, and demonstrate that the proposed methods do mitigate this and often successfully refine the saliency maps to focus on the correct anatomy. However they do not consistently prevent overfitting.

heading: 2 RELATED WORK
text: It is a well-documented phenomenon that convolutional neural networks (CNNs), instead of building object-level representations of the input data, tend to find convenient surface-level statistics in the training data that are predictive of class (Jo & Bengio, 2017). Previous work has attempted to reduce the model’s proclivity to use distractor features by randomly masking out regions of the input (DeVries & Taylor, 2017). By randomly removing information from the inputs to the network, this method helped the network learn representations that aren’t dependent on single feature types in the image. However, this regularization approach gives no control over the kinds of representations learned by the model.
Recently, the Gradmask (Simpson et al., 2019b) and CARE methods (Zhuang et al., 2019) both proposed to control feature representations by penalizing the model for utilizing gradients outside of regions of interest. CARE was additionally designed to deal with class imbalances by increasing the impact of the gradients inside region of interest of the under-represented class.
In contrast to these two methods, we propose a new method which does not work with a saliency map, which can be noisy due to the ReLU activations allowing irrelevant features to pass through the activation function (Kim et al., 2019). Instead, this method operates directly on the activations themselves, encouraging the model to produce similar activation patterns in the presence of, and absence of, information outside of the region of interest.

heading: 3 METHODS
text: Actdiff Loss: To mitigate the effect of xd, we propose to explicitly regularize the network to ignore xd at test time by minimizing the distance between the network activations of the model when presented with a full input image and one where the information outside of some mask on the image has been corrupted. The model is only directly trained on unmasked examples. This encourages the network to build features which appear inside the masked regions even though it always sees the full image during training. The method requires having access to masks drawn by an expert who can distinguish between interesting and non-interesting discriminatve features, as is often the case in medical imaging. The actdiff regularization term is
Lact = 1
n n∑ l=1 ||ol(xmasked)− ol(x)||2, (1)
where ol(x) are the pre-activation outputs for layer l of the n-layer encoder f(x) when the network is presented with the original data x, and ol(xmasked) are the pre-activations outputs for layer l when presented with masked data xmasked. We call this the actdiff penalty. xmasked should be constructed by randomizing the indices of all pixels that fall outside of the mask, destroying any spatial information available in those regions of the image, but retaining the distribution of intensities
found in the data. Furthermore to retain important context around a masked region, we always dilate the mask by a set number of pixels. This loss has previously been used in style transfer to transform a random noise image into one containing the same layer-wise activations as some target image Gatys et al. (2015).
Reconstruction Loss: The pressure from Lact to not represent aspects of the input data x can be balanced with a reconstruction term that encourages the model to represent the data. Concretely, we can employ an auto-encoder architecture x̂ = g(f(x)), where f is a convolutional encoder and g is a deconvolutional decoder, and require the model to minimize the reconstruction loss Lrecon = 1m ∑m j=1 ||x(j) − x̂(j)||2, where j indexes the pixel-wise difference between the input x and reconstruction x̂. In our experiments a reconstruction term is helpful for guiding the network toward building useful representations of the data while employing Lact. All autoencoder models presented in this paper employed this loss term. Furthermore, the reconstruction task can be used to indirectly control feature representation learning by asking the model to reconstruct a masked version of the input given the full input x̂masked = g(f(x)) and minimizing Lrecon masked = 1m ∑m j=1 ||x (j) masked − x̂ (j) masked||2, which we used as a replacement for Lrecon.
Gradmask Loss: Gradmask is a recently proposed Simpson et al. (2019b) method for controlling which regions of the input are desirable for determining the class label using saliency maps. Saliency maps, or “input feature attribution”, can be calculated as ∂|ŷi|∂x for each input x (Zeiler & Fergus, 2013; Simonyan et al., 2014; Lo et al., 2015). In these experiments we minimize the ’contrast’ saliency between healthy and non-healthy classes (labels y0 and y1 respectively), as we expect that input variance which increases the distinction between the two classes leads to overfitting and is what we want to regularize. Therefore, we minimize
Lgrad = ∑ x∈D Lclf + ∣∣∣∣∂ |ŷ1 − ŷ0|∂x · (1− xseg) ∣∣∣∣ 2 , (2)
where ŷ0 and ŷ1 are the predicted outputs for our two classes and (1 − xseg) is a binary mask that covers everything outside the defined regions of interest.
Control Experiments: For many experiments, we evaluated the effect of simply training a model using masked data (and evaluating it using unmasked data). We call these experiments Classify Masked. We present this as a naive alternative to actdiff. Additionally, we evaluated what would happen if our auto encoder models were asked to reconstruct a masked version of the input (Lrecon masked above). These experiments are reported as Reconstruct Masked.

heading: 4 SYNTHETIC DATASET
text: Method: To evaluate the proposed methods for combating overfitting in the presence of a distractor variable, we generated a dataset following the description provided earlier (Figure 1) with 500 training, 128 validation, and 128 test examples respectively. The position of the distractors was perfectly correlated with class label in the training set and the logic governing this relationship was inverted for the validation and test sets. In cases where the model relies on the distractor to make a class prediction, we expected 0.0 AUC for the validation and test sets.
To evaluate the effect of the actdiff loss, gradmask loss, and the reconstruction penalty during training, we constructed a simple CNN architecture that optionally deconvolved the final layer to generate a reconstruction, or optionally did so in the style of a UNet (Ronneberger et al., 2015), see Figure 2. We additionally tested all non-reconstructing approaches using a simple 18-layer ResNet model (He et al., 2016). As control experiments, we also evaluated classifier performance when simply trained using masked versions of x, but evaluated on unmasked examples. All models were trained using Adam for 500 epochs with a learning rate of 10e−4, batch size of 32, with the weights of our regularizers in the loss term (denoted as λ) being λact = 10, λgrad = 10, and λrecon = 10 when applicable. The classification weight was set to 1. Before masking, masks were blurred using a Gaussian filter using a σ = 0.5, in order for some context to be included around the masked area. For reconstruction, the binary cross entropy loss was used as the input images were binary.
Results: The results of all experiments are shown in Table 1, with the architectures that successfully avoid overfitting in bold. All results are the average of 10 random model initilizations and data splits. To determine the effect of actdiff on feature representations in the network, we display saliency maps on the validation set during the final epoch of training in Figure 3.
First we demonstrate a CNN overfitting on Dtrain, using a simple CNN architecture trained using only the cross entropy loss Lclf = − ∑N i=1 y log ŷi. This model achieves 1.0 AUC on Dtrain and 0.0 AUC on Dvalid. Note that the model is attributing all saliency to the distractor. When the model is trained using masked inputs from Dtrain and trained using Lclf , the model performs well on Dtrain but is unsure how to handle the distractor in Dvalid, leading to performance similar to chance.
A CNN trained using both the classification and actdiff loss fails to learn any useful representations, scoring an AUC of 0.5 during train and validation (see ”CNN Actdiff” from Table 1). If this model instead a Convolutional Autoencoder (Conv AE) or UNet, and is additionally trained using a reconstruction loss, it can successfully learns to ignore the distractor when classifying the image x (”Conv AE Actdiff” and ”UNet Actdiff”). Note in Figure 3 that the UNet reconstruction, taken from the validation set, places the distractor in the location seen in the training set. This implies that the actdiff loss guided the model to learn to ignore the distractor before the classification layer, but retained some representation of it in an early layer of the network. Training using reconstruction without Actdiff does not prevent overfitting (”Conv AE Classify” and ”UNet Classify”). Experiments with an 18-layer ResNet demonstrate good performance using classification and actdiff alone (”ResNet Actdiff”).
Gradmask never produces a model with good generalizaion performance, with generally unstable training (Figure 7). We suspect this is because the saliency map is always non-zero everywhere on the input, leading to a constant source of noise in the loss function.
The best performing models took many more epochs to reach the optimal solution than would be expected for such a simple dataset. Note the ”Best Epoch” for overfit models is misleading as the best performance is chance. We found that models using a reconstruction loss more slowly approach the vicinity of their optimum than a ResNet (Figure 7), and the ResNet model takes longer to each its best epoch (see Table 1).

heading: 5 SINGLE-SITE MEDICAL DATASET WITH SEGMENTATIONS
text: Method: We applied all previous methods to three medical imaging datasets from the Medical Segmentation Decathlon (MSD) where one task is to detect the presence of an organ in the image, and for each organ we have a segmentation. Classifiers were previously shown to use the incorrect organ when performing this task (Simpson et al., 2019a). We tested our approaches on the liver detection task in CT, cardiac left atrium detection task in MRI, and pancreas detection task in CT. All results are the average of 10 independent seeds, which also lead to independent splits of the data. For each seed 128 training samples, 256 valid samples, and 256 test samples were randomly selected. A sample was a 2D image containing the segmentation. The mask blur factor was σ = 16
for the Liver and Pancreas datasets, and σ = 8 for the cardiac dataset. All images were resized to 100× 100 pixels. These blur values were selected as they provided the best gradmask performance, and are higher for the medical datasets when compared with the synthetic dataset experiments as the images used here are much larger (5122 vs 242). All models were trained with an Adam optimizer with a learning rate of 10e−4 for the Pancreas and Liver datasets, and 4e−3 for the Cardiac dataset, which were found to be the optimal learning rates using a hyperparameter search. All models were trained with a batch size of 32 and batch shuffling. The weights of our regularizers in the loss term (denoted as λ) were set to λact = 1, λgrad = 1, and λrecon = 1, when applicable. The classification loss weight was always set to 1. The reconstruction loss was the mean squared error.
Results: We present in all test AUCs for the best valid AUC over 500 epochs of training in Table 2 alongside the experiments from the previous section. For each model, the best-performing (or otherwise notable) configurations are in bold.
For CNN-based models, classification alone gave best performance. In contrast, the ResNet model generally benefited from the addition of gradmask, in contrast to the synthetic dataset results. The one notable exception was for the pancreas dataset, where the combination of classification, actdiff, and gradmask gave the best performance.
The best-performing model at baseline was the ResNet model, which is unsurprising given its superior expressive power over the simple CNN architecture we used for all other experiments. For the Liver and Cardiac datasets, classification with gradmask outperformed the baseline and all other models. In the case of the Pancreas dataset, classification with both actdiff and gradmask performed best. In all cases, classification with actdiff alone performed as well as, or worse than, the baseline.
Surprisingly, the auto-encoding models showed the best performance when trained to reconstruct a masked version of the input for the Liver and Pancreas dataset. For the Cardiac dataset, the best performing method was classification with gradmask for the Convolutional AutoEncoder and with actdiff for the UNet, and each achieve similar performance. This is likely because actdiff is too strong of a regularizer for CNN models, so the skip connections in the UNet allow the model to greatly reduce the actdiff penalty in the deeper layers of the encoder. The variance of model performance across seeds is higher if trained with gradmask than actdiff. Again, training with actdiff and gradmask outperforms either approach alone in the Liver and Pancreas datasets.
Individual saliency map examples for the Liver, Cardiac, and Pancreas datasets can be found in Figure 9, and the mean saliency over 100 test-set examples for all datasets can be found in Figure 4. In both the the Liver and Cardiac dataset, actdiff encourages the ResNet model to focus on the correct
anatomy, but this does not lead to an increase in test AUC performance over baseline. In contrast, gradmask less consistently encourages the model to focus its attention on the specified anatomy, but results in consistent test AUC performance improvements relative to baseline. The combination of the two methods (“Actgrad”) also produces improved feature attribution in the absence of improved generalization performance. In the Pancreas dataset, we see both actdiff and gradmask both focus the saliency maps of the model broadly across the anatomy. The gradmask and combination actdiff and gradmask models improved over baseline, but there is no clear reason why this would be true from the saliency maps. We therefore conclude an inconsistent relationship between improved generalization performance and refined saliency maps, where improved saliency does not nessicarily lead to improved performance.

heading: 6 MULTI-SITE X-RAY DATASET
text: Method: In an attempt to replicate the results of the synthetic dataset in a real world application, we constructed an X-Ray dataset using a combination of the PadChest (Bustos et al., 2019) dataset and the NIH Chestx-Ray8 (Wang et al., 2017) dataset. A site-driven overfitting signal has previously been reported when combining these datasets around the border of the images (Zech et al., 2018). We also observed this in regions far from the lungs (see the mean X-ray from each dataset in Figure 8), and therefore hypothesized we could improve overfitting performance by masking out the edges of the image using a circular mask. We constructed a joint dataset that allowed us to define a site-pathology correlation in the training set, and then produce validation and test set where the reverse relationship is true. In the training set, 90% of the unhealthy patients were drawn from the PadChest dataset and the remaining 10% of the unhealthy patients were drawn from the NIH dataset, and the reverse logic was followed for the validation and test sets. In all splits the classes and site distributions were always balanced, making it tempting for the classifier to use a site-specific feature when predicting the class in the presence of site-pathology correlation. We chose emphysema detection as the detection task, resulting in 998 samples for training, 498 samples for validation and 504 samples for test. All images were resized to 128×128 pixels. All experiments trained a 18-layer
ResNet model using an Adam optimizer with a learning rate of 10e−3 for 100 epochs. All results were averaged over 10 seeds. We trained a classifier on the same dataset with no site-pathology correlation as a baseline, and compare these results that with the same classifier in the face of a site-pathology correlation of 90%. We train two models using Lact and Lgrad with masks for all input images in the training set. We also trained a normal classifier with no Lact and Lgrad penalty, but where the masked region was set to zero in the train, valid, and test splits, to obtain an upper bound on the expected performance of our model.
Results: See Table 5b. A ResNet trained on the dataset mixed, without a site-pathology bias scores a test AUC of 0.7, while one trained in the presence of a strong site-pathology bias scores below chance on the test set (0.44). Both actdiff and gradmask improve performance of the model, but only actdiff scores above chance, and performs similarly to a model trained with the areas outside of the mask completely removed. However, the saliency maps of the ResNet trained with actdiff shows strong feature attribution from outside of the lungs. The model appears to be paying attention to the brightest regions of the image, which might be predictive of the scanner. Since a region of high intensity is available in the center of each X-Ray, actdiff when trained with these masks cannot handle this case of overfitting. Models trained with gradmask appropriately attribute more saliency to the lungs, but score below-chance level on the test set. Generally, the poor performance of all models in the presence of a site-pathology bias suggests that there is no regional source of the sitebias. This bias likely exists in almost every pixel of the dataset and therefore methods such as actdiff or gradmask are not well-suited to controlling overfitting in these scenarios.

heading: 7 CONCLUSION
text: We hypothesized that poor generalization performance could be partially attributable to classifiers exploiting spatially-distinct distractor features, and proposed the actdiff regularizer that prevents this behaviour on a synthetic dataset. We compare the performance of this method against previouslyproposed methods operating on saliency maps and demonstrate that the methods influence feature construction and generalization performance in a dataset-dependent manner. We conclude that while our methods successfully control the features constructed from the data, and solve the overfitting problem in a synthetic setting where the distracting feature is spatially distinct from the discriminative features, in real data we found no evidence of a spatially-distinct signal that can be reliably removed to mitigate overfitting. We now doubt the validity of using saliency maps for diagnosing whether a model is overfit because improving them does consistently improve generalization. Improved generalization performance observed in saliency-map based approaches may be more due to the fact that these approaches add useful noise into the updates, similarly to cutout (DeVries & Taylor, 2017). We leave this conjecture to future work.

heading: A MASK REQUIREMENT
text: Actdiff ’s requirement for hand-drawn masks can be a detriment in practice as they are costly to acquire from human experts. To determine whether actdiff has applicability in the setting where only a subset of the training set has masks, we repeated our experiments detailed above, retaining either 20%, 40%, 60%, 80%, or 100% of the masks in the training set. We analyzed the resulting final test AUC (Figure 6a) and the number of epochs (Figure 6b) required to reach this level of performance on the two best-performing actdiff models: the UNet and the ResNet18 model, averaging across 5 seeds. In general, the ResNet18 model appears to be more robust to missing masks, although across datasets, there does not seem to be a direct correlation between more masks and better performance. In fact, the addition of more masks can decrease performance, suggesting that the quality of the masks used is more important than the quantity (Figure 6a). There was no consistent effect of the number of masks used during training and the number of epochs required to reach the best epoch. We suspect having a small set of very precise masks is sufficient to guide the model toward developing good representations of the anatomy given that more compute time is available.

heading: B ARCHITECTURE OF THE CNN, AUTOENCODER, AND UNET MODEL
text: The encoder of the AutoEncoder and UNet model was shared with the CNN model and was 4 layers deep, with each layer consisting of a double convolution (kernel size of 3 and a stride of 1). All predictions were made off of the deepest layer of the network. The number of input channels was 16 for the synthetic dataset and 64 for the medical datasets, and doubled for each subsequent layer. All activations for actdiff were saved before applying the ReLU activation during the forward pass. During reconstruction, a sigmoid activation was optionally applied to the output to assist in the binary output case (for the synthetic dataset). In the decoder path of the autoencoding models, upsampling was applied using bilinear interpolation before each double convolution.
(a) Mean and standard deviation of the test AUC (as selected by the valid AUC) across 10 seeds for all experiments, under the condition that only some percentage of the masks are available. Only shown for the two best-performing Actdiff models (UNet and ResNet18).
(b) Mean and standard deviation of the best epoch (out of 500) based on the valid AUC across 10 seeds for all experiments, under the condition that only some percentage of the masks are available. Only shown for the two best-performing Actdiff models (UNet and ResNet18).
Figure 6: Results of the maximum masks experiments. (a) Best Test AUC for the best Valid AUC for each of the maximum masks conditions. (b) Best valid Epoch for each of the maximum masks conditions.

Review summaries:
This paper attempts to tackle the overfitting problem due to the focusing on the distracting information. By utilizing the dataset with masks, the authors propose a simple method to ignore the distracting features. By designing actdiff loss and reconstruction loss, the authors demonstrate that classifiers are likely to predict using features unrelated to the task and their losses can mitigate this problem.