Fast Axiomatic Attribution for Neural Networks

1 Introduction . Many traditional machine learning ( ML ) approaches , such as linear models or decision trees , are inherently explainable [ 4 ] . Therefore , an ML practitioner can comprehend why a method yields a particular prediction and correct the method if the explanation for the result is flawed . The prevailing ML architectures in use today [ 12 , 25 ] , namely deep neural networks ( DNNs ) , unfortunately , do not come with this inherent explainability . This can cause models to depend on dataset biases and spurious correlations in the training data . For real-world applications , e.g. , credit score or insurance risk assignment , this can be highly problematic and potentially lead to models discriminating against certain demographic groups [ 3 , 20 ] . To mitigate the dependence on spurious correlations in DNNs , attribution priors have been recently proposed [ 7 , 22 , 23 ] . By enforcing priors on the feature attribution of a DNN at training time , they allow actively controlling its behavior . As it turns out , attribution priors are a very flexible tool , allowing even complex model interventions such as making an object recognition model focus on shape [ 22 ] or less sensitive to high-frequency noise [ 7 ] . However , their use brings new challenges over regular training . First , computing the feature attribution of a DNN is a nontrivial task . It is critical to use an attribution method that faithfully reflects the true behavior of the deep network and ideally satisfies the axioms proposed by Sundararajan et al . [ 34 ] . Otherwise , spurious correlations may go undetected and the attribution prior would be ineffective . Second , since the feature attribution is used in each training step , it needs to be efficiently computable . Existing work incurs a trade-off between high-quality attributions for which formal axioms hold [ 34 ] and the time required to compute them [ 7 ] . Prior work on attribution priors thus had to choose whether to rely on high-quality feature attributions or allow for efficient training . In this work , we obviate this trade-off . Code and additional resources at https : //visinf.github.io/fast-axiomatic-attribution/ . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Specifically , we make the following contributions : ( i ) We propose to consider a special class of DNNs , termed efficiently axiomatically attributable DNNs , for which we can compute a closed-form axiomatic feature attribution that satisfies the axioms of Sundararajan et al . [ 34 ] , requiring only one gradient evaluation . As a result , we can compute axiomatic high-quality attributions with only one forward/backward pass , and hence , require only a fraction of the computing power that would be needed for regular DNNs , which involve a costly numerical approximation of an integral [ 7 , 34 ] . This significant improvement in efficiency makes our considered class of DNNs particularly well suited for scenarios where the attribution is used in the training process , such as for training with attribution priors . ( ii ) We formally prove that nonnegatively homogeneous DNNs ( termed X -DNNs ) are efficiently axiomatically attributable DNNs and establish a new theoretical connection between Input×Gradient [ 27 ] and Integrated Gradients [ 34 ] for nonnegatively homogeneous DNNs of different degrees . ( iii ) We show how X -DNNs can be instantiated from a wide range of regular DNNs by simply removing the bias term of each layer . While this may seem like a significant restriction , we show that the impact on the predictive accuracy in two different application domains is surprisingly minor . In a variety of experiments , we demonstrate the advantages of X -DNNs , showing that they ( iv ) admit accurate axiomatic feature attributions at a fraction of the computational cost and ( v ) beat state-of-the-art generic attribution methods for training regular networks with an attribution prior . 2 Related work . Attribution methods can roughly be divided into perturbation-based [ 13 , 39 , 41 , 42 ] and backpropagation-based [ 1 , 5 , 28 , 30 , 34 ] methods . The former repeatably perturb individual inputs or neurons to measure their impact on the final prediction . Since each perturbation requires a separate forward pass through the DNN , those methods can be computationally inefficient [ 28 ] and consequently inappropriate for inclusion into the training process . We thus consider backpropagationbased methods or , more precisely , gradient-based and rule-based attribution methods . They propagate an importance signal from the DNN output to its input using either the gradient or predefined rules , making them particularly efficient [ 28 ] , and thus , well suited for inclusion into the training process . Gradient-based methods have the advantage of scaling to high-dimensional inputs , can be efficiently implemented using GPUs , and directly applied to any differentiable model without changing it [ 2 ] . The saliency method [ 30 ] , defined as the absolute input gradient , is an early gradient-based attribution method for DNNs . Shrikumar et al . [ 27 ] propose the Input×Gradient method , i.e. , weighting the ( signed ) input gradient with the input features , to improve sharpness of the attributions for images . Bach et al . [ 5 ] introduce the rule-based Layerwise Relevance Propagation ( LRP ) , with predefined backpropagation rules for each neural network component . As it turns out , LRP without modifications to deal with numerical instability can be reduced to Input×Gradient for DNNs with ReLU [ 19 ] activation functions [ 1 , 27 ] , and hence , can be expressed in terms of gradients as well . DeepLIFT [ 28 ] is another rule-based approach similar to LRP , relying on a neutral baseline input to assign contribution scores relative to the difference of the normal activation and reference activation of each neuron . Generally , rule-based approaches have the disadvantage that each DNN component requires custom modules that may have no GPU support and require a model re-implementation . Axiomatic attributions . As it is hard to empirically evaluate the quality of feature attributions , Sundararajan et al . [ 34 ] propose several axioms that high-quality attribution methods should satisfy : Sensitivity ( a ) is satisfied if for every input and baseline that differ in one feature but have different predictions , the differing feature should be given a non-zero attribution . Sensitivity ( b ) is satisfied if the function implemented by the deep network does not depend ( mathematically ) on some variable , then the attribution to that variable is always zero . Implementation invariance is satisfied if the attributions for two functionally equivalent networks are always identical . Completeness is satisfied if the attributions add up to the difference between the output of the network for the input and for the baseline . Linearity is satisfied if the attribution of a linearly composed deep network aF1 + bF2 is equal to the weighted sum of the attributions for F1 and F2 with weights a and b , respectively . Symmetry preservation is satisfied if for all inputs and baselines that have identical values for symmetric variables , the symmetric variables receive identical attributions . [ 34 ] shows that none of the above methods satisfy all axioms , e.g. , the saliency method and Input×Gradient can suffer from the well-known problem of gradient saturation , which means that even important features can have zero attribution . To overcome this , [ 34 ] introduces Integrated Gradients , a gradient-based backpropagation method that provably satisfies these axioms ; it is considered a high-quality attribution method to date [ 7 , 14 ] . Its crucial disadvantage over previous methods is that an integral has to be solved , which generally requires an approximation based on ∼ 20–300 gradient calculations , making it correspondingly computationally more expensive than , e.g. , Input×Gradient . Attribution priors . The above attribution methods can not only be used for explaining a model ’ s behavior but also to actively control the behavior . To that end , the training objective can be formulated as θ∗ = arg min θ 1 |X| ∑ ( x , y ) ∈X L ( Fθ ; x , y ) + λΩ ( A ( Fθ , x ) ) , ( 1 ) where a model Fθ with parameters θ is trained on the annotated dataset X . L denotes the regular task loss , and Ω is a scalar-valued loss of the feature attribution A , which is called the attribution prior [ 7 ] ; λ controls the relative weighting . For example , by forcing certain values of the attribution to be zero , we can mitigate the dependence on unwanted features [ 23 ] . But also more complex model interventions , such as making an object recognition model focus on shape [ 22 ] or less sensitive to high-frequency noise [ 7 ] , can be formulated using attribution priors . An early instance of this idea is the Right for the Right Reasons ( RRR ) approach of Ross et al . [ 23 ] , which uses the input gradient of the log prediction to mitigate the dependence on unwanted features . While this is more stable than simply using the input gradient , it still suffers from the problem of saturation . RRR may thus not reflect the true behavior of the model , and therefore , miss relevant features . Subsequent work addresses this issue using axiomatic feature attribution methods , specifically Integrated Gradients [ 6 , 14 , 34 ] , which allows for more effective attribution priors [ 7 ] but incurs significant computational overhead , rendering them impractical for many scenarios . For example , Liu and Avci [ 14 ] report a thirty-fold increase in training time compared to standard training . Rieger et al . [ 22 ] propose an alternative attribution prior based on a rule-based contextual decomposition [ 18 , 31 ] ( CD ) as attribution method . This allows to consider clusters of features [ 7 ] instead of individual features and to define attribution priors working on feature groups . However , computing the attribution for individual features becomes computationally inefficient [ 7 ] . Additionally , since CD is a rule-based attribution method , it requires custom modules and can not be applied to all types of DNNs [ 7 ] . The very recently proposed Expected Gradients [ 7 ] method reformulates Integrated Gradients as an expectation , allowing a sampling-based approximation of the attribution . Erion et al . [ 7 ] argue that similar to batch gradient descent , where the true gradient of the loss function is approximated over many training steps , the sampling-based approximation allows to approximate the attribution over many training steps . This results in better attributions while using fewer approximation steps . Even using as little as one reference sample , i.e. , only one gradient computation , can yield advantages over the regular input gradient . However , we show that using only one reference sample still does not yield the same attribution quality as an axiomatic feature attribution method , reflected in less effective attribution priors . Schramowski et al . [ 26 ] propose a human-in-the-loop strategy to define appropriate attribution priors while training . Our attribution method is complementary and could be used within their framework . 3 Efficiently axiomatically attributable DNNs . Formally , given a function F : Rn 7→ R , representing a single output node of a DNN , and an input x ∈ Rn , the feature attribution for the prediction at input x relative to a baseline input x′ is a vector A ( F , x , x′ ) ∈ Rn , where each element ai is the contribution of feature xi to the prediction F ( x ) [ 34 ] . In this work we seek an attribution method that is particularly well suited for scenarios where such an attribution is used at training time , e.g. , training with attribution priors . As such , the attribution method should be of high quality , while being efficiently computable in a single forward/backward pass . Since attributions obtained from Integrated Gradients [ 34 ] have strong theoretical justifications and are known to be of high-quality [ 14 ] , they will serve as our starting point . In general , however , Integrated Gradients are expensive to compute for arbitrary DNNs . Therefore , in this work , we restrict our attention to a special sub-class of DNNs , termed efficiently axiomatically attributable DNNs , that require only a single forward/backward pass to compute a closed-form solution of Integrated Gradients . We show that nonnegatively homogeneous DNNs belong to this class and use this insight to guide the design of a concrete instantiation of efficiently axiomatically attributable DNNs . While there may be several such instantiations , we chose this particular one as it can be easily constructed from a wide range of regular DNNs by simply removing the bias term of each layer . This ensures comparability to prior work and allows for an easy adaptation of existing network architectures . Definition 3.1 . We call a DNN F : Rn 7→ R efficiently axiomatically attributable w.r.t . a baseline x′ ∈ Rn , if there exists a closed-form solution of the axiomatic feature attribution method Integrated Gradients IGi ( F , x , x′ ) along the ith dimension of x ∈ Rn , requiring only one forward/backward pass . Note that all differentiable models are efficiently axiomatically attributable w.r.t . the trivial baseline x′ = x . However , using such a baseline is not helpful . Instead , commonly chosen baselines are some kind of averaged input features or baselines such that F ( x′ ) = 0 , which allow an interpretation of the attribution that amounts to distributing the output to the individual input features [ 34 ] . Proposition 3.2 . For a DNN F : Rn 7→ R there exists a closed-form solution of IGi ( F , x,0 ) w.r.t . the zero baseline 0 ∈ Rn , requiring only one forward/backward pass , if F is strictly positive homogeneous of degree k ∈ R≥1 , i.e. , F ( αx ) = αkF ( x ) for α ∈ R > 0 . Proof . Sundararajan et al . [ 34 ] define the axiomatic feature attribution method Integrated Gradients ( IG ) along the ith dimension for a given model F , input x , baseline 0 , and straightline path γ ( α ) = αx as IGi ( F , x,0 ) = ∫ 1 0 ∂F ( γ ( α ) ) ∂γi ( α ) ∂γi ( α ) ∂α dα = ∫ 1 0 ∂F ( αx ) ∂αxi ∂αxi ∂α dα . ( 2 ) Assuming F is strictly positive homogeneous of degree k ≥ 1 , we can write Integrated Gradients in Eq . ( 2 ) as IGi ( F , x,0 ) = lim β→0 ∫ 1 β ∂F ( αx ) ∂αxi xi dα = lim β→0 ∫ 1 β αk−1 ∂F ( x ) ∂xi xi dα = 1 k xi ∂F ( x ) ∂xi . ( 3 ) The gradient expression in Eq . ( 3 ) can be computed using a single forward/backward pass . While Ancona et al . [ 1 ] already found that Input×Gradient equals Integrated Gradients with the zero baseline for linear models or models that behave linearly for a selected task , our Proposition 3.2 is more general : We only require strictly positive homogeneity of an arbitrary order k ≥ 1 . This allows us to consider a larger class of models including nonnegatively homogeneous DNNs , which generally are not linear . Definition 3.3 . We call a DNN F : Rn 7→ R nonnegatively homogeneous , if F ( αx ) = αF ( x ) for all α ∈ R≥0 . Corollary 3.4 . Any nonnegatively homogeneous DNN is efficiently axiomatically attributable w.r.t . the zero baseline 0 ∈ Rn and a closed-form solution of the axiomatic feature attribution method Integrated Gradients , requiring only one forward/backward pass , exists . Proof . Corollary 3.4 follows directly from Proposition 3.2 and Definitions 3.1 and 3.3 . Definition 3.5 . We let X -DNN denote a nonnegatively homogeneous DNN . Further , for any X -DNN F : Rn 7→ R , we let X -Gradient ( XG ) be an axiomatic feature attribution method relative to the zero baseline 0 ∈ Rn defined as XGi ( F , x ) = IGi ( F , x,0 ) = xi ∂F ( x ) ∂xi . ( 4 ) Note that while the formulas for the existing attribution method Input×Gradient [ 27 ] and our novel X -Gradient are equal , X -Gradient is only defined for the subclass of X -DNNs and provably satisfies axioms that are generally not satisfied by Input×Gradient . Additionally , from the nonnegative homogeneity of X -DNNs it follows that X -Gradient attributions are also nonnegatively homogeneous . This allows us to define another desirable axiom that is in line with intuition about how attribution should work and that is satisfied by X -Gradient . Definition 3.6 . An attribution method A satisfies nonnegative homogeneity if A ( F , αx , αx′ ) = α A ( F , x , x′ ) for all α ∈ R≥0 . For an overview of the axioms [ 34 ] that are satisfied by popular gradient-based attribution methods , see Table 1 . The right-hand side methods use only one gradient evaluation , and therefore , have similar computational expense . The left-hand side methods generally require multiple gradient evaluations until convergence , making them correspondingly computationally more expensive . Note that X -Gradient satisfies all the axioms satisfied by Integrated Gradients and Expected Gradients [ 7 ] , assuming convergence of the latter , while requiring only a fraction of the computational cost , however being only defined for X -DNNs . Existing methods that have similar computational expense as X -Gradient generally do not satisfy all of the axioms , and therefore , are likely to produce lower quality attributions , which can be misleading and less effective for imposing attribution priors . Constructing X -DNNs . With this motivation in mind , we will now study concrete instantiations of nonnegatively homogeneous DNNs . Note that this class of DNNs has already been considered by Zhang et al . [ 40 ] , however , neither at the same level of detail nor in the context of feature attributions . We define the output of a regular feedforward DNN F : Rn 7→ Ro , for an input x ∈ Rn , as a recursive sequence of layers l that are applied to the output of the respective previous layer : Fl ( x ) = { ψl ( φl ( WlFl−1 ( x ) + bl ) ) if l ≥ 1 x if l = 0 , ( 5 ) withWl and bl being the weight matrix and bias term for layer l , φl being the corresponding activation function , and ψl being the corresponding pooling function . Both φl and ψl are optional ; alternatively they are the identity function . For simplicity , we assume that the last task-specific layer , e.g. , the softmax function for classification tasks , is part of the loss function . Further , for a cleaner notation that aligns with [ 34 ] , we assume w.l.o.g . that we are only considering one output node at a time , e.g. , the logit of the target class for classification tasks . This yields the DNN F : Rn 7→ R that we consider and allows us to directly compute the derivative of the model w.r.t . an input feature xi . Importantly , the above formalization comprises many popular layer types and architectures . For example , fully connected and convolutional layers are essentially matrix multiplications [ 36 ] , and therefore , can be expressed by Eq . ( 5 ) . Skip connections can also be expressed as matrix multiplication by appending the identity matrix to the weight matrix so that the input is propagated to later layers [ 36 ] . This allows us to describe even complex architectures such as the ResNet [ 9 ] variant proposed by [ 40 ] . As the above definition of a DNN includes models that are generally not nonnegatively homogeneous , we have to make some assumptions . Assumption 3.7 . The activation functions φl and pooling functions ψl in the model are nonnegatively homogeneous . Formally , for all α ∈ R≥0 : αφl ( z ) = φl ( αz ) and αψl ( z ) = ψl ( αz ) . ( 6 ) Proposition 3.8 . Piecewise linear activation functions with two intervals separated by zero satisfy Assumption 3.7 . For z = ( z1 , . . . , zn ) ∈ Rn , these activation functions φl : Rn 7→ Rn are defined as φl ( z ) = ( φ ′ l ( z1 ) , . . . , φ ′ l ( zn ) ) with φ ′ l ( zi ) = { al,1zi if zi > 0 al,2zi if zi ≤ 0 . ( 7 ) Proposition 3.9 . Linear pooling functions or pooling functions selecting values based on their relative ordering satisfy Assumption 3.7 . For z = ( z1 , . . . , zn ) ∈ Rn , these pooling functions ψl : Rn 7→ Rm are defined as ψl ( z ) = ( ψ ′ l ( z ′ 1 ) , . . . , ψ ′ l ( z ′ m ) ) , ( 8 ) with z′i being a grouping of entries in z based on their spatial location and ψ ′ l : Rm 7→ R being linear or a selection of a value based on its relative ordering , e.g. , the maximum or minimum value . For proofs of Propositions 3.8 and 3.9 , please refer to the supplemental material . Activation functions in Proposition 3.8 include ReLU [ 19 ] , Leaky ReLU [ 16 ] , and PReLU [ 8 ] . Linear pooling functions in Proposition 3.9 include average pooling , global average pooling , and strided convolutions . Other pooling functions in Proposition 3.9 include max pooling and min pooling [ 32 ] , where the largest or smallest value is selected . Therefore , DNN architectures satisfying Assumption 3.7 include , inter alia , AlexNet [ 11 ] , VGGNet [ 29 ] , ResNet [ 9 ] as introduced in [ 40 ] , and MLPs with ReLU activations . They alone have been cited well over one hundred thousand times , showing that we are considering a substantial fraction of commonly used DNN architectures . However , these architectures are generally still not nonnegatively homogeneous . It is easy to see that even for a simple linear model F ( x ) = ax+ b that can be expressed by Eq . ( 5 ) and that satisfies Assumption 3.7 , nonnegative homogeneity does not hold , because 0F ( x ) = 0 6= b = F ( 0x ) . Therefore , in a final step we set the bias term of each layer to zero . As this may seem like a significant restriction , we show in Sec . 4 that the impact on the predictive accuracy in two different application domains is surprisingly minor . Corollary 3.10 . Any regular DNN given by Eq . ( 5 ) satisfying Assumption 3.7 can be transformed into an X -DNN by removing the bias term of each layer . Proof . A DNN F with L layers given by Eq . ( 5 ) with all biases bl set to 0 can be written as F ( x ) = ψL ( φL ( WL ( ... ( ψ1 ( φ1 ( W1x ) ) ) ) ) ) . As all the pooling functions ψl , activation functions φl , and matrix multiplicationsWl in F are nonnegatively homogeneous , it follows that F ( αx ) = αF ( x ) for all α ∈ R≥0 . Further discussion . We additionally note that our results have interesting consequences for DNNs in certain application domains , e.g. , in computer vision , as they allow to relate efficient axiomatic attributability to desirable properties of DNNs : Remark 3.11 . If a DNN F : Rn 7→ R , taking an image x ∈ Rn as input , is equivariant w.r.t . to the image contrast , it is efficiently axiomatically attributable . This observation follows directly from the fact that contrast equivariance implies nonnegative homogeneity . Consequentially , contrast-equivariant DNNs for regression tasks , such as image restoration or image super-resolution , are automatically efficiently axiomatically attributable . For classification tasks , such as image classification or semantic segmentation , contrast equivariance of the logits at the output implies efficient axiomatic attributability . If the classification is done using a softmax , then this also implies contrast invariance of the classifier output . In other words , there is a close relation between efficient axiomatic attributability and the desirable property of contrast equi-/invariance . We further illustrate this experimentally in Sec . 4.4 . Limitations . So far , we have discussed the advantages of X -DNNs such as being able to efficiently compute high-quality feature attributions . However , we also want to mention the limitations of our method . First , our method can only be applied to certain DNNs satisfying Assumption 3.7 . Although this is a large class of models , our method is not completely model agnostic as other gradient-based attribution methods . Second , removing the bias term might be disadvantageous in certain scenarios . Intuitively , the bias term can be seen like an intercept in a linear function , and therefore , is important to fit given data . As a matter of fact , a DNN without bias terms will always produce a zero output for the zero input , which might be problematic . Additionally , prior work argues that the bias term is an important factor for the predictive performance of a DNN [ 36 ] . However , these theoretical foundations are somewhat contradictory to our own practical findings . In Sec . 4 , we show that removing the bias term has less of a negative impact than perhaps expected , indicating that removing it can be a plausible intervention on the model architecture . As we can not guarantee that this holds for all DNNs , we recommend that practitioners who plan on using our method , first make a preliminary analysis of whether removing the bias from the model at hand is plausible . Third , our method uses implicitly the zero baseline 0 . As F ( 0 ) = 0 , this is a reasonable choice because it can be interpreted as being neutral [ 34 ] . Nevertheless , other baselines could produce attributions that are better suited for certain tasks [ 21 , 33 , 37 ] . For example , the zero baseline will generally assign lower attribution scores to features closer to zero , which can result in misleading attributions . Whether the advantages outweigh the disadvantages must be decided for each application , individually . In Sec . 4 we demonstrate the advantages of X -DNNs , beating state-of-the-art generic attribution methods for training with attribution priors .

Summary: The paper addresses the question of attributing the predictions of deep neural networks to the input features. It proposes to combine a particular restriction on the prediction function: non-negative homogeneity (e.g. a deep neural networks without biases) with the fast gradient x input explanation technique. With this combination, it is possible to produce explanations that satisfy the completeness axiom that can be computed quickly (from a single function evaluation). The paper also argues that non-negative homogeneity makes the model invariant to contrast, which is a useful property in image recognition.