Generalizing Backpropagation for Gradient-Based Interpretability (2023)

Abstract: Many popular feature-attribution methods for interpreting deep neural networks rely on computing the gradients of a model’s output with respect to its inputs. While these methods can indicate which input features may be important for the model’s prediction, they reveal little about the inner workings of the model itself. In this paper, we observe that the gradient computation of a model is a special case of a more general formulation using semirings. This observation allows us to generalize the backpropagation algorithm to efficiently compute other interpretable statistics about the gradient graph of a neural network, such as the highest-weighted path and entropy. We implement this generalized algorithm, evaluate it on synthetic datasets to better understand the statistics it computes, and apply it to study BERT’s behavior on the subject–verb number agreement task (SVA). With this method, we (a) validate that the amount of gradient flow through a component of a model reflects its importance to a prediction and (b) for SVA, identify which pathways of the self-attention mechanism are most important.

## introduction 1 one of the key contributors to the success of deep learning in nlp has been backpropagation @xcite , a dynamic programming algorithm that efficiently computes the gradients of a scalar function with respect to its inputs @xcite . backpropagation works by constructing a directed acyclic computation graph foot_1 that describes a function as a composition of various primitive operations, e.g., +, ×, and exp(•), whose gradients are known, and subsequently traversing this graph in topological order to incrementally compute the gradients. since the runtime of backpropagation is linear in the number of edges of the computation graph, it is possible to quickly perform vast numbers of gradient descent steps in even the most gargantuan of neural networks. while gradients are arguably most important for training, they can also be used to analyze and interpret neural network behavior. for example, feature attribution methods such as saliency maps @xcite and integrated gradients @xcite exploit gradients to identify which features of an input contribute most towards the model's prediction. however, most of these methods provide little insight into how the gradient propagates through the computation graph, and those that do are computationally inefficient, e.g., @xcite give an algorithm for computing the highest-weighted gradient path that runs in exponential time. in this paper, we explore whether examining various quantities computed from the gradient graph of a network, i.e., the weighted graph whose edge weights correspond to the local gradient between two nodes, can lead to more insightful and granular analyses of network behavior than the gradient itself. to do so, we note that backpropagation is an instance of a shortest-path problem @xcite over the (+, ×) semiring. this insight allows us to generalize backpropagation to other semirings, allowing us to compute statistics about the gradient graph beyond just the gradient, all while retaining backpropagation's linear time complexity. 3 in our experiments, the first semiring we consider is the max-product semiring, which allows us to identify paths in the computation graph which carry most of the gradient, akin to @xcite influence paths. the second is the entropy semiring @xcite , 4 which summarizes how dispersed the gradient graph is, i.e., whether the gradient flows in a relatively focalized manner through a small proportion of possible paths or in a widely distributed manner across most paths in the network. with experiments on synthetic data, we validate that the max-product semiring results in higher values for model components we expect to be more critical to the model's predictions, based on the design of the transformer @xcite architecture. we further apply our framework to analyze the behavior of bert @xcite in a subject-verb agreement task (sva; @xcite . in these experiments, we find that the keys matrix for subject tokens carries most of the gradient through the last layer of the self-attention mechanism. our results suggest that semiringlifted gradient graphs can be a versatile tool in the interpretability researcher's toolbox. ## gradient-based interpretability neural networks are often viewed as black boxes because their inner workings are too complicated for a user to understand why the model produced a particular prediction for a given input. this shortcoming has spawned an active field of research in developing methods to better understand and explain how neural networks work. for example, feature attribution methods aim to measure the sensitivity of a model's predictions to the values of individual input features. many of these methods quantify feature attribution as the gradient of the model's output with respect to an input feature @xcite @xcite . we note that while the general reliability and faithfulness of gradient-based methods has been a contentious area of research @xcite @xcite , gradient-based methods have nonetheless continued to be widely used @xcite @xcite . other works have applied feature attribution methods to not only highlight sensitive input features but also uncover important internal neurons. @xcite define influence as the gradient of a quantity of interest with respect to a neuron, averaged across a collection of inputs of interest. @xcite further define and analyze the notion of influence paths, i.e., paths in the computation graph between the neuron of interest and the output that on average carry most of the gradient. by applying this method to analyze the behavior of @xcite lstm language model on the sva task, they draw conclusions about which internal components of the lstm are most sensitive to the concept of number agreement based on the paths with the greatest amount of influence. however, @xcite method exhaustively enumerates all paths in the computation graph and ranks them by the amount of influence along each one. as the number of paths in a computation graph is usually exponential in the depth of a neural network, this quickly becomes intractable for larger networks @xcite . therefore, this method is limited to computing influence paths for networks with very small numbers of paths. indeed, while @xcite computed the influence along 40 000 paths for a 2-layer lstm, follow-up work that attempted to apply this method to bert had to use an approximation which might not find the correct paths @xcite . the method we propose does not exhibit this issue and scales to any network one can train using backpropagation. ## generalizing backpropagation in this section, we build toward our generalization of backpropagation as a semiring-weighted dynamic program. at a high level, we observe that if we replace the addition and multiplication operations in the typical backpropagation algorithm with similar operations that satisfy the necessary properties, then the resulting algorithm will compute other useful statistics about the network's gradient graph in the same runtime as backpropagation. in the remainder of this section, we make this notion of swapping operations precise by formulating backpropagation as a semiring algorithm, and later in §4 we describe how different semirings yield different, useful, views of the gradient graph. ## experiments to demonstrate the utility of semiring backpropogation, we empirically analyze their behavior on two simple transformer models (1-2 layers) on well-controlled, synthetic tasks. we also explore semiring backpropogation on a larger model, bert @xcite , on the popular analysis task of subject-verb agreement to understand how our method can be useful for interpreting language models in more typical settings. to implement semiring backpropagation, we developed our own python-based reverse-mode automatic differentiation library, building off of the pedagogical library brunoflow (ritchie, 2020) and translating it into jax @xcite ). 13 ## conclusion we presented a semiring generalization of the backpropagation algorithm, which allows us to obtain an alternative view into the inner workings of a neural network. we then introduced two semirings, the max-product and entropy semirings, which provide information about the branching points of a neural network and the dispersion of the gradient graph. we find that gradient flow reflects model component importance, gradients flowing through the self-attention mechanism for the subject token pass primarily through the keys matrix, and the entropy has some relationship with the difficulty of learning a task. future work will consider semirings outside the scope of this work, e.g., the top-k semiring @xcite to track the top-k gradient paths, as well as computing semirings online for control during training.