On Smoother Attributions using Neural Stochastic Differential Equations

1 Introduction

Deep neural networks (DNNs) are increasingly deployed in healthcare, security, autonomous driving, and other safety-critical applications. The responsible use of deep learning in such high-assurance applications necessitates the capability to explain their decisions. However, while a plethora of attribution and explanation techniques have been developed [Simonyan et al., 2013; Sundararajan et al., 2017; Shrikumar et al., 2017; Lundberg and Lee, 2017; Kim et al., 2018], these techniques are known to be noisy and not robust to small input perturbations [Ghorbani et al., 2019]. The noise in pixel-level attributions (Figure 1, center) produced by path-integral-based methods [Sturmfels et al., 2020; Smilkov et al., 2017; Xu et al., 2020; Dombrowski et al., 2019] has recently been analyzed, and factors such as high manifold curvature and choice of baselines have been identified.

In this paper, we focus on attributions computed over deep learning models that can be analyzed using the theory of dynamical systems [Chen et al., 2015; Lu et al., 2018]. Residual neural networks have been modeled using neural ordinary differential equations (ODEs) [Chen et al., 2018], and stochastic variants of ResNets are described using neural stochastic differential equations (SDEs) [Wang et al., 2019; Liu et al., 2020; Wang et al., 2019]. We leverage methods from the study of stochastic dynamical systems to show theoretically and empirically that a suitable injection of noise into the residual layers during training produces robust attributions. While addition of noise in the inputs during inference has empirically been shown to improve attributions [Smilkov et al., 2017], we use the theory of stochastic differential equations to inject noise at different residual layers and improve the robustness of attributions computed by path-integrals [Lundberg and Lee, 2017].

We make the following contributions towards computing robust attributions of deep neural network decisions:

1. We perform a worst-case analysis to show that the total attribution computed on neural SDEs trained using noise are more robust than attributions computed on neural ODEs. We are the first to use the connection between dynamical systems and residual neural networks to study the quality of attributions.
2. This robustness of the attributions of our SDE approach is agnostic to the choice of path-integral attribution methods (see tables 1, 2 and 3). We quantitatively demonstrate improvement in the computed attributions using the sensitivity metric [Yeh et al., 2019].
3. We experimentally illustrate that attributions computed using neural SDEs are qualitatively better, smoother and visually sharper than those obtained from neural ODEs. This improvement is demonstrated on several attribution methods including DeepLIFT SHAP, Smooth-Grad, and Integrated Gradients (see figures 1 through 4).

atic

ec-

has-

we

Es

ttri-

lish

lud-

the

im-

be

s an

i be

s on

ach

the

For

ting

the

that

ell-

is is

sed

un-

the

ith

line

d↵

(1)

ong

for

es.

(2)

and

and

re-

and of neural networks. In particular, we use this connection in Sec. 3.4 to show that neural SDEs lead to more robust attriage can be the baseline for images. The baseline can also be a set of random inputs where attribution is computed as an expected value.

butions that neural ODEs.

Dynamical systems like ordinary and stochastic differential

equations can be used to model residual networks (ResNets)

as such dynamical systems can represent the process of in-

Let the attribution for the j-th feature and output label i be

Aij (x). The attribution for the j-th input feature depends on

ference and training in ResNets [Chen et al., 2015; Chang

et al., 2017; Sonoda and Murata, 2017; Chen et al., 2018;

the complete input x and not just xj. The treatment for each

logit is similar, and so, we drop the logit/class and denote the

2 Overview Lu et al., 2018]. This connection between dynamical systems and neural networks gives us the opportunity to theoretically network output simply as F(·) and attribution as Aj(x). For simplicity, we use the baseline input xb = 0 for computing

founded:

scribed using the following equation:

1. The attribution is dominated by the linear term. This is

2016] with the residual mapping R(X(i),W(i)) can be de-

scribed using the following equation:

X(i + 1) = X(i) + R(X(i), W (i))

(3)

block and X(i+1) is the corresponding output that serves as

an input to the next building block. The weights of the neural

network layers in this ResNet building block are denoted by

W(i). In particular, X(0) denoted the input x in this notation

and the output F of a network with depth T is denoted by

dX(t)

We present a brief summary of our observations and sketch our theoretical analysis as well as our experimental results. study the robustness of neural network attributions using results from dynamical systems. That is the key focus of our investigations in this paper. A building block of a residual neural network [He et al., 2016] with the residual mapping R(X(i),W(i)) can be deattributions. We make the following two assumptions on the DNN model and the attributions, which reflect the fact that the model is well-trained and the attribution method is wellX(T) in this framework. After taking suitable limits, the evolution of the residual network can be described by the ResNet ordinary differential equation (ODE):

also an assumption made by attribution methods based

Here, X(i) is the input to the ith residual network building

block and X(i+1) is the corresponding output that serves as

an input to the next building block. The weights of the neural

network layers in this ResNet building block are denoted by

W(i). In particular, X(0) denoted the input x in this notation

xb to the input x, that is,

and the output F of a network with depth T is denoted by

X(T) in this framework.

After taking suitable limits, the evolution of the residual

network can be described by the ResNet ordinary differential

Aij (x) = (xj - xbj ) ⇥ Z 1

equation (ODE):

the j-th feature is denoted by @jF i (·).

any input x and the baseline input xb :

the input to the neural network.

Here, G(X(t),W(t)) = limδt!0 R(X(t),W(t))

The residual network is extended to a stochastic residual

network by adding a noise term N(i):

F(x) - F(xb ) =

The neural differential equation corresponding to such

stochastic residual networks is obtained by generalizing to a

stochastic differential equation. Here, the noise is modeled

Lee, 2017] satisfy this axiom too.

by weighing a Brownian motion term B(t) with a suitable

diffusion coefficient σ(X(t), t):

lationship between the robustness of attributions @Aij (x)

A natural question to ask is: Can we theoretically and

experimentally show that attributions corresponding to neu-

the output F(x) of the model.

ral stochastic differential equations (stochastic ResNets) are

inherently more robust that neural ordinary differential equa-

tions (ResNets)?

We will first mathematically compute the impact of a small

change in the input on the attribution of a traditional ResNet

or neural ordinary differential equation in Sec. 3.3. Next, we

will theoretically compute the impact of a small change in

the input on the attribution of a stochastic ResNet or neural

stochastic differential equation in Sec. 3.4, and mathematically demonstrate that neural stochastic differential equations generate more robust attributions. Figure 2: An overview of our observations and results: ResNets with stochastic noise injected into the residual layers or neural SDEs create more robust attributions. We show that the sum of change in attributions is smaller for neural SDEs than that for neural ODEs. The input image of a green snake on a green background and its attributions are shown; the integrated gradient attribution with noise tunnel for our neural SDE approach (bottom right) is visually sharper than the integrated gradient (IG) attribution as well as IG coupled to a noise tunnel for the neural ODE (bottom left).

Given a traditional residual network

the layers i of the network can be viewed as a time discretization of a continuous neural ODE

We suggest the use of neural SDEs

or ResNets with noise N(i) injected at each residual layer

This injection of noise in each residual layer of a residual network or a neural SDE leads to visually sharper and more robust attributions than the traditional residual networks or neural ODEs. Experimentally, we find that attribution methods create visually sharper attributions for neural SDEs. For example, Figure 2 shows how integrated gradients for neural SDE clearly identifies the green snake in the bottom right figure while both integrated gradients and integrated gradients with noise tunnel [Kokhlikyan et al., 2020] produce

dt = G(X(t),W(t))

(4)

and X(0) is

(5)

dX(t) = G(X(t),W(t)) dt +σ(X(t), t) dB(t) (6)

diffused attributions in the bottom left figures. The snake in the input image and in the attributions is highlighted using a dashed red box.

A natural question to ask is:

Can we theoretically and experimentally show that attributions corresponding to neural SDEs are more robust to small input perturbations than those obtained from neural ODEs?

Brief Summary of Results. Our main result is theoretical and empirical demonstration that neural SDEs create more robust attributions than neural ODEs. We consider a neural network such that a change in the input ϵx causes a change of ϵF in the output layer of the neural SDE or ODE. In most applications, linear layers are then used for projecting to a different dimension such as the number of classes of the overall deep learning classification model. In this paper, we focus on the residual layers to study the incremental benefit of using neural SDEs, and the output of the SDE or ODE model is of the same dimension as that of the input.

If the evolution of the dynamical system is contracting in the neighborhood of an input x, and ϵF is less than ϵx, then the dynamics is inherently robust to perturbations in the input. But in most applications such as classification, the dynamics is expected to increase ϵF with the increase in input perturbation helping create easy-to-learn classification boundaries between the inputs of different classes. This points to a natural trade-off between robustness and accuracy of deep learning models. While an increase in the ratio ϵF/ϵx helps improve the model accuracy, it also indicates a reduced robustness as very small changes ϵx in the input will produce large changes ϵF in the output representation. The difference in the sum of attribution over the features of an input x and that over a perturbed input ˆx can be computed as the difference in the output representations when the attribution technique satisfies the completeness axiom [Sundararajan et al., 2017]. We observe that robustness of the attribution for an input x appears to be at odds with model accuracy. This is not surprising since models with high accuracy are naturally expected to have multiple refined cues helping in making decisions, and a small change in the input can alter the cues responsible for the decision even if the decision itself remains unchanged. This presence of different set of fine-grained cues in high-accuracy models would generate different attributions over the input features.

3 Approach and Theoretical Results

We show that the attributions of neural stochastic differential equations are inherently more robust than the attributions of their deterministic counterparts. In Sec. 3.1, we briefly recapitulate the axiomatic definition of integrated gradients and then review the connection between residual networks with noise and neural stochastic differential equations (SDEs) in Section 3.2 before presenting our theoretical results.

3.1 Axiomatic Attributions

In this section, we discuss the axiomatic definition of attributions that will be used in subsequent sections to establish the robustness for neural SDEs. Attribution methods including integrated gradients and Shapley values often employ the notion of a baseline input xb ; for example, the all dark image can be a baseline for images. The baseline can also be a set of random inputs where attribution is computed as an expected value.

Let the attribution for the j-th feature and output label i be denoted by Aij (x). The attribution for the j-th input feature depends on the complete input x and not just xj . The treatment for each logit is similar, and so, we can drop the logit/class and denote the network output simply as F(·) and attribution as Aj(x). One can readily make the following two assumptions on the DNN model and the attributions, which reflect that the model is well-trained and the attribution method is well-founded:

1. The attribution is dominated by the linear term. This is also an assumption made by attribution methods based on Shapley values such as Integrated Gradient [Sundararajan et al., 2017] which define attribution as the path integral of the gradients of the DNN output with respect to that feature along the path from the baseline xb to the input x, that is,

where the gradient of i-th logit output of the model along the j-th feature is denoted by ∂jF i (·).

1. Attributions are complete, that is, the following property is true for any input x and the baseline input xb :

Shapley value methods such as Integrated Gradient and DeepShap [Sundararajan et al., 2017; Lundberg and Lee, 2017] satisfy this completeness axiom.

Our proof for robustness of attributions uses the completeness axiom satisfied by the attribution methods. In particular, we use this axiom to establish that the pathwise change in the sum of attributions ∑nk=1 (Ak(ˆx) -Ak(x)) for two nearby inputs x and ˆx can be related to the pathwise change in the output representation F(ˆx)-F(x) of the neural ODE/SDE model.

3.2 Neural Stochastic Differential Equations

We connect with the recent literature on dynamical systems for neural networks to analyze the robustness of attributions of neural networks. Dynamical systems such as ordinary and stochastic differential equations can be used to model residual networks (ResNets). The evolution of these dynamical systems represents the process of inference in ResNets [Chen et al., 2018] where the layers can be viewed as discretization of the continuous inference dynamics. This continuous dynamics view of inference and learning has been shown to improve efficiency as well as analyzability of the deep learning models [Chang et al., 2017; Chen et al., 2018; Lu et al., 2018]. In this paper, we exploit this connection between dynamical systems and neural networks to theoretically study the robustness of neural network attributions using results from dynamical systems. We show that neural SDEs where noise is injected at different layers of the neural network produce more robust attributions than neural ODEs in Sec. 3.3.

A building block of a residual neural network [He et al., 2016] with the residual mapping R(X(i),W(i)) can be described using the following equation:

Here, X(i) is the input to the ith residual network building block and X(i+1) is the corresponding output that serves as an input to the next building block. The weights of the neural network layers in this ResNet building block are denoted by W(i). In particular, X(0) denoted the input x in this notation and the final output representation F of the residual neural network with depth T is denoted by X(T ).

After taking suitable limits, the evolution of the residual network can be described by the continuous evolution of a ResNet ordinary differential equation (ODE):

Here, G(X(t),W(t)) = limδt→0 R(X(t),W(t)) δt and X(0) is the input to the neural network.

The residual network is extended to a stochastic residual network by adding a noise term N(i):

The neural dynamical system corresponding to such stochastic residual networks is obtained by generalizing to a stochastic differential equation [Wang et al., 2019; Liu et al., 2018; Liu et al., 2020; Wang et al., 2019]. Here, the noise is modeled by weighing a Brownian motion term B(t) with a suitable diffusion coefficient σ(X(t), t):

We show that the attributions corresponding to such neural stochastic differential equations are more robust to small perturbations in the input than those corresponding to the neural ODEs. We analyze the relative robustness of attributions for ResNets or neural ODEs and the stochastic ResNets or neural SDEs in Section 3.3.

3.3 Robustness of Attributions

In a neural ODE or SDE, the input x can be denoted by X(0) and the output representation F(x) by X(t). Given a perturbed input ˆx, the inferred output representation F(ˆx) of the network may be different from its output representation F(x) to the original input x.

Theorem 1. Consider a neural ODE with final output representation F where ϵx denotes a change in the input and ϵF denotes the corresponding change in the output F. Similarly, for a neural SDE, we denote the final output representation by ˜ F, the change in input by ˜ ϵx, and the change in the final output representation by ˜ ϵF. The attribution over the inputs for the neural ODE and the neural SDE model is denoted by A and ˜ A, respectively. The total change in attributions for a neural SDE is smaller than the total change in attributions for a neural ODE for the worst-case input x perturbed to ˆx, that is,

Proof. Given a neural ODE, one can compute a bound on the influence of a small change in the initial input x on the final output representation F. X(t) denotes the trajectory starting from the initial state x. Let ˆ CF ≥ 0 be a constant bounding the change in the dynamics, that is, for all t ∈ [0, T], the following holds:

From the completeness of attributions for the inputs x and ˆx, we know the following:

Subtracting attributions for x from those of ˆx, we get the following result:

We can similarly bound the change in attributions for the neural SDE. In order to compare a neural ODE and SDE for their relative robustness of attribution, we require them them to be trained with the same data and similar accuracy, except for the insertion of noise in SDEs. Due to this insertion of noise, the dynamics for the training data for neural SDEs will satisfy the following equation

where N depends on the injected noise. In order to compare a neural ODE and a neural SDE, they need to be trained on the same data and for all data points (including the extreme points where the above inequality is actually an equality), the model needs to have the same accuracy. Thus,

because eNT ≥ 1 since N is positive for any added noise in training the neural SDE.

4 Experimental Results

Our attribution analysis and model training is performed on a system with four NVIDIA V100 32GB GPUs using the ResNet-50, WideResNet-101-2 and ResNeXt-101 models [He et al., 2016] on the ImageNet benchmark. The neural SDE models are trained by injecting a normalized noise into each layer of the residual network. We use the Adam optimizer and the cross-entropy loss in our experimental studies.

We use the Captum tool [Kokhlikyan et al., 2020] for computing attributions and use a variety of different attribution methods including integrated gradients, Smoothgrad, DeepLIFT and DeepLIFT Shap to demonstrate the generality of our result. The improvement in attributions for neural SDEs is agnostic to the choice of attribution method. Further, we investigate qualitative improvement in the form of smooth robust attributions which are visibly sharper and well-defined, as well as improvement in more recently proposed quantitative metrics for attributions such as the sensitivity metric [Yeh et al., 2019].

4.1 Qualitative Results

We analyzed attributions for 100 images from the ImageNet validation data set using both Neural ODEs and Neural SDEs for ResNet-50, WideResNet-101-2 and ResNeXt-101 models, and repeated the analysis 5 times to eliminate any significant statistical variations in our quantitative results. Each image has been analyzed using five different attribution methods: integrated gradients [Sundararajan et al., 2017], integrated gradients with SmoothGrad [Smilkov et al., 2017], Saliency Maps [Simonyan et al., 2013], DeepLIFT [Shrikumar et al., 2017] and DeepLIFT SHAP [Lundberg and Lee, 2017]. The Adam optimizer with a learning rate of 0.0001 and 100 epochs is used to train stochastic residual network models with a normalized injected noise of magnitudes 0.1, 0.25, 0.5 and 0.75. Attributions are computed on inferences without any noise injection using the Captum library [Kokhlikyan et al., 2020].

We briefly highlight a few ResNet-50 examples in this section. Figure 5 shows the attributions obtained using integrated gradients and integrated gradients with a traditional noise tunnel [Smilkov et al., 2017] for the input in the middle two images. The rightmost image is the attribution obtained using integrated gradients with a noise tunnel on the Neural SDE model with a normalized noise of 0.5; our neural SDE approach produces a visually sharper image.

The middle image of Figure 6 visualizes the attributions obtained by using DeepLIFT SHAP [Lundberg and Lee, 2017; Kokhlikyan et al., 2020] on the neural ODE model.

The rightmost figure is an outcome of applying the same attribution on a neural SDE model with a normalized noise of magnitude 0.5. The two planes are more clearly visible in the neural SDE explanation.

Figure 7 illustrates the difference between the DeepLIFT attributions computed by neural SDEs and neural ODEs for the image of a humming bird from the ImageNet data set. The humming bird is clearly visible in the neural SDE attribution while the neural ODE attribution is diffused and spread over the tree and the bird in the image.

Figure 4 shows the impact of using Neural SDEs with different degrees of injected noise. The following two observations are relevant here:

1. A noise of small magnitude has an impact on the visual quality of attributions. However, we have not quantitatively studied the impact of the magnitude of the noise on the sensitivity metric [Yeh et al., 2019].
2. 97 2. The visual quality of the attributions improves as we increase the noise being injected. However, an increase in the magnitude of the noise also reduces the accuracy of the network. This trade-off in accuracy and robust attribution needs further investigation.

4.2 Quantitative Results

We quantitatively study the robustness of the computed attributions for both ResNet-50, WideResNet-101-2 and ResNeXt-101 models using the sensitivity metric [Yeh et al., 2019] implemented in the Captum library [Kokhlikyan et al., 2020]. We repeated our sensitivity analysis experiments 5 times on 100 correctly labeled but random images from the validation set of the ImageNet benchmark. We report the average values for ResNet-50 in Table 1 below.

97

97

Table 2 shows the average sensitivity metric obtained from five repetitions of our sensitivity analysis experiments on the WideResNet-101-2 model using both neural ODEs and This, our experiments demonstrate that the neural SDEs have lower sensitivity metric than neural ODEs for each of the attribution methods: saliency maps, integrated gradients, Deep LIFT, and DeepLIFT SHAP, and this result generalizes across ResNet-50, WideResNet-101-2 and ResNeXt-101 models. As long as we have a residual neural network architecture that can be trained as a neural SDE, the attributions are robust, smoother, and qualitatively sharper.

Table 2: The sensitivity metrics for neural SDEs are smaller than those for neural ODEs. Our experiments used the WideResNet-101-2 model and the ImageNet benchmark.

Table 3: Neural SDE produces a lower sensitivity metric than neural ODE for ImageNet using the ResNeXt101 32x8d model.

5 RelatedWork

Model interpretability and attribution methods. A number of techniques [Lundberg and Lee, 2017; Sundararajan et al., 2017; Li and Yu, 2015; Jha et al., 2017; Jha et al., 2019] for explaining deep neural networks have recently been proposed in the literature. These methods either provide a complete logical explanation for the output of the model or assign quantitative importance (attributions) to input features for a given model decision. Many of these methods are based on different analyses of the gradient of the predictor function with respect to the input [Simonyan et al., 2013; Selvaraju et al., 2017; Sundararajan et al., 2017]. A number of attribution methods have been compared in [Adebayo et al., 2018] and the sensitivity of these attributions to perturbations in the input has been studied in [Ghorbani et al., 2019]. These studies highlight the need for additional methods that can make attributions of neural networks more robust to input perturbations. Our observation that neural SDEs produce more robust attributions with a smaller sensitivity score [Yeh et al., 2019] is a step in this direction. The challenge of quantifying the improvement in robustness with the addition of noise during training of neural SDEs remains open.

Dynamical systems for neural networks. Dynamical systems models of neural networks have been the subject of several recent investigations with a particular emphasis on residual networks [Chang et al., 2017; Weinan, 2017; Lu et al., 2018; Tabuada and Gharesifard, 2020]. The theory of partial differential equations has been used to obtain dynamical system models of ResNets [Chang et al., 2017; Weinan, 2017; Lu et al., 2018]. Stochastic variants of residual neural networks have been described using neural stochastic differential equations [Wang et al., 2019; Liu et al., 2018; Liu et al., 2020; Wang et al., 2019]. Our main contribution in this paper is to leverage the connection between dynamical systems and neural networks to analyze the robustness and quality of attributions over the input features for a prediction by a deep learning model.

Connection between SmoothGrad and our work. SmoothGrad [Smilkov et al., 2017] uses attributions over multiple noisy variants of an input image to generate visually sharp attributions. Our qualitative experimental results show that neural SDEs lead to SmoothGrad attributions that are visually sharper than SmoothGrad attributions obtained using neural ODEs. In our experiments, we use the noise tunnel implementation of SmoothGrad, as implemented in Captum [Kokhlikyan et al., 2020]. Our approach can be seen as a generalization of the SmoothGrad [Smilkov et al., 2017] approach where noise is injected not just in the input but in all the internal representations of the deep neural network.

6 Conclusions

We make three key observations in this paper. First, we mathematically analyze the attributions computed on neural SDEs trained using noise and show that these are more robust than attributions computed on the deterministic neural ODEs. Second, we experimentally show that this improvement in the computed attributions is agnostic to the choice of path-integral attribution method. We demonstrate improvement over state-of-the-art attribution methods including Saliency Maps, DeepLIFT SHAP and Integrated Gradients (see Fig. 1 through Fig. 4). Finally, we experimentally illustrate that attributions computed using neural SDEs have lower sensitivity scores than those computed using neural ODEs for ResNet-50, WideResNet-101-2 and ResNeXt-101 models.

Acknowledgements

The authors acknowledge support from the National Science Foundation awards #2113307, #1908471, and #1740079, the DARPA cooperative agreement #HR00112020002, ONR grant #N000142112332, and the U.S. Army Research Laboratory Cooperative Research Agreement W911NF-17-2-0196. The views, opinions and/or findings expressed are those of the authors and should not be interpreted as representing the official views or policies of the Department of Defense or the U.S. Government.