A short post laying out our reasoning for using integrated gradients as attribution method. It is intended as a stand-alone post based on our LIB papers [1] [2]. This work was produced at Apollo Research.

Context

Understanding circuits in neural networks requires understanding how features interact with other features. There’s a lot of features and their interactions are generally non-linear. A good starting point for understanding the interactions might be to just figure out how strongly each pair of features in adjacent layers of the network interacts. But since the relationships are non-linear, how do we quantify their ‘strength’ in a principled manner that isn’t vulnerable to common and simple counterexamples? In other words, how do we quantify how much the value of a feature in layer $l + 1$ should be attributed to a feature in layer $l$ ?

This is a well-known sort of problem originally investigated in cooperative game theory. A while ago it made its way into machine learning, where people were pretty interested in attributing neural network outputs to their inputs for a while. Lately it’s made its way into interpretability in the context of attributing variables in one hidden layer of a neural network to another.

Generally, the way people go about this is setting up a series of ‘common-sense’ axioms that the attribution method should fulfil in order to be self-consistent and act like an attribution is supposed to act. Then they try to show that there is one unique method that satisfies these axioms. Except that (a) people disagree about what axioms are ‘common-sense’, and (b) the axioms people maybe agree most on don’t quite single out a single method as unique, just a class of methods called path attributions. So no attribution method has really been generally accepted as the canonical ‘winner’ in the ml context yet. Though some methods are certainly more popular than others.

Integrated Gradients

Integrated gradients is a computationally efficient attribution method (compared to activation patching / ablations) grounded in a series of axioms. It was originally proposed the context of economics (Friedman 2004), and recently used to attribute neural networks outputs to their inputs(Sundararajan et al. 2017). Even more recently, they started being used for internal feature attribution as well (Marks et al. 2024, Redwood Research (unpublished) 2022).

Properties of integrated gradients

Suppose we want to explain to what extent the value of an activation $f_{i}^{l_{2}}$ in a layer $l_{2}$ of a neural network can be ‘attributed to’ the various components of the activations $f^{l_{1}} = [f_{0}^{l_{1}}, \dots, f_{d}^{l_{1}}]$ in layer $l_{1}$ upstream of $l_{2}$ .^[1] For now, we do this for a single datapoint only. So we want to know how much $f_{i}^{l_{2}} (x)$ can be attributed to $f_{j}^{l_{1}} (x)$ . We’ll write this attribution as $A_{i, j}^{l_{2}, l_{1}} (x)$ .

There is a list of four standard requirements of properties attribution methods should satisfy that single out path attributions as the only kind of attribution methods that can be used to answer this question. Integrated gradients, and other path attribution methods, fulfil all of these (Sundararajan et al. 2017).

Implementation Invariance: If two different networks have activations $f^{l_{2}}$ , $g^{l_{2}}$ such that $f_{i}^{l_{2}} (f^{l_{1}}) = g_{i}^{l_{2}} (f^{l_{1}})$ for all possible inputs $f^{l_{1}}$ , then the attributions for any $f_{j}^{l_{1}}$ in both networks is the same.
Completeness: The sum over all attributions equals the value of $f_{i}^{l_{2}} (x)$ , that is $\sum_{j} A_{i, j}^{l_{2}, l_{1}} (x) = f_{i}^{l_{2}} (x)$ .
Sensitivity: If $f_{i}^{l_{2}}$ does not depend (mathematically) on $f_{j}^{l_{1}}$ , the attribution of $f_{j}^{l_{1}}$ for $f_{i}^{l_{2}}$ is zero.
Linearity: Let $g = a_{1} f_{i_{1}}^{l_{2}} + a_{2} f_{i_{2}}^{l_{2}}$ . Then the attribution from $f_{j}^{l_{1}}$ to $g$ should equal the weighted sum of its attributions for $f_{i_{1}}^{l_{2}}$ and $f_{i_{2}}^{l_{2}}$ .

If you add on a fifth requirement that the attribution method behaves sensibly under coordinate transformations, integrated gradients are the only attribution method that satisfies all five axioms:

Consistency under Coordinate Transformations: If we transform layer $l_{1}$ into an alternate basis of orthonormal coordinates, where the activation vector is one-hot ( $f^{l_{1}} (x) = [| | f^{l_{1}} (x) | |, 0, \dots, 0]$ )^[2] then the first component $f_{0}^{l_{1}} (x)$ should receive the full attribution $f_{i}^{l_{2}} (x)$ , and the other components should receive zero attribution.

In other words, all the attribution should go to the direction our activation vector $f^{l} (x)$ actually lies in. If we go into an alternate basis of coordinates such that one of our coordinate basis vectors $e_{1}$ lies along $f^{l_{1}} (x)$ , $e_{1} = \frac{f^{l_{1}} (x)}{| | f^{l_{1}} (x) | |}$ , then the component along $e_{1}$ should get all the attribution at data point $x$ , because the other components aren’t even active and thus obviously can’t influence anything.

We think that this is a pretty important property for an attribution method to have in the context of interpreting neural network internals. The hidden layers of neural networks don’t come with an obvious privileged basis. Their activations are vectors in a vector space, which we can view in any basis we please. So in a sense, any structure in the network internals that actually matters for the computation should be coordinate independent. If our attribution methods are not well-behaved under coordinate transformations, they can give all kinds of misleading results, for example by taking the network out of the subspace the activations are usually located in.

Property 4 already ensures that the attributions are well-behaved under linear coordinate transformations of the target layer $l_{2}$ . This 5th axiom ensures they’re also well-behaved under coordinate transforms in the starting layer $l_{1}$ .

We will show below that adding the 5th requirement singles out integrated gradients as the canonical attribution method that satisfies all five requirements.

Integrated gradient formula

The general integrated gradient formula to attribute the influence of feature $f_{j}^{l_{1}} (x)$ in a layer $l_{1}$ on feature $f_{i}^{l_{2}} (x)$ in layer $l_{2}$ is given by an integral along a straight-line path $C$ in layer $l_{1}$ activation space. To clarify notation, we introduce a function which maps activations from layer $l_{1}$ to $l_{2}$ $F^{l_{2}, l_{1}} : R^{d^{l_{1}}} \to R^{d^{l_{2}}}$ . For example, in an MLP (bias folded in) we might have $F^{l_{2}, l_{1}} (f^{l}) = ReLU (W^{l_{1}} f^{l_{1}})$ . Then we can write the attribution from $f_{j}^{l_{1}} (x)$ to $f_{i}^{l_{2}} (x)$ as

\begin{matrix} A_{i j}^{l_{2}, l_{1}} (x) & := \int_{C} d z_{j} {[\frac{\partial}{\partial z_{j}} (F_{i}^{l_{2}, l_{1}} (z))]}_{z = α f^{l_{1}} (x) + (1 - α) b^{l_{1}}} = f_{j}^{l_{1}} (x) \int_{0}^{1} d α {[\frac{\partial}{\partial z_{j}} (F_{i}^{l_{2}, l_{1}} (z))]}_{z = α f^{l_{1}} (x) + (1 - α) b^{l_{1}}} . \end{matrix}

where $z$ is a point in the layer $l_{1}$ activation space, and the path $C$ is parameterised by $α \in [0, 1]$ , such that along the curve we have $z (α) = α f^{l_{1}} (x) + (1 - α) b^{l_{1}}$ .^[2]

Intuitively, this formula asks us to integrate the gradient of $f_{i}^{l_{2}} (x)$ with respect to $f_{j}^{l_{1}} (x)$ along a straight path from a baseline activation $b^{l_{1}}$ to the actual activation vector $f^{l_{1}} (x)$ , and multiply the result with $f_{j}^{l_{1}} (x)$ .

We illustrate the integrated gradient attribution with a two-dimensional example. The plot shows a feature $f_{1}^{l_{2}}$ in layer $l_{2}$ that we want to attribute to the two features $f_{1}^{l_{1}}$ and $f_{2}^{l_{1}}$ in layer $l_{1}$ . The attribution to $f_{1}^{l + 1}$ (or $f_{2}^{l_{1}}$ ) is calculated by integrating the gradient of $f_{1}^{l_{2}}$ with respect to $f_{1}^{l_{1}}$ (or $f_{2}^{l_{1}}$ ) along a straight line from the baseline activation $b^{l_{1}}$ , here chosen to be $b^{l_{1}} = (0, 0)$ , to the activation vector $(f_{1}^{l_{1}} (x), f_{2}^{l_{1}} (x))$ and multiplying the result by the activation $f_{1}^{l_{1}} (x)$ (or $f_{2}^{l_{1}} (x)$ ).

Proof sketch: Integrated Gradients are uniquely consistent under coordinate transformations

Friedman 2004 showed that any attribution method satisfying the first four axioms must be a path attribution of the form

\begin{matrix} A_{i j}^{l_{2}, l_{1}} (x) := \int_{C} d z_{j} ⎡ ⎣ \frac{\partial}{\partial z_{j}^{l_{1}}} (F_{i}^{l_{2}, l_{1}} (z^{l_{1}})) ⎤ ⎦ with z^{l_{1}} (α) : R \to R^{n_{l_{1}}}, z^{l_{1}} (0) = b^{l_{1}}, z^{l_{1}} (1) = f^{l_{1}} (x), \end{matrix}

or a convex combination (weighted average with weights $c_{k}$ ) of these

\begin{matrix} A_{i j}^{l_{2}, l_{1}} (x) := \sum k c_{k} \int_{C_{k}} d z_{k, j} ⎡ ⎢ ⎣ \frac{\partial}{\partial z_{k, j}^{l_{1}}} (F_{i}^{l_{2}, l_{1}} (z_{k}^{l_{1}})) ⎤ ⎥ ⎦ with z_{k}^{l_{1}} (α) : R \to R^{n_{l_{1}}}, z_{k}^{l_{1}} (0) = b^{l_{1}}, z_{k}^{l_{1}} (1) = f^{l_{1}} (x), \sum k c_{k} = 1, c_{k} \geq 0 . \end{matrix}

Each term is a line integral along a monotonous path $C_{k}$ in the activation space of layer $l_{1}$ that starts at the baseline $b^{l_{1}}$ and ends at the activation vector $f^{l_{1}} (x)$ .

Claim: The only attribution that also satisfies the fifth axiom is the straight line from $b^{l_{1}}$ to $f^{l_{1}} (x)$ . That is, $c_{k} = 0$ for all the paths in the sum except for the path parametrised as

z_{1}^{l_{1}} (α) = b^{l_{1}} (1 - α) + α f^{l_{1}} (x) .

Proof sketch: Take $f^{l_{2}} (f^{l_{1}} (x)) = b^{l_{1}} + \sum_{k} U_{1, k} (f_{k}^{l_{1}} (x) - b_{k}^{l_{1}}) e^{- z \sum_{{i | i > 1}} {(\sum_{j} U_{i, j} (f_{j}^{l_{1}} (x) - b_{j}^{l_{1}}))}^{2}}$ as the mapping between layers $l_{1}$ and $l_{2}$ , with $U \in R$ an orthogonal matrix $U U^{T} = 1$ , and $U_{1, k} = \frac{f_{k}^{l_{1}} (x) - b_{k}^{l_{1}}}{| | f^{l_{1}} (x) - b^{l_{1}} | |}$ . Then, for any monotonous paths $C_{k}$ which are not the straight line $z_{1}^{l_{1}} (α)$ , at least one direction $v$ in layer $l_{1}$ with $v \cdot f^{l_{1}} (x) = 0$ will be assigned an attribution $> 0$ .

Since no monotonous paths lead to a negative attribution, the sum over all paths must then also yield an attribution $> 0$ for those $v$ , unless $c_{k} = 0$ for every path in the sum except $z_{1}^{l_{1}} (α) = b^{l_{1}} (1 - α) + α f^{l_{1}} (x)$ .

The problem of choosing a baseline

The integrated gradient formula still has one free hyperparameter in it: The baseline $b^{l}$ . We’re trying to attribute the activations in one layer to the activations in another layer. This requires specifying the coordinate origin relative to which the activations are defined.

Zero might look like a natural choice here, but if we are folding the biases into the activations, do we want the baseline for the bias to be zero as well? Or maybe we want the origin to be the expectation value of the activation $E (f^{l})$ over the training dataset? But then we’d have a bit of a consistency problem with axiom 2 across layers, because the expectation value of a layer $E (f^{l + 1})$ often will not equal its activation at the expectation value $E (f^{l})$ of the previous layer, $E (f^{l + 1}) \neq F^{l + 1, l} (E (f^{l}))$ . So, with this baseline the attributions to the activations in a layer $l$ would not add up to the activations in layer $l + 1$ . In fact, for some activation functions, like sigmoids for example, $0 \neq F^{l + 1, l} (0)$ , so baseline zero potentially has this consistency problem as well.

We don’t feel like we have a good framing for picking the baseline in a principled way yet.

Attributions over datasets

We now have a method for how to do attributions on single data points. But when we’re searching for circuits, we’re probably looking for variables that have strong attributions between each other on average, measured over many data points. But how do we average attributions for different data points into a single attribution over a data set in a principled way?

We don’t have a perfect answer to this question. We experimented with applying the integrated gradient definition to functionals, attributing measures of the size of the function $f_{i}^{l_{2}} : x \to f_{i}^{l_{2}} (x)$ to the functions $f_{j}^{l_{1}} :\to f_{i}^{l_{2}} (x)$ but found counter-examples to those (e.g. cancellation between negative and positive attribution). Thus we decided to simply take the RMS over attributions on single datapoints

$A_{i, j}^{l_{2}, l_{1}} (D) = \sqrt{\sum_{x} {A_{i, j}^{l_{2}, l_{1}} (x)}^{2}} .$

This averaged attribution does not itself fulfil axiom 2 (completeness), but it seems workable in practice. We have not found any counterexamples (situations where $A_{i, j}^{l_{2}, l_{1}} (D) = 0$ even though $f_{j}^{l_{1}}$ is obviously important for $f_{i}^{l_{2}}$ ) for good choices of bases (such as LIB).

Acknowledgements

This work was done as part of the LIB interpretability project [1] [2] at Apollo Research where it benefitted from empirical feedback: the method was implemented by Dan Braun, Nix Goldowsky-Dill, and Stefan Heimersheim. Earlier experiments were conducted by Avery Griffin, Marius Hobbhahn, and Jörn Stöhler.

^
The activation vectors here are defined relative to some baseline $b$ . This can be zero, but it could also be the mean value over some data set.
^
Integrated gradients still leaves us a free choice of baseline relative to which we measure activations. We chose 0 for most of this post for simplicity, but e.g. the dataset mean of the activations also works.

Interpretability: Integrated Gradients is a decent attribution method