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: This paper proposes to use non-negatively homogeneous DNN to perform feature attribution with better computational efficiency at a minimal performance cost. The so-called non-negatively homogeneous DNN is obtained by removing the bias term from most common neural networks and using a Relu/Leaky_Relu activation function. By doing this, the model would be linearly scaled by the input data (F(ax) = a*F(x)). With this property, the integrated Gradient (IG) method can be easily computed with one call of gradient computing, which improves its original computational efficiency. In the experiment, this model shows comparable performance with the IG method but much less computation effort.