Improving Deep Learning Interpretability by Saliency Guided Training

1 Introduction . Deep Neural Networks ( DNNs ) have been widely used in a variety of different tasks [ 31 , 26 , 43 , 37 ] ; yet interpreting complex networks remains a challenge . Reliable explanations are necessary for critical domains like medicine , neuroscience , finance , and autonomous driving [ 9 , 34 ] . Explanations are also useful for model debugging [ 63 , 35 ] . As a result , various interpretability methods were developed to understand DNNs [ 5 , 49 , 22 , 53 , 52 , 32 , 51 ] . A common approach for understanding model decisions is to identify features in the input that highly influenced the final classification decision [ 6 , 63 , 53 , 52 , 47 , 36 , 64 ] . Such approaches , known as saliency maps , often use gradient calculations to assign an importance score to individual features , reflecting their influences on the model prediction . Saliency methods aim to highlight meaningful input features in model predictions to humans ; however , the maps produced are often noisy ( i.e. , contain visual noise ) . To improve the faithfulness of saliency maps , explanations methods that depend on more than one or higher-order gradient calculations were developed . For example , SmoothGrad [ 52 ] reduces saliency noise by adding noise to the input multiple times and then taking the average of the resulting saliency maps for each input . Integrated gradients [ 53 ] , DeepLIFT [ 47 ] and Layer-wise Relevance Propagation [ 5 ] backpropagate through a modified gradient function [ 3 ] while Singla et al . [ 51 ] studies the use of higher-order gradients in saliency maps . ∗Authors contributed equally †Code : https : //github.com/ayaabdelsalam91/saliency_guided_training 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this paper , we take a different approach to improve the interpretability of deep neural networks— instead of developing yet another saliency method , we propose a new training procedure that naturally leads to improved model explanations using current saliency methods . Our proposed training procedure , called saliency guided training , trains models that produce sparse , meaningful , and less noisy gradients without degrading model performance . This is done by iteratively masking input features with low gradient values ( i.e. , less important features ) and then minimizing a loss function that combines ( a ) the KL divergence [ 27 ] between model outputs from the original and masked inputs , and ( b ) the appropriate loss function for the model prediction . This procedure reduces noise in model gradients without sacrificing its predictive performance . To demonstrate the effectiveness of our proposed saliency guided training approach , we consider a variety of classification tasks for images , language , and multivariate time series across diverse neural architectures , including Convolutional Neural Networks ( CNNs ) , Recurrent Neural Network ( RNNs ) , and Transformers . In particular , we observe that using saliency guided training in image classification tasks leads to a reduction in visual saliency noise and sparser saliency maps , as shown in Figure 2 . Saliency guided training also improves the comprehensiveness of the produced explanations for sentiment analysis , and fact extraction tasks as shown in Table 1 . In multivariate time series classification tasks , we observe an increase in the precision and recall of saliency maps when applying the proposed saliency guided training . Interestingly , we also find that the saliency guided training reduces the vanishing saliency issue of RNNs [ 20 ] as shown in Figure 6 . Finally , we note that although we use the vanilla gradient for masking in the saliency guided training procedure , we observe significant improvements in the explanations produced after training by several other gradient-based saliency methods . 2 Background and Related Work . Interpretability is a rapidly growing area with several diverse lines of research . One strand of interpretability considers post-hoc explanation methods , aiming to explain why a trained model made a specific prediction for a given input . Post-hoc explanation methods can be divided into gradient-based methods [ 6 , 53 , 52 , 47 , 36 , 46 ] that can be reformulated as computing backpropagation for a modified gradient function and perturbation-based approaches [ 63 , 54 , 42 , 57 ] that perturb areas of the input and measure how much this changes the model output . Perkins et al . [ 40 ] uses gradients for feature selection through Grafting . Another line of works aims to measure the reliability of interpretability methods . This can be done by creating standardize benchmarks with interpretability metrics [ 18 , 21 , 11 , 56 , 45 , 41 ] or debugging explanations [ 1 , 24 , 15 , 2 ] by identifying test cases where explanations fail . Others [ 4 , 12 , 44 , 61 , 20 ] focus on modifying neural architectures for better interpretability . Similar to our line of work , Ghaeini et al . [ 14 ] and Ross et al . [ 44 ] incorporate explanations into the learning process . However , Ghaeini et al . [ 14 ] relies on the existence of the ground truth explanations while Ross et al . [ 44 ] relies on the availability of annotations about incorrect explanations for a particular input . Our proposed learning approach does not rely on such annotations ; since most datasets only have ground truth labels , it may not be practical to assume the availability of positive or negative explanations . Input level perturbation during training has been previously explored . [ 33 , 19 , 60 , 50 ] use attention maps to improve segmentation for weakly supervised localization . Wang et al . [ 59 ] incorporates attention maps into training to improve classification accuracy . DeVries and Taylor [ 10 ] masks out square regions of input during training as a regularization technique to improve the robustness and overall performance of convolutional neural networks . Our work focuses on a different task which is increasing model interpretability through training in a self-supervised manner . In this paper , we evaluate our learning procedure with the following saliency methods : Gradient ( GRAD ) [ 6 ] is the gradient of the output w.r.t the input . Integrated Gradients ( IG ) [ 53 ] calculates a path integral of the model gradient to the input from a non-informative reference point . DeepLIFT ( DL ) [ 47 ] compares the activation of each neuron to a reference activation ; the relevance is the difference between the two activations . SmoothGrad ( SG ) [ 52 ] samples similar input by adding noise to the input and then takes the average of the resulting sensitivity maps for each sample . Gradient SHAP ( GS ) [ 36 ] adds noise to the input , then selects a point along the path between a reference point and input , and computes the gradient of outputs w.r.t those points . We demonstrate the effectiveness of our training procedure using several neural network architectures : Convolution neural networks ( CNNs ) including VGG-16 [ 48 ] , ResNet [ 16 ] and Temporal Convolu- tional Network ( TCN ) [ 38 , 28 , 7 ] , a CNN that handles sequences ; Recurrent neural networks ( RNNs ) including LSTM [ 17 ] and LSTM with Input-Cell Attention [ 20 ] ; as well as Transformers [ 58 ] . 3 Notation . First , consider a classification problem on the input data { ( Xi , yi ) } ni=1 such that each X = [ x1 , . . . , xN ] ∈ RN has N features and y is the label . Let fθ denote a neural network parameterized by θ . The standard training of the network involves minimizing the cross-entropy loss L over the training set as follows : minimize θ 1 n n∑ i=1 L ( fθ ( Xi ) , yi ) ( 1 ) The gradient of the network output fθ ( X ) with respect to the input X is given by ∇Xfθ ( X ) . Let S ( . ) be a sorting function such that Se ( Z ) is the eth smallest element in Z . Hence , S ( ∇Xfθ ( X ) ) is the sorted gradient . We define the input mask function Mk ( . ) such that Mk ( S ( X ) , X ) replaces all xi where S ( xi ) ∈ { Se ( xi ) } ke=0 with a mask distribution , i.e. , Mk ( S ( X ) , X ) removes the k lowest features from X based on the order provided by S ( X ) . For a language input , we useX = [ x1 , . . . , xN ] where xi ∈ Rd is the feature embedding representing the ith word of the input . In that case , S ( X ) would sort elements ofX based on the sum of the gradient of the embeddings for each word x and Mk ( S ( X ) , X ) would mask the bottom k words according to that sorting . For a multivariate time series input , we use X = [ x1,1 , . . . , xF,1 , . . . , xF , T ] ∈ RF×T where T is the number of time steps and F is the number of features per time step . xi , t is the input feature i at time t ; sorting and masking would be done at the xi , t level . For two discrete probability distributions P and Q defined on the same probability space X , the Kullback–Leibler ( KL ) divergence [ 27 ] ( or , relative entropy ) from Q to P is given as DKL : DKL ( P ‖ Q ) = ∑ x∈X P ( x ) log ( P ( x ) Q ( x ) ) . ( 2 ) 4 Saliency Guided Training . Existing gradient-based methods can produce noisy saliency maps as shown in Figure 1 . The saliency map noise may be partially due to some uninformative local variations in partial derivatives . Using a standard training procedure based on ERM ( expectation risk minimization ) , the gradient of the model w.r.t . the input ( i.e. , ∇Xfθ ( X ) ) may fluctuate sharply via small input perturbations [ 52 ] . If gradient-based explanation methods faithfully interpret the model ’ s predictions , irrelevant features should have gradient values close to zero . Building on this intuition , we introduce saliency guided training , a procedure to train neural networks such that input gradients computed from trained models provide more faithful measures to downstream ( gradient-based ) saliency methods . Saliency guided training aims to reduce gradient values of irrelevant features without sacrificing the model performance . During saliency guided training , for every input X , we create a new input X̃ by masking the features with low gradient values as follows : X̃ =Mk ( S ( ∇Xfθ ( X ) ) , X ) ( 3 ) X̃ is then passed through the network which results in an output fθ ( X̃ ) . In addition to the classification loss , the saliency guided training minimizes the KL divergence between fθ ( X ) and fθ ( X̃ ) to ensure that the trained model produces similar output probability distributions over labels for both masked and unmasked inputs . The optimization problem for the saliency guided training is : minimize θ 1 n n∑ i=1 [ L ( fθ ( Xi ) , yi ) + λDKL ( fθ ( Xi ) ‖ fθ ( X̃i ) ) ] ( 4 ) where λ is a hyperparameter to balance between the cross-entropy classification loss and the KL divergence term . Since this loss function is differentiable with respect to θ , it can be optimized using existing gradient-based optimization methods . The KL divergence term encourages the model to produce similar outputs for the original input X and masked input X̃ . For this to happen , the model will need to learn to assign low gradient values to irrelevant features in model predictions . This potentially results in sparse and more faithful gradients as shown in Figure 1 . Masking functions : In images and time series data , features with low gradients are replaced with random values within the feature range . In language tasks , the masking function replaces the low salient word with the previous high salient word . This allows us to emphasize on high salient words and remove non-salient ones while maintaining the sentence length . The selection of k is dataset-dependent . It depends on the amount of irrelevant information in a training sample . For example , since most pixels in MNIST are uninformative , a larger k is desired . Detailed hyperparameters used is available in the appendix . Note that , only input features are masked during the saliency guided training . Limitations : ( a ) Compared to traditional training , our proposed training procedure is more computationally expensive . Specifically , the memory needed is doubled since now in addition to storing the batch , we are storing the masked batch as well . Similar to adversarial training , this training process is slow and takes a larger number of epochs to converge . For example , the standard training of a CIFAR-10 model usually takes on average 118 epochs to converge where each epoch is roughly 24 seconds . Using the saliency guided training , the convergence takes about 124 epochs where each epoch takes roughly 75 seconds ( all experiments on the same GPU ) . ( b ) Our training procedure requires two hyperparameters k and λ which might require a hyperparameter search ( we find that λ = 1 works well in all of our experiments ) . Algorithm 1 : Saliency Guided Training Given : Training samples X , # of features to be masked k , learning rate τ , hyperparameter λ 1 Initialize fθ 2 for i← 1 to epochs do 3 for minibatch do 4 Compute the masked input : Get sorted index I for the gradient of output with respect to the input . I = S ( ∇Xfθi ( X ) ) Mask bottom k features of the original input . X̃ =Mk ( I , X ) 5 Compute the loss function : Li = L ( fθi ( X ) , y ) + λDKL ( fθi ( X ) ‖ fθi ( X̃ ) ) 6 Use the gradient to update network parameters : fθi+1 = fθi − τ ∇θiLi 7 end 8 end

Summary: The authors propose an interpretable training procedure, in which the input features with the bottom k input gradients are masked. The loss term is comprised of the original task loss, plus a KL divergence between the model outputs on the original and masked inputs. The authors perform experiments on 3 image datasets, 3 language datasets, and a suite of synthetic time series datasets. They use model accuracy drop as a metric to evaluate saliency map quality.