Reversed Attention: On The Gradient Descent Of Attention Layers In GPT (2025)

Abstract: The success of Transformer-based Language Models (LMs) stems from their attention mechanism. While this mechanism has been extensively studied in explainability research, particularly through the attention values obtained during the forward pass of LMs, the backward pass of attention has been largely overlooked.In this work, we study the mathematics of the backward pass of attention, revealing that it implicitly calculates an attention matrix we refer to as “Reversed Attention”.We visualized Reversed Attention and examine its properties, demonstrating its ability to elucidate the models’ behavior and edit dynamics.In an experimental setup, we showcase the ability of Reversed Attention to directly alter the forward pass of attention, without modifying the model’s weights, using a novel method called “attention patching”.In addition to enhancing the comprehension of how LMs configure attention layers during backpropagation, Reversed Attention maps contribute to a more interpretable backward pass.

## introduction the widespread use of automatic gradient technologies such as autograd @xcite to obtain the gradients may cause the explicit derivations of these gradients to be overlooked. in this work we explore the equations that govern the backpropagation of the popular transformer-based @xcite gpt @xcite architecture. this explicit analysis leads to several surprising discoveries. first, while during the forward pass the attention mechanism explicitly creates triangular attention matrices, from multiplying the query and key vectors, during the backward pass it implicitly creates triangular matrices that determine the gradients of the queries and keys. due figure 1 : in this paper we examine the attention maps obtained from the backward pass, which we named "reversed attention" (ra). this example present the forward and backward pass of a single attention head of gpt2-xl when prompt with "i like italy and france, i visited the city of". after the model answer "florence", a city in italy, we apply a backward pass with "florence" as the target for the loss and produce the ra maps. between all the 1200 attention heads this model has, the presented head has the highest ra map's norm. compared to the forward attention map, the ra map is more sparse and interpretable. this ra demonstrates how the backpropagation attempts to amplify the information from the token "italy" (red) while reducing the influence of "france" (blue). the similarity of these triangular matrices to the forward attention, we refer to them as the reversed attention (ra) matrices, figure 1 . secondly, we study the effect induced by gd and how it tries to increase or decrease attention scores between the queries and keys of each attention head. based on these discoveries, we explore the use of ra in lm explainability. while the forward pass attention, which fails in providing a clear explanation of the model's behavior (jain and wallace, 2019), we demonstrate ra's ability to achieve competitive results in concert with methods such as causal mediation @xcite in perturbation analysis. furthermore, one can view the ra as a correction term to the attention, given the loss of the backward pass. as an application, we inject the ra scores directly into the forward pass of attention, in order to modify the model's predictions. this novel method, which we call "attention patching", does not involve any parameter updates and offers a new perspective on how interventions can be performed on lms. our main contributions are as follows: (i) we provide a mathematical walk-through and interpretation of the gradients and vector-jacobian products (vjps) governing the backpropagation of gpt. (ii) we identify the attention-softmax derivative as an implicit attention map, which we term reversed attention (ra). (iii) we visualize and qualitatively explore the interpretability of ra. (iv) we conduct a perturbation test to quantify the explainability of ra. (v) we demonstrate a novel patching method that uses ra to edit lm's predictions. ## related work all leading deep learning models are trained using variants of gradient descent (gd), an implementation of the backpropagation algorithm. while much research examines the impact of gd on gpts' performance, the internal computations of this process often remain a black box @xcite @xcite . some studies simplify the transformer architecture to understand gd, for instance, reducing multi-head attention to a single head or linear attention in toy models @xcite @xcite . notably, the literature on the backward pass of full multi-head attention is limited. our work addresses this gap by examining assumption-free full gpt models, with a specific emphasis on detailing the mathematical computation of gradients. previous investigations into weight updates via gd have focused primarily on the mlp layers or the data it was trained on @xcite . @xcite has revealed that gradients can be interpreted as tokens' embeddings. specifically, the vector jacobian products (vjps) that are passed by the residual from the last layer's loss to earlier layers can be seen as an intermediate state toward the model as it adjusts its weights. gradients have been leveraged to localize the source of specific model behavior for a given input @xcite or specific components of the models @xcite . in our study, we explore the use of ra as a means to quantify the relative influence of the various components of a model. we compare our method with the current leading technique for such localization, causal mediation (cm) @xcite . cm involves probing the effect of altering components during the forward pass, which necessitates significantly more computation compared to the ra approach. recent works explore the patching technique, where one model's intermediate state is integrated into the forward pass of another @xcite @xcite . while they focus on activation patching, we propose a novel approach: directly injecting attention scoring maps of each attention head, without parameter updates. as far as we can ascertain, we are the first to identify, visualize, and explore the dynamics of the backward pass using ra maps. ## background this section provides the necessary background and establishes the notation used. generative pre-trained transformer (gpt) is an auto-regressive architecture of multiple transformer blocks connected via a residual stream. as input, a gpt model receives a sequence of n tokens (a prompt) and predicts a single token. an embedding matrix at the start of the model embeds the token into vectors @xmath0, • • • , x n ] ∈ r n×d where d is an embedding dimension. at the end of the model, the embedded predictions are projected back into tokens using a decoding matrix. each transformer block consists of two sub-blocks: multi-head attention (attn) and a multi-layer perceptron (mlp), interconnected by a residual stream. the attention mechanism is executed using matrices w q , w k , w v , w @xmath1, named query, key, value, and output, respectively. this calculation is performed after splitting the matrices vertically into h non-overlapping parts, called heads. we denote the attention matrices for the l head in ŵ l , hence ŵ l q , ŵ l v , ŵ l,⊤ @xmath2, and for example the first three matrices are used to project the input into queries, keys, and values: together, the queries and keys are used to calculate the forward attention scores: where @xmath3: . the output of each head is calculated by multiplying the attention scores with the values and projecting the result back using the output matrix. the attention block output is the sum of the output of all heads: attn(x) = h @xmath4, where head @xmath5. the mlp block is a pair of fully connected matrices f f 1 , f f ⊤ 2 ∈ r d×dm and an activation function f . the output of this block is: mlp(x) = f (xf f 1 )f f t 2 . lastly, the forward pass of the i-th transformer block on its input hidden state, note that gpt models also include layer norms. for simplicity and due to the relatively low contribution to the gradients and inconsistency when they are placed within different architectures, we omit them from this explanation. gradient decent, backward pass and vjps gd's backward pass is the execution of @xcite , the process of applying the chain rule to compute a model's gradients. a backward pass is initiated after the model executes a forward pass; it computes a loss score l, comparing the model's output with a desired target. this loss score is propagated back through the model's layers as an error signal, in the reverse order of the forward pass. the error signal can be represented as a vector that is used as an intermediate state of the backward pass, similar to the hidden state in the forward pass. given a model's parameter w , which is used to compute @xmath6, where @xmath7, @xmath8, the error signal is the loss with respect to the layer's output, δ = ∂l ∂@xmath9. this vector is known as the vector-jacobian product (vjp) of z. at the last layer of the model, the vjp is calculated directly by the loss function. for earlier layers, the vjp is calculated using the backward step, where the vjp of the next layer is used to compute the vjp of those that precede it. for instance, the output of a sequential layer l is the input of l + 1, meaning z @xmath10. given those layers are weight matrices, the vjp of the l layer is computed by the following step: finally, the gradient of each weight matrix w is the outer product of the layer's input x and the vjp δ, which updates the weights using a learning rate η ∈ r: in models such as gpts, each forward pass includes a sequence of inputs @xmath11, • • • , x @xmath12. in this case, each input has its own vjp, δ @xmath13, and the full matrix's gradient is the sum of the outer products of each input and its vjp: in this section, we examine the vjps and gradient matrices for each of the attention layer matrices. the purpose of this examination is to reveal properties of the gradient matrices, as well as to fill a missing gap in the literature. in this mathematical walk-through, we examine the submatrices of each attention head, ŵ l q , ŵ l k , ŵ l v , ŵ l o , dropping the head index l for brevity. the gradients of the full matrices are the concatenation of the heads' gradients (same as the concatenation of weight matrices). throughout this section, we denote the forward pass input and output vectors of each submatrix ŵ ∈ r d× d h by @xmath14, i.e., @xmath15. the only exception is w o , where the dimensions are swapped, ŵo ∈ r d h ×d . a i,j denotes the forward pass attention score from the i-th token to the j-th token. lastly, considering the model's input as a sequence of n tokens, instead of explicitly writing the gradient of every matrix and token, we will focus on determining the vjp for a single token @xmath16, • • • , n}, annotating its input and vjp with x j , δ j . given the inputs and vjps, the gradients are the outer product of the two equation 7 . the output projection matrix ŵo the vjps of ŵo , denoted as δ j o , are obtained directly from the residual stream at the end of the attention block of ŵo , denoted by attn(x) in section 3. viewing gd as an application of the chain rule, we consider δ j o as the layer's intermediate editing target. only for this matrix, we will demonstrate that when we are given its vjps, we can infer its gradient and the effect of gd updating. this explanation holds for all further matrices, too. the gradient updates introduced by only the j-th token (left equation) or by all n tokens (right equation) are: the value projection matrix ŵv in the forward pass, the value vector associated with the j-th token is given by v @xmath17. to calculate the vjp for the j token, we consider, due to causality, the subsequent tokens @xmath18. for every such token, the backward pass propagates its own vjp from ŵo by computing the error signal from the l-th token to token j: softmax derivative although the softmax function is not a weight parameter of the model, comprehending its role in the backward pass is crucial for understanding the editing dynamics. in the forward pass, the attention scores for the j-th forward pass, a @xmath19, with a scaling factor of d h . these scores are then multiplied by the attention values, @xmath20. the backward pass performs a reverse operation: each token's error signal from ŵo is multiplied by the attention values and scaled as follows: the query projection matrix ŵq during the j-th forward pass, the j-th query is generated by computing q @xmath21. the query is then multiplied by all key vectors, @xmath22, • • • , k n ] ∈ r n× d h (and masking the ones that followed it), to obtain the raw attention scores (logits). therefore, the backward pass calculates the vjp of the query by calculating: the key projection matrix ŵk the gradients of ŵk are computed similarly to those of ŵq , except that for the j-th forward pass, we utilize the queries ## reversed attention in equation 16 we defined r as the softmax derivative. r shares many properties with a, the forward attention (fa): these shared properties between the fa and the softmax derivative suggest that the softmax derivative serves as an implicit attention matrix. we called r the "reversed attention" (ra). in the following sections, we delve into some properties of the ra and explore its potential uses in explaining and controlling gpts. ## conclusions the self-attention component of transformers is perhaps the most distinctive part of this architecture. its role when performing inference has been extensively studied and shown to provide insights into the inner workings of transformers. here, we explore a dual entity we call reversed attention (ra), which plays a role when transformers learn. we present qualitative samples of the way learning occurs and show that ra can help identify the most influential heads at inference time. finally, figure 5 : attention patching using reversed attention (ra): first we collect the ra maps of the model without applying any model editing (without changing its weights). later, for each attention head, we add its corresponding ra map to the forward pass attention. the focus of this work was to introduce ra and prove that it is interpretable rather than merely an artifact of the backward pass. beyond enriching our understanding of how lms work, we hope our conclusions will serve future research on dynamically editing attention layers as well as explaining how information is stored in them.