Masked convolutional neural network for supervised learning problems

Abstract

Convolutional neural networks (CNNs) have exhibited superior performance in various types of classification and prediction tasks, but their interpretability remains to be low despite years of research effort. It is crucial to improve the ability of existing models to interpret deep neural networks from both theoretical and practical perspectives and to develop new neural network models with interpretable representations. The aim of this paper is to propose a set of novel masked CNN (MCNN) models with better ability to interpret networks and more accurate prediction. The key ideas behind MCNNs are to introduce a latent binary network to extract informative regions of interest that contain important signals for prediction and to integrate the latent binary network with CNNs to achieve better prediction in various supervised learning problems. Extensive numerical studies demonstrate the competitive performance of the proposed MCNN models.

1 |. INTRODUCTION

In the past decade, convolutional neural network (CNN) models have received much attention due to their competitive performance in various tasks, including object classification (He et al., 2016; Arvidsson et al., 2018; Xie et al., 2020) and semantic segmentation (Badrinarayanan, Kendall, & Cipolla, 2017; Chen, Papandreou, Kokkinos, Murphy, & Yuille, 2018; Yang, Shan, Kolen, & De With Peter, 2019; Falk et al., 2019). For instance, in object classification problems, we may be interested in identifying objects in the images that are associated with class labels. Many CNN-based models facilitate learning data-driven, highly representative, and layered hierarchical image features from complex datasets; an example is ImageNet (Deng et al., 2009). These models are quite robust to the large variation of object locations and sizes and tend to aggregate information over the whole image. However, it can be difficult to interpret most high-level image features extracted from CNNs; thus, interpretability remains to be a major challenge (Zhang & Zhu, 2018).

To address this challenge, we propose to modify standard CNNs to achieve better model interpretability and prediction in certain supervised learning problems without any additional human supervision. Our proposed approach differs from most existing methods in the direction of understanding neural network representations but can be regarded as a set of neural network models with interpretable/disentangled representations. A comprehensive review of various methods for improving the model’s ability to interpret CNNs can be found in Zhang and Zhu (2018).

The aim of this paper is to propose a set of masked CNN (MCNN) models with high model interpretability and better prediction. The key components of MCNNs are shown in Figure 1 and summarized as follows:

• Introduce a latent binary network to extract regions of interest (ROIs) for each image that contain informative signals for prediction.

• Simultaneously learn the latent binary network with CNNs for achieving better prediction in various supervised learning problems.

FIGURE 1.

A representative structure of MCNNs consisting of the latent binary mapping module with U-net and a classification network based on VGG (Simonyan & Zisserman, 2014). Here, ⊗ denotes the element-wise product between the masking matrix and the input image

Our MCNN can be regarded as a novel extension of the standard two-stage computer-aided diagnostic approach, which consists of segmenting objects of interest (e.g., a tumour) in the first stage and using segmented objects for prediction in the second stage. Specifically, the MCNN is a simultaneous segmentation-prediction approach that integrates a semantic segmentation network and CNNs into a single neural network model to improve prediction accuracy.

Compared with the existing methods in the literature, three major methodological contributions in this paper are summarized as follows:

• First, the MCNN introduces the latent binary network to carry out population semantic segmentation across all samples. It focuses on objects that are highly predictive of the response of interest, whereas standard semantic segmentation networks are developed to identify various ROIs in individual images that represent different objects. Moreover, although semantic segmentation networks are able to deliver pixel-wise annotations, they require extensive human labelling in preparing training samples, which are expensive to acquire (Yu & Koltun, 2015; Long, Shelhamer, & Darrell, 2015; Chen, Papandreou, Schroff, & Adam, 2017).

• Second, we propose to learn the latent binary network with CNNs to improve its interpretability, which can be widely applicable to CNNs with different architectures. Moreover, an MCNN focuses on the informative object learned from the latent binary network by ruling out irrelevant regions, so it may enhance predictive signals, subsequently leading to better prediction.

• Third, the MCNN is applicable to three major types of scenarios that are clustered according to whether accurate annotated mask images are available or not. This is in general enough to cover most real applications.

The rest of this paper is organized as follows. In Section 2, we discuss three different scenarios that the MCNN is mainly designed for. In Section 3, we introduce the technical details of the proposed MCNN. We demonstrate the performance of MCNNs in two synthetic experiments in Section 4 and two real applications in Section 5. Section 6 concludes the paper with some discussion.

2 |. DATA STRUCTURE

We consider three different scenarios that commonly appear in real applications.

• Scenario 1 assumes that we observe mask images that accurately capture the true predictive regions of interest.

• Scenario 2 assumes that we observe mask images that roughly capture the true predictive regions. This scenario is in general enough to cover the case without any masked information.

• Scenario 3 assumes that imaging data are mixed either with annotated and unannotated mask images or with high- and low-quality annotations. It can be regarded a mixture of Scenarios 1 and 2.

The proposed MCNN efficiently handles all three scenarios by assigning various weights for different samples. Details about various strategies for handling each scenario are discussed in Section 3.3.

3 |. MASKED CONVOLUTIONAL NEURAL NETWORK

An MCNN consists of two connected modules: a segmentation module that estimates the latent mask maps and a prediction module to perform regression or classification. We denote the network architecture of the MCNN as

$$\widehat{y}={\mathcal{F}}_1\left(X\otimes {\mathcal{F}}_0(X)\right),$$

where X represents the input image; ${\mathcal{F}}_0$ and ${\mathcal{F}}_1$ respectively correspond to the segmentation and prediction networks; $\widehat{y}$ denotes the predicted value, which can be a vector of K probabilities for a K-category classification problem or real values in the regression setting; and “⊗” represents the pixel-wise multiplication.

3.1. Segmentation module

The segmentation module enhances the predictive signals by masking off “background noises.” It estimates a latent binary map (or mask image) $\widehat{M}$ from the input image X such that we have $\widehat{M}={\mathcal{F}}_0(X)$. For instance, we set F₀(·) to be the U-net architecture that consists of a symmetric structure of auto-encoders and decoders (Ronneberger, Fischer, & Brox, 2015). In this case, the module utilizes convolutional and down-sampling layers to aggregate the information over the whole imaging space, gradually expands the feature maps by deconvolutional or up-sampling operations, and eventually constructs a probability map $\widehat{M}$ as the mask. Compared with standard semantic segmentation methods, the MCNN does not require pixel-wise annotated images, even though knowing such annotated images may substantially increase the discrimination power. Furthermore, the individual mask images $\widehat{M}$ explicitly localize important ROIs that contribute most to prediction outputs at the pixel level. As shown in numerical examples, the individual mask images $\widehat{M}$ can efficiently handle objects with large variation in terms of both location and size across samples. Thus, the use of $\widehat{M}$ dramatically improves the model’s ability to interpret deep neural networks.

3.2 |. Prediction module

The prediction module $\widehat{y}={\mathcal{F}}_1(\widehat{M}\otimes X)$ is directly connected with the output end of the segmentation module. Specifically, it takes the mask images as inputs and estimates the classification probabilities or numerical responses according to the type of learning problems. This module can adopt any CNN architecture. It sequentially processes the mask input images by using convolutional and down-sampling operations and utilizes fully connected layers to estimate the final predicted responses.

3.3 |. Loss functions

The loss function of an MCNN is the weighted sum of a segmentation loss and a prediction loss:

$$I(\widehat{y},\widehat{M})=I_p(\widehat{y})+\lambda l_s(\widehat{M}),$$

where λ balances the two modules. Moreover, in Scenario 1, we have mask images with precise pixel-wise annotations, so we may assign a relatively large λ to emphasize the segmentation module. In Scenario 2 with “imprecise” pixel-wise annotations, we use small values for λ. In the case of mixed samples (Scenario 3), we may apply adaptive weights according to the precision of annotations that is often available in practice.

The segmentation loss measures the similarity between the estimated probability map $\widehat{M}$ and the given mask image M. Popular options include the cross-entropy loss and dice coefficients. The cross-entropy loss measures the similarity between $\widehat{M}$ and M at the pixel level, that is, $I_s(\widehat{M})={\sum }_{1\le i\le d_1}{\sum }_{1\le j\le d_2}{M}_{ij}\text{log}{\widehat{M}}_{i,j}+\left(1-M_{i,j}\right)\text{log}\left(1-{\widehat{M}}_{i,j}\right)$, whereas the dice coefficient loss quantifies the overlap between the estimated map $\widehat{M}$ and the training mask M and is given by $I_s(\widehat{M})=-\frac{{\sum }_{i,j}{M}_{i,j}{\widehat{M}}_{i,j}}{{\sum }_{i,j}{M}_{i,j}+{\sum }_{i,j}{\widehat{M}}_{i,j}-{\sum }_{i,j}{M}_{i,j}{\widehat{M}}_{i,j}}$.

The prediction loss varies depending on the type of learning problems. For classification problems, we adopt a cross-entropy loss that measures the Kullback-Leibler divergence between the estimated probabilities and the observed values, that is, $l_p(\widehat{y})=-\sum _{k=1}^K\left\{y_k\text{log}{\widehat{y}}_k+\left(1-y_k\right)\text{log}\left(1-{\widehat{y}}_k\right)\right\}$, where $\widehat{y}={\left[{\widehat{y}}_1,\dots ,{\widehat{y}}_K\right]}^T$ represents the vector of predicted probabilities for the K class labels and y_k is the categorical indicator, that is, y_k = 1 if the object belongs to the kth category. For regression problems, we may apply the squared loss, that is, $l_p(\widehat{y})={\left(\widehat{y}-y\right)}^2$.

3.4 |. Implementation

We adopt a U-net (Ronneberger et al., 2015) structure as our segmentation module. We use a convolutional layer with kernel size (3, 3) or (3, 3, 3) for 2D or 3D input images, respectively, and apply zeros as padding to maintain fixed image sizes during convolutions. We add batch normalization (Ioffe & Szegedy, 2015) layers after each convolutional layer and activate the feature maps using rectifier liner units (Nair & Hinton, 2010). In the down-sampling phase of the network, we apply the maximum pooling layers with kernel size (2, 2) or (2, 2, 2) and use repetition up-sampling with the same kernel size to increase the resolution of feature maps in the up-sampling phase. The segmentation loss is the pixel-wise cross-entropy function.

For the prediction module, we use a VGG structure (Simonyan & Zisserman, 2014) with batch normalized convolutional layers for 2D images and a ResNet structure (He et al., 2016) for 3D images. For 2D images, we use λ = 10⁻⁴ as the weight parameter of the segmentation module and λ = 10⁻⁶ for 3D images when rough annotations are provided.

We conduct the numerical experiments with training, validation, and test datasets and use the stochastic gradient descent algorithm to train the networks (see Ruder, 2016, for details about the algorithm). In particular, we use the training set to estimate model parameters and evaluate the model on the validation set at the end of each epoch. We initialize the learning rate of the algorithm with 10⁻⁴ and gradually decrease it if the validation loss does not decrease for ten consecutive epochs. The model with the minimum validation loss is outputed as the final model and is used to evaluate the predictive accuracy on the test set.

In the numerical studies, we compare our model with the corresponding prediction network without the segmentation module, that is, the model with the same structure as the prediction module in the MCNN. We denote such a model as CNN to distinguish it from MCNN. For Scenarios 1–3, the MCNN models are denoted as MCNN 1–3, respectively, that is, MCNN 1 for the case with precise annotations, MCNN 2 for imprecise annotations, and MCNN 3 for mixed samples.

4 |. SYNTHETIC SIMULATION EXPERIMENTS

In this section, we conduct numerical studies with synthetic images, in which the segmentation ground truth is known. We generate the training mask according to the three scenarios discussed in Section 2 and apply the strategies in Section 3 to assign the loss weight.

4.1 |. Synthetic image regression

In this experiment, we simulate a set of symbolic images, each containing three ROIs: a circle, a square, and a triangular region. The ROIs vary randomly in size and in location within a 32 × 32 pixel grid. The responses are generated according to the radius of the circles and the area of the squares. The triangular regions are not associated with the responses. We add Gaussian noise with standard deviation 1 to each response and impose background noise to the covariate images according to a Gaussian random field. We generate 40,000 training, 10,000 validation, and 10,000 test samples. The precise training masks in Scenario 1 are the images with pixel-level annotations of the corresponding ROIs. The rough mask images in Scenario 2 cover the predictive signals with larger irregular regions. For Scenario 3, only 20% of the training samples have pixel-wise annotations, and the rest do not have any annotation information.

Some results are illustrated in Figure 2. The estimated masks clearly capture the predictive ROIs while ruling out the irrelevant triangular regions when precise annotations are provided. With rough training masks, the predictive ROIs are still identifiable by the model, and the background noise is reduced. In each of the three scenarios, the non-predictive triangular regions are effectively ruled out.

FIGURE 2.

Estimation results for the synthetic image regression. In each of the two panels, the image on the left is the original input image. The three images on the top are the estimated masks for Scenario 1–3 respectively. The images on the bottom are the corresponding mask images

4.2 |. Noisy MNIST

The MNIST (LeCun, Bottou, Bengio, & Haffner, 1998) dataset was constructed from a number of scanned documents collected by the National Institute of Standards and Technology. Each image is of size 28 × 28 pixels. There are 50,000 training samples, 10,000 validation samples, and 10,000 test samples in the dataset. The original MNIST images have clean backgrounds and are relatively easy to classify with high accuracy, for example, over 99.5% in Wan, Zeiler, Zhang, Le Cun, and Fergus (2013). In order to evaluate the improvement of the proposed network structure, we conducted our experiment with noisy MNIST images that we created by adding random noise to the original MNIST images and randomly resizing the digits and shifting them within a 32 × 32 pixel grid. The precise training masks cover the digit regions in each image. The rough masks contain broader regions around the digits. For mixed samples, we assign 20% of the training samples with pixel-wise annotations and the rest with no annotation information.

Figure 3 illustrates some of the estimation results, from which we can see that the estimated mask images can effectively block out background noise and accurately extract signals of digits from the noisy images. With the precise annotations, the extracted ROIs are more accurate, even when such annotations are available only for a small portion of samples.

FIGURE 3.

The estimation results for the noisy MNIST experiment. The top row is the plot of the input images with digits 0 to 9 from left to right. The following rows are the corresponding estimated masks from Scenarios 1 to 3, respectively

5 |. REAL APPLICATIONS

We apply MCNN models in real image prediction problems, in which the ground truth of the mask images is unknown. We apply our prior knowledge about the data and create binary images that roughly cover the predictive regions as training mask images. Those masks are unique across different samples and are refined accordingly in the estimated results. Therefore, these experiments can be categorized to Scenario 2 in Section 2.

5.1 |. Street view house number

In this experiment, we classify the street view house number images (Netzer et al., 2011). This dataset consists of colour images of house numbers collected by Google Street View. Each image is of size 32 × 32 pixels and may contain multiple digits. Our target is to classify the digit in the centre of the image. The other digits must be ignored. There are 73,257 images in the training set and 26,032 images in the test set. We further split the training images into 50,000 training samples and 23,257 validation samples for model selection. We applied the training mask images that cover the middle part and ruled out the side regions of the input images. With such mask images, the model tends to focus on the centre of the images so that any interference from the digits on the side can be reduced.

We illustrate the estimation results in Figure 4 and the training history in Figure 5. The estimated masks are able to highlight the target digits in the centre while masking the digits on the sides. The test accuracy is improved from 90.73% to 95.13% compared with that achieved by the CNN model. Some of the incorrect classification of images is potentially due to incorrectly annotated labels (e.g., the first row in the right group of images) or incomplete digit regions (see the second row in the right group). The training history also indicates a stable validation error and improved classification accuracy.

FIGURE 4.

The estimation results for the street view house number experiment. The left group of images shows three correctly classified images, and the right group corresponds to three misclassified images. Each of the six panels consists of three images: the original noisy image, the estimated mask, and the mask image

FIGURE 5.

Training history of the CNN and MCNN models in the street view house number experiment. The red and blue lines represent the loss function values of the CNN and MCNN models, respectively. The solid and dashed lines correspond to the training and validation losses, respectively. MCNN, masked convolutional neural network

5.2 |. Alzheimer’s Disease Neuroimaging Initiative magnetic resonance imaging classification

In this experiment, we aim to classify the structural magnetic resonance imaging (MRI) data from the Alzheimer’s Disease Neuroimaging Initiative (ADNI) study (see Figure 6). The ADNI study was launched in 2003 as a large-scale, long-term project to collect MRI and positron emission tomography images, cerebrospinal fluid, and blood biomarkers, among other data. The goal of the ADNI study is to track the progression of Alzheimer’s disease using these biomarkers and assess the brain’s structural and functional changes over different disease states. More information about this study can be found at the ADNI website (http://www.adni-info.org/).

FIGURE 6.

Timeline of the ADNI1, ADNI2, and ADNI GO studies. Image source: http://adni.loni.usc.edu/about/. ADNI, Alzheimer’s Disease Neuroimaging Initiative

We utilize the RAVENS-maps of the T1-weighted MRI images from different phases of the study, including ADNI1, ADNI2, and ADNI GO. The total number of participants in this study is 749, including participants with Alzheimer’s disease, mild cognitive impairment, and healthy status. For each participant, multiple images are collected at different time points, and their disease status may vary as well. We use the disease status at the time of image acquisition as its corresponding class label. After dropping the images with no diagnostic results or low quality, we collect a total of 3,021 images in our study. We generate the RAVENS-maps by following the pipeline in Liu, Liu, Zhu, and Alzheimer’s Disease Neuroimaging Initiative (2018) and further down-sample the maps to the resolution of 64 × 64 × 64 for the consideration of computational load. We randomly split the samples into the training (80%), validation (10%), and test (10%) sets in the modelling procedure and apply a mask that covers the whole brain region as the training image for each sample.

This experiment involves 3D images, and all the networks are built with 3D operations, including 3D convolution, up-sampling, and pooling. In consideration of the model size, we use the ResNet structure (He et al., 2016) for the prediction phase.

Figure 7 illustrates some of the estimation results. We can see that the estimated masks tend to select most of the brain regions, while focusing on the frontal cortex, temporal gyrus, hippocampus, and fornix regions. These parts of the brain have been well studied in the literature and shown to relate to planning, logical thinking, and memory (Bordi et al., 2016; Lozano et al., 2016; Alzheimer’s Association, 2019). With such masks, the classification accuracy is improved from 90.04% to 92.74% compared with that of the CNN model. From the plot of training losses in Figure 8, we can see that the loss of MCNN decreases more slowly at the beginning but decreases more quickly after 60 epochs, and the validation results become more stable as well. This is mainly due to the larger model size compared with that of the CNN model. Once a stable estimation from the segmentation module is achieved, the classification results become better and more stable.

FIGURE 7.

Estimation results for the Alzheimer’s Disease Neuroimaging Initiative experiment. The three groups of brain images respectively show typical samples of Alzheimer’s disease (AD), mild cognitive impairment (MCI), and healthy status (NC). From top to bottom, each column shows three images: the original image, the estimated mask, and the mask image

FIGURE 8.

Training history of the CNN and MCNN models in the Alzheimer’s Disease Neuroimaging Initiative experiment. The red and blue lines represent the loss function values of the CNN and MCNN models, respectively. The solid and dashed lines correspond to the training and validation losses, respectively. MCNN, masked convolutional neural network

We summarize the prediction results of all the experiments in Table 1. Note that with the segmentation module added, the MCNN models tend to deliver better results in both regression and classification. Moreover, the more precise annotations, the better prediction we can have.

TABLE 1.

Summary of results in the numerical experiments

Method	SIM	MNIST	SVHN	ADNI
CNN	2.27	5.30%	9.27%	9.60%
MCNN1	1.31	2.02%	N/A	N/A
MCNN2	1.60	2.33%	4.87%	7.26%
MCNN3	1.51	2.28%	N/A	N/A

Note. The mean squared prediction errors are reported for the regression problem, and the misclassification errors are reported for the classification problems.

6 |. DISCUSSION

In this paper, we propose an MCNN model that can carry out simultaneous segmentation-prediction procedures. More importantly, the proposed model can generate mask images that are able to select the predictive ROIs in the input images and mask off the background noise. This can potentially enhance the target signals and improve the predictive accuracy.

The segmentation module in the MCNN model functions as the pre-whitening process for the input images. This is beneficial when the backgrounds are actual noises and not informative for the prediction. In some cases, if the background signals and target objects were highly correlated, then the MCNN might not significantly improve the prediction.

We have focused on cases with non-informative backgrounds in the numerical experiments. The main purpose of these studies is to demonstrate the improved prediction achieved by the MCNN due to the segmentation module. Thus, instead of comparing our proposed model’s results with benchmark results, we have mainly compared the MCNN models with their corresponding CNN models with the same structure as our prediction module. We do not have a specific requirement for the network structure. Any segmentation and prediction networks can be used to construct an MCNN. However, including the segmentation module increases the size of the model. This can be a potential issue, but it may be solved by parameter sharing, that is, using the same structure and parameters in the down-sampling phase of both the segmentation and prediction modules. The competitive performance of the MCNN both in terms of interpretability and accuracy suggests that this is a promising area for future research.