Backward Lens: Projecting Language Model Gradients into the Vocabulary Space (2024)

Abstract: Understanding how Transformer-based Language Models (LMs) learn and recall information is a key goal of the deep learning community. Recent interpretability methods project weights and hidden states obtained from the forward pass to the models’ vocabularies, helping to uncover how information flows within LMs. In this work, we extend this methodology to LMs’ backward pass and gradients. We first prove that a gradient matrix can be cast as a low-rank linear combination of its forward and backward passes’ inputs. We then develop methods to project these gradients into vocabulary items and explore the mechanics of how new information is stored in the LMs’ neurons.

## introduction deep learning models consist of layers, which are parameterized by matrices that are trained using a method known as backpropagation. this process involves the creation of gradient matrices that are used to update the models' layers. backpropagation has been playing a major role in interpreting deep learning models and multiple lines of study aggregate the gradients to provide explainability @xcite @xcite @xcite . recent interpretability works have introduced methods to project the weights and intermediate activations of transformer-based lms @xcite into the vocabulary space. the seminal "logit lens" method (nostalgebraist, 2020) has paved the way to explaining lms' behavior during inference @xcite @xcite , including directly interpreting individual neurons @xcite . our work is the first, as far as we can ascertain, to project lm gradients to the vocabulary space. furthermore, modern lms contain figure 1 : an illustration depicting the tokens promoted by a single lm's mlp layer and its gradient during the forward and backward pass when editing the model to answer "paris" for the prompt "lionel messi plays for". the gradients (in green) of the first mlp matrix, f f 1 , attempt to imprint into the model's weight (in blue) the information that f f 1 encountered during the forward pass. utilizing a vocabulary projection method, we reveal that this information represents the token "team". the gradients of the second mlp matrix, f f 2 , aim to shift the information encoded within f f 2 towards the embedding of the new target. despite the popularity of lms, our understanding of their behavior remains incomplete @xcite , particularly regarding how lms acquire new knowledge during training and the mechanisms by which they store and recall it @xcite @xcite . we identify a mechanism we refer to as "imprint and shift", which captures how information is stored in the feed-forward (mlp) module of the transformer layer. this module has two fully connected layers, f f 1 and f f 2 . the "imprint" refers to the first layer, to or from which the learning process adds or subtracts copies of the intermediate inputs encountered during the forward pass. the "shift" refers to the second matrix, where the weights are shifted by the embedding of the target token; see figure 1 . in summary, our contributions are: (i) analyzing the rank of gradients. (ii) interpreting gradients by inspecting relatively small spanning sets. in particular, we examine the vector-jacobian product (vjp) obtained during the backward pass, which is the equivalent of the hidden states of the forward pass. (iii) investigating the embedding of these sets by projecting them into tokens. (iv) revealing a two-phase mechanism by which models learn to store knowledge in their mlps, which we termed "imprint and shift". (v) exploring a novel editing method based solely on a single forward pass. ## related work developing methods to explain lms is central to the interpretability community @xcite . initially inspired by interpretability efforts in vision models @xcite @xcite @xcite , lms have also benefited from the ability to operate in the language domain. this includes leveraging projections of vectors into readable concepts @xcite or clustering them into the idioms they promote @xcite @xcite . reverse engineering the gradient's role in shifting model behavior has been a primary method to comprehend the mechanics of deep learning models. recent work @xcite @xcite demonstrates that clustering the weights that models learn during training or fine-tuning reveals patterns that connect the tasks and their training data. other approaches employing saliency maps @xcite explore the relationship between a gradient matrix and parts of its corresponding forward pass input, where we have observed that gradients are spans, linear combinations, of those inputs. our work differs from previous studies in two key aspects: (1) we interpret gradients based on the components that generate them, primarily the vjps, rather than analyzing entire gradient matrices, and (2) we elucidate the information stored in the gradients and project them into tokens, instead of solely examining their impact on the model's predictions or their connection to the training data. to the best of our knowledge, we are the first to explore both projecting gradients into tokens and examining them through vjps, offering a perspective that has yet to be explored. our experiment regarding lm editing adds to the line of work that utilizes interpretability for knowledge editing. the closest idea to our implementation was introduced by @xcite , who identified activated neurons for specific idioms in encoder lms and altered them by injecting the embedded target. we show that gradients work in a very similar way. other state-of-the-art model editing methods include @xcite and @xcite . our work analyzes the gradient's rank. it was previously known that gradients are low-rank @xcite , but the utilization of this characteristic for interpretability study or predicting the rank of an edited prompt have remained unexplored. optimizers and adaptors, such as lora @xcite , were created to constrain the rank of gradients, making fine-tuning faster. in contrast, our work demonstrates that the ranks of gradients can be predicted based on the edited prompt length. ## background we first provide background on transformers, focusing on the components that play a role in our analysis and omitting other concepts, such as layer norms and positional embedding, which are explored in full by @xcite and @xcite . we then provide the necessary background on the backward pass; a more comprehensive description is given by @xcite and @xcite . finally, we discuss the building blocks of the logit lens method. ## backward lens in this section, we detail the methods we developed to analyze gradients based on our understanding of how each gradient matrix is formed. in the following sections, we consider the editing of a single sequence (prompt) to answer a single token (target) using nll loss and without employing batching. ## experiments we conduct a series of experiments to support the results of sec. 4 and 5, as well as to briefly demonstrate their application to lm analysis. we employ gpt2 @xcite and llama2-7b @xcite in our experiments. we randomly sampled 100 prompts and their corresponding editing targets from the coun-terfact dataset @xcite . for each model and prompt, we conducted a single backpropagation using sgd and without scaling optimizers, such as adam (kingma, 2014), and no batching. the rank of the gradients to examine lemma 4.1, we measure the rank of each layer gradient matrix. as depicted in figure 3 , for every prompt with the length of n tokens, the model's gradient matrices are almost always exactly rank n. the only exceptions are the last mlp layers, which have a rank of 1, as predicted in section 4. although unnoticeable from the figure, once in a few dozen examples, there is a drop of one or two in the rank of the gradients, indicating linear dependency in x i or δ i , see section 4.1. this is not a result of a repeated token, since the positional encoding would still lead to a different x i . logit lens of gradients next, we present examples of our gradients' interpretation through ll in figure 4 and in appendix e. in each cell of these plots, the ll projections of the chosen spanning set (f f 1 's x i and f f 2 's δ i ) are presented for a specific layer and a token from the prompt that was used for the editing. prior studies that projected the forward pass examine the ll projections of hidden states, highlighting the gradual change in the projected tokens between layers (nostalgebraist, 2020; @xcite . similarly, figure 4a presents a gradual change in the backward pass' vjp. across most layers, ll reveals that the gradients represent the embedding of "paris". other projections have semantics that are related to "paris" such as "macron", the family name of the president of france. the norm of the vjp is indicted by color, and, in the top layers, the only meaningful updates are for the token "paris". some of the edits in the lower layers are harder to explain, similarly to the situation in those layers for the vanilla ll of the forward pass. another example, presented in figure 4b , shows similar patterns: when attempting to edit the model to suggest that jack dorsey, known as one of twitter's founders, founded ibm, we observe that most gradients correspond to the embedding of "ibm". ## application: editing based on the "shift" mechanism prior work @xcite @xcite introduced editing methods that change only f f 2 's matrices. in section 5.2 we identify the "shift" mechanism of editing f f 2 . in section 6 we observe that the dominant components in constructing the gradients are derived from the outer product of the last token's input and vjp, x ⊤ n • δ n , and that δ n contains the embedding of the target token. we hypothesize that we can edit lms' internal knowledge by updating only a single f f 2 matrix with a single forward pass and an approximation of the vjp, thereby eliminating the need for the backward pass. based on lemma 5.1, the embedding of the editing target is annotated by d ⊤ [t], where d is the decoding matrix and t is the index of the target token. our experimental method works as follows: (i) we choose an mlp layer we wish to edit, f f 2 (predefined as a hyper-parameter according to appendix i). (ii) we run a single forward pass with the prompt whose output we want to edit. (iii) during the forward pass, we collect the last token input for the layer we want to edit, x n . (iv) we collect the embedding of the target token d ⊤ [t] and (v) update the mlp matrix by f f we examined our method on 1000 samples from counterfact ( ## conclusions other ll-type interpretability contributions shed light on lms through the forward pass. here, we show that gradients can be projected into the vocabulary space and utilize the low-rank nature of the gradient matrices to explore the backward pass in an interpretable way. as we show, the gradients are best captured by a spanning set that contains either the input to each layer, or its vjp. these two components, which are accessible during the forward and backward passes, are used to store information in the mlp layers, using a mechanism we call "imprint and shift". we provided experimental results to substantiate the results of our analysis, including an editing method that only requires a single forward pass, but is on par with the sota knowledge editing methods.