One-sentence Summary

Researchers at Technion and Tel Aviv University propose Backward Lens, which extends vocabulary-space projection to language model gradients by proving that gradient matrices can be expressed as low-rank linear combinations of forward and backward pass inputs, thereby revealing how new information is stored in transformer neurons.

Key Contributions

• Logit Lens-style interpretability is extended to the backward pass by proving that gradient matrices are low-rank linear combinations of inputs from the model’s forward and backward passes.

• Gradient components are projected into vocabulary space, revealing an “imprint and shift” mechanism in which layer inputs and vector-Jacobian products store new information in Transformer MLP neurons.

• Experiments support the analysis and introduce a knowledge-editing method that requires only a single forward pass and performs on par with state-of-the-art knowledge-editing methods.

Introduction

Transformer-based language models are trained with backpropagation, which updates weight matrices through gradient information, making gradients central to how models learn, store, and apply knowledge. Interpreting these gradients is important because it could clarify how new facts are written into model parameters and enable more transparent knowledge editing. Prior interpretability work has projected weights and forward-pass activations into vocabulary space, for example through Logit Lens methods, but it largely leaves backward-pass gradients unexplored, even though modern layers contain thousands of neurons and gradients are distributed across them. The authors leverage the structure of gradient matrices, which are generated from forward-pass inputs and Vector-Jacobian Products, and their typically low-rank nature, to project gradient subspaces into token space and analyze what information they carry. Based on this, they reveal an “imprint and shift” mechanism in MLP layers that explains how language models acquire and update knowledge, and they propose a simple knowledge-editing method based on a single forward pass.

Method

The framework centers on the backward pass of an autoregressive Transformer language model, treating gradients not merely as optimization signals but as low-rank geometric objects that reveal how information is written into a model. In a standard Transformer block, the MLP component is the main editable module under study. Given hidden states $X^l$Xl, the block update can be summarized as:

$$X^{l+1} = X^l + \mathrm{Attn}(X^l) + \mathrm{MLP}(\mathrm{Attn}(X^l) + X^l),$$Xl+1=Xl+Attn(Xl)+MLP(Attn(Xl)+Xl),

where the MLP is formed by two fully connected matrices, $FF_1$FF1​ and $FF_2$FF2​, with a non-linear activation between them. This structure makes the two matrices the primary targets for interpretation and editing, because the forward pass writes intermediate inputs into the MLP parameters, while the backward pass modifies those parameters according to gradients computed from a prompt and a desired target.

The first analytical step is to characterize the shape of the gradient matrices. For a linear or matrix-valued parameter, a backward pass produces a gradient that is a sum of outer products between forward-pass hidden states and backward-pass gradients, commonly called vector-Jacobian products. For a sequence of $n$n tokens, the gradient of a parameter matrix $W$W can be written as:

$$\frac{\partial L}{\partial W} = \sum_{i=1}^{n} x_i^{\top} \cdot \delta_i,$$∂W∂L​=i=1∑n​xi⊤​⋅δi​,

where $x_i$xi​ is the input vector observed in the forward pass and $\delta_i$δi​ is the corresponding backward-pass vector. Because this gradient is a sum of $n$n rank-one matrices, its rank is at most $n$n, and it may be smaller when some vectors are linearly dependent. In a Transformer, this implies that gradients induced by a single prompt are structurally simple and can be analyzed through a small number of forward and backward vectors rather than through the full high-dimensional weight matrix.

This low-rank perspective leads to a key interpretability choice: each gradient matrix can be viewed in two dual ways, either as a span of the forward inputs $x_i$xi​ or as a span of the backward vectors $\delta_i$δi​. The two views are mathematically equivalent for a single outer product, but they become analytically distinct once the vectors are projected into the vocabulary space. Refer to the figure below:

For $FF_1$FF1​, the forward inputs $x_i$xi​ are used as the spanning set. The $j$j-th neuron of the gradient can then be interpreted as a linear combination of the hidden states:

$$\frac{\partial L}{\partial FF_1}[j] = \sum_{i=1}^{n} x_i^{\top} \cdot \delta_i[j].$$∂FF1​∂L​[j]=i=1∑n​xi⊤​⋅δi​[j].

Here, $\delta_i[j]$δi​[j] controls the magnitude with which each forward input $x_i$xi​ contributes to that neuron. For $FF_2$FF2​, the spanning choice is reversed: the gradient is analyzed as a combination of the backward vectors $\delta_i$δi​:

$$\frac{\partial L}{\partial FF_2}[j] = \sum_{i=1}^{n} \delta_i \cdot x_i[j].$$∂FF2​∂L​[j]=i=1∑n​δi​⋅xi​[j].

In this formulation, the forward inputs act as scalar coefficients that scale the backward vectors. This choice is important because the two matrices receive vectors of different semantic roles: $FF_1$FF1​ tends to capture representations of the current prompt, while $FF_2$FF2​ tends to encode the backward signal that steers the output toward the target token.

To understand the backward signal itself, the framework examines the VJPs at the top layer. The final model parameter is the decoding matrix $D$D, which maps hidden states to vocabulary scores. If $\hat{y}_n$y^​n​ is the output logits for the last prompt token, and $\hat{p} = \mathrm{Softmax}(\hat{y}_n)$p^​=Softmax(y^​n​) is the predicted distribution, the initial backward vector passed to the preceding layer is:

$$\frac{\partial L}{\partial x_n} = \delta_n D^{\top}.$$∂xn​∂L​=δn​D⊤.

This vector can be decomposed as a weighted sum of token embeddings:

$$\delta_n D^{\top} = \sum_{k=1}^{|\mathrm{vocabulary}|} \delta_n[k] D^{\top}[k].$$δn​D⊤=k=1∑∣vocabulary∣​δn​[k]D⊤[k].

The coefficient for the target token $t$t is approximately $\hat{p}[t] - 1$p^​[t]−1, which is typically negative, while coefficients for non-target tokens are positive probabilities $\hat{p}[k]$p^​[k]. Therefore, the initial backward vector contains a negative embedding of the desired target and positive embeddings of currently plausible alternative tokens. When this vector is propagated backward, it carries the semantic pressure that pushes the model from the old prediction toward the new target.

The central mechanism discovered by the authors is “imprint and shift.” During a single backward pass, the gradient update does not act on $FF_1$FF1​ and $FF_2$FF2​ in the same way. Instead, each matrix plays a different functional role. For $FF_1$FF1​, the update introduces scaled copies of the forward inputs $x_i$xi​ into the weight neurons. This is the imprint phase: the layer learns to associate the current input pattern with the update, changing how strongly certain hidden activations will trigger downstream neurons. For $FF_2$FF2​, the update introduces scaled backward vectors $\delta_i$δi​ into the weight neurons, usually with a negative sign because gradient descent moves against the gradient. This is the shift phase: the layer alters its output representation so that the projection of the final hidden state moves toward the target token. As shown in the figure below:

In this mechanism, $FF_1$FF1​ adjusts which internal features are selected from the prompt, while $FF_2$FF2​ adjusts how those selected features are written into the output vocabulary space. The result is a two-step update: the model imprints the input context into the first MLP matrix, and then shifts the second MLP matrix so that the final decoding path favors the edited target. This explains why editing knowledge in Transformer LMs is not merely a scalar weight adjustment but a structured interaction between prompt representations and target-specific backward signals.

The framework also turns this insight into a lightweight editing procedure. Because the backward vector near the top layer is strongly related to the target token embedding, a full backward pass can be approximated by using the target embedding directly. The editing process proceeds as follows. First, choose an MLP layer containing $FF_2$FF2​. Then run a single forward pass on the prompt to be edited, collecting the last hidden input $x_n$xn​ that reaches the chosen layer. Next, retrieve the target token embedding from the decoding matrix as $D^{\top}[t]$D⊤[t]. Finally, update $FF_2$FF2​ using the outer product between the prompt input and the target embedding:

$$FF_2 \leftarrow FF_2 + \eta \cdot x_n^{\top} \cdot D^{\top}[t],$$FF2​←FF2​+η⋅xn⊤​⋅D⊤[t],

where $\eta$η is a learning-rate-like scalar controlling the strength of the update. This method removes the need to compute gradients through multiple Transformer layers, because the target embedding replaces the backward vector that would normally be produced by the loss gradient. The underlying rationale is that $FF_2$FF2​ is the matrix most directly associated with the “shift” phase, so adding a prompt-conditioned target vector is sufficient to bias the layer toward the desired output.

Overall, the method provides a geometric interpretation of model editing. The authors treat gradients as sums of forward and backward vectors, separate the distinct roles of the two MLP matrices, and use vocabulary-space projections to interpret what the gradients are writing into the model. The resulting pipeline connects the structure of backpropagation to an interpretable mechanism: $FF_1$FF1​ imprints the input, $FF_2$FF2​ shifts the output, and both together determine how a single prompt can change the internal knowledge of a Transformer language model.

Experiment

The experiments use GPT2 and Llama2-7B on sampled CounterFact editing examples to examine the structure, interpretability, and application of model gradients. They support the predicted gradient rank behavior, show through logit lens analysis that gradients largely encode the editing target, and reveal that update contributions are concentrated in subject tokens in early layers and the final prompt token in middle layers. A further application test on CounterFact indicates that a forward pass based update can match strong single editing baselines with much lower complexity, though it is weaker on generalization and specificity.

The evaluated single editing approach changes GPT2 XL behavior on counterfactual prompts with efficacy close to strong baseline methods while preserving generation fluency at a broadly comparable level. It is less robust than the strongest baseline on paraphrased prompts and shows lower specificity on nearby prompts that should remain unchanged. Edited models substantially increase the success rate on directly edited prompts compared with the unedited model. The forward pass based editing method trades some generalization and specificity advantages for a much lower runtime editing procedure.

The experiments evaluate single-editing methods on counterfactual prompts, comparing them with strong baselines in terms of editing efficacy, fluency, robustness to paraphrases, and specificity on nearby prompts. The results show that the evaluated approach achieves competitive efficacy and fluency and greatly improves success on edited prompts, though it is less robust and specific than the strongest baseline. The findings also indicate that forward pass based editing trades some generalization and specificity for much faster runtime.