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: This paper proposes a method (interpretable training) to improve the interpretability of saliency methods. The interpretable training method iteratively masks the features with small vanilla gradients and maximizes the similarity of the model outputs for both masked and unmasked inputs. They demonstrate this training procedure can improve interpretability in various data sets from CV, NLP, and time series across neural architectures, including RNNs, CNNs, and Transformers. They also find this method can reduce the vanishing saliency issue of RNNs. Lastly, they demonstrate improvements of the explanations produced after training by several gradient-based saliency methods.