PaperlayerView exact revision on arXiv

Source-linked AI summary

Backward Lens: Projecting Language Model Gradients into the Vocabulary Space

Shahar Katz, Yonatan Belinkov, Mor Geva, Lior Wolf

arXiv:2402.12865v1cs.CLcs.AIcs.LG

TL;DR

Understanding how Transformer language models acquire, store, and recall knowledge remains incomplete, especially for backward-pass mechanisms. This paper proves that gradients have a low-rank structure tied to forward and backward inputs, projects them into vocabulary items, and uses them to study MLP knowledge storage. It identifies “imprint and shift” and reports a single-forward-pass editing method on par with state-of-the-art knowledge-editing methods.

  • Problem

    How Transformer language models acquire, store, and recall new knowledge remains incompletely understood, particularly in their backward pass and neuron-level storage mechanisms.

  • Method

    The paper decomposes low-rank gradient matrices into forward-pass inputs and backward-pass VJPs, projects these components into vocabulary items, and analyzes their role in MLP storage and editing.

  • Results

    The gradients are almost always rank n for prompts of length n, with last-MLP-layer exceptions of rank 1, and a single-forward-pass editing method is reported on par with state-of-the-art knowledge-editing methods.

  • Takeaways & Limitations

    The results support an “imprint and shift” account in which FF1 incorporates forward inputs while FF2 shifts information toward the target-token embedding.

  • Takeaways & Limitations

    The analysis uses the original Logit Lens projection despite uncertainty about whether earlier-layer gradients share the same embedding space, and it ignores optimizer scaling beyond SGD.

Abstract

from arXiv · show

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.

1 Introduction

The paper extends vocabulary-space interpretability from forward-pass weights and activations to backward-pass gradients, addressing how language models store and recall newly learned information. It identifies an “imprint and shift” mechanism in MLP layers and develops low-rank, vocabulary-based analyses of gradients.

  • The work is the first described effort to project language-model gradients into vocabulary space, extending prior projections of weights and intermediate activations.
  • Figure 1 illustrates an edit toward “Paris” for “Lionel Messi plays for,” where FF1 gradients reveal the token “team” and FF2 gradients shift information toward the new target.
  • Gradient matrices are analyzed through provably low-rank decompositions that capture information from forward and backward passes despite distributed features across many neurons.
  • The proposed “imprint and shift” mechanism stores information in MLPs by adding or subtracting forward-pass inputs in FF1 and shifting FF2 toward the target-token embedding.
  • The paper interprets gradients through small spanning sets, projects those sets into tokens, examines backward-pass VJPs, and explores single-forward-pass knowledge editing.

2 Related Work

Related work has used vocabulary projections, gradient analyses, and interpretability for knowledge editing, but primarily studied static weights or forward-pass states. This paper instead exploits gradients’ existing low-rank structure and backward-pass VJPs.

  • Prior interpretability work projects vectors into readable concepts or clusters, while gradient studies often inspect full matrices or connect learned weights to tasks and training data.
  • The paper differs from earlier forward-pass and static-weight approaches by interpreting the backward pass through its Vector-Jacobian Product.
  • Knowledge-editing methods have identified activated neurons or modified model parameters, whereas this work shows that gradients can provide a similar route for injecting target information.
  • Although low-rank gradients were previously known, this work uses that property for interpretability and for predicting the rank associated with an edited prompt.

3 Background

The background introduces Transformer and MLP computations, backpropagation and VJPs, and Logit Lens vocabulary projections. These components provide the forward- and backward-pass quantities used to interpret gradients.

  • A GPT maps an n-token prompt through embedding dimension d, transformer blocks, and a decoding matrix that scores vocabulary tokens.
  • Each transformer block combines attention with an MLP, whose two matrices FF1 and FF2 surround an activation function and operate through the residual stream.
  • Backpropagation applies the chain rule after a forward prediction and loss calculation to compute gradients for updating model parameters.
  • The VJP is the backward-pass analogue of a hidden state and represents the error factor propagated from later layers.
  • A layer gradient is formed from the outer product of its forward-pass input x and the output VJP δ.
  • Logit Lens applies the decoding transformation to intermediate hidden states to obtain vocabulary predictions, with later-layer projections generally more interpretable than earlier ones.
  • Existing neuron studies use vocabulary projections for MLP neurons, while this paper applies the approach to the backward pass rather than only static weights or forward states.

4 Backward Lens

The section shows that LM gradient matrices are low-rank sums of forward- and backward-pass vectors, enabling vocabulary-space analysis through small spanning sets.

  • 4 Backward Lens: The method applies Logit Lens projections to these spanning vectors to interpret gradient neurons in vocabulary space while avoiding exhaustive inspection of every matrix dimension.
  • 4.1 Gradients as Low-Rank Matrices: A gradient from a sequence of n inputs has rank at most n, with lower rank arising from linear dependencies among forward or backward vectors.For a single input, the outer-product gradient has rank 1, or rank 0 if either vector is zero.
  • 4.1 Gradients as Low-Rank Matrices: The final transformer layer has a rank-one gradient, making it a particularly compact case for backward-pass analysis.
  • 4 Backward Lens: Each gradient matrix can be examined as a sum of n outer products, reducing analysis from the full MLP dimension to n much smaller vectors.The outer product can be viewed simultaneously as a span of forward inputs x_i or backward vectors δ_i.
  • 4.2 Applying Logit Lens to Gradient Matrices: For FF1, the forward inputs x_i form the spanning set because each gradient neuron is a linear combination of them.For FF2, the dimensions are reversed, so the backward vectors δ_i are selected as the spanning set.

5 Understanding the Backward Pass

The backward pass is analyzed through VJPs that can be projected into token embeddings, revealing how a single edit imprints inputs in FF1 and shifts FF2 toward a target.

  • 5 Understanding the Backward Pass: The initial backward vector is a weighted sum of token embeddings, dominated by the target token’s embedding with a negative coefficient while other tokens receive positive coefficients.This VJP is the vector passed backward from the decoding matrix and spans the last MLP gradient.
  • 5 Understanding the Backward Pass: The target token is expected to receive the lowest Logit Lens probability for the initial VJP, while related tokens can be affected by embedding and logit similarities.
  • 5.2 Storing Knowledge in LMs: Updating FF1 adds or subtracts forward inputs, adjusting future activation strength, while updating FF2 subtracts VJPs and amplifies their presence in the output.These effects are established when the edited layer is rerun with the same prompt inputs.
  • 5.2 Storing Knowledge in LMs: The resulting two-phase “imprint and shift” mechanism first imprints an input through FF1 and then shifts FF2’s output toward the new target.The mechanism describes knowledge storage during a single backward pass.
  • 5.2 Storing Knowledge in LMs: For FF1, Logit Lens rankings emphasize the most probable tokens, whereas subtracting FF2’s VJPs reverses rankings and promotes previously least probable tokens.

6 Experiments

Experiments show that gradient matrices are usually prompt-length rank and that their vocabulary-projected backward-pass components concentrate on a few prompt tokens, layers, and target-related directions.

  • The rank of the gradients: For prompt length n, gradient matrices are almost always exactly rank n, except the last MLP layers, whose rank is 1.Ranks occasionally drop by one or two because forward- or backward-pass vectors become linearly dependent.
  • Logit Lens of Gradients: Across most layers, FF2 VJPs project to the target “Paris” or related tokens such as “Macron,” while top-layer updates mainly represent “Paris”.Lower-layer projections are harder to interpret, paralleling limitations of vanilla forward-pass Logit Lens.
  • Impact of Different Segments of the Prompt: Most gradient contributions have near-zero VJP norms, and individual gradient neurons’ projections correlate with only 1–2 spanning-set vectors.Thus, although all prompt tokens contribute formally, only a small subset contributes substantially to the interpretable gradient structure.
  • Impact of Different Segments of the Prompt: Predominant FF2 updates come from subject tokens in initial layers and the last prompt token around the second quarter of layers.Most other tokens have near-zero norms throughout the network and have almost no effect on the update; changes to the last subject token may involve information transferred through attention.
  • Logit Lens of Gradients: The last token’s VJP consistently ranks the target among the least probable tokens, while other prompt-token VJPs generally show the same behavior.Rank degradation in the first and last layers is associated with near-zero VJP norms; normalization improves target-token visibility, while early-layer readability may remain limited.

7 Application: Editing Based on the “Shift” Mechanism

The authors use the shift mechanism to derive a forward-pass-only editing rule that injects the target embedding into FF2. In CounterFact single editing, it matches established methods on the edited prompt but is weaker on paraphrase generalization and specificity.

  • Editing Based on the “Shift” Mechanism: The dominant gradient component is the outer product x_n^T · δ_n, and δ_n contains the target-token embedding.This observation motivates replacing the backward-pass component with D^T[t] in the proposed update.
  • Editing Based on the “Shift” Mechanism: Forward pass shifting updates FF2 with η · x_n^T · D^T[t], combining the last-token input with the target token’s decoding-matrix embedding.The method approximates the relevant VJP and avoids a backward pass.
  • Editing Based on the “Shift” Mechanism: The method uses a single forward pass and has lower runtime complexity than multi-step editing methods.The authors present it as evidence that fine-tuning may admit shortcuts based on directly injecting tokens into LM layers.
  • Editing Based on the “Shift” Mechanism: Single editing is on par with MEND, ROME, and MEMIT for editing the given prompt, but falls short of ROME on paraphrase generalization and specificity.The comparison uses GPT2-xl single-editing results on CounterFact, where EFF measures editing success, PAR measures paraphrase accuracy, and N-gram measures fluency.

8 Conclusions

The paper projects gradients into vocabulary space by exploiting their low-rank structure and interprets backward-pass VJPs as counterparts to forward-pass inputs. It uses this view to describe MLP knowledge storage and motivate forward-pass-only editing.

  • 8 Conclusions: Gradients can be projected into vocabulary space and interpreted through small spanning sets formed from layer inputs or backward-pass VJPs.The analysis uses these components to make gradient structure readable in token space.
  • 8 Conclusions: The paper identifies an “imprint and shift” mechanism in which forward-pass inputs and VJPs store information in MLP layers.The mechanism assigns imprinting to FF1 and shifting to FF2.
  • 8 Conclusions: Experimental results include an editing method requiring only a single forward pass that is on par with state-of-the-art knowledge-editing methods.The conclusion presents this as an application of the gradient analysis rather than as a complete replacement for existing methods.

9 Limitations

The analysis is scoped to specific architectures, editing settings, model components, and simplified theoretical assumptions. Its gradient projections are also less reliable for earlier layers, and the editing case study has limited evaluation scope.

  • 9 Limitations: Logit Lens explanations are limited for earlier layers because the appropriate embedding-space transformation remains unclear.The paper uses the original projection despite this uncertainty and notes that later-layer token embeddings are more directly justified.
  • 9 Limitations: The analysis ignores optimizer scaling beyond SGD, although optimizers such as Adam can alter gradient rank and weights.The backward-pass VJPs remain unaffected, but the resulting gradient matrices do not.
  • 9 Limitations: The knowledge-storage analysis assumes single editing with a constant embedding, whereas typical fine-tuning uses multiple prompts and iterations.Training from scratch also updates embeddings, which this analysis does not cover.
  • 9 Limitations: Forward pass shifting is presented as a case study, and benchmark outputs may not fully capture effects on similar prompts.The authors caution that this issue can obfuscate the distinction between editing and output shifting.
  • 9 Limitations: The study focuses on MLP layers and excludes possible information storage in attention layers.Editing MLPs and attention jointly could produce effects different from those analyzed.
  • 9 Limitations: The theoretical analysis omits Dropout, Layer Norm, positional embeddings, and bias vectors, whose effects on backward-pass interpretation may differ.These omissions simplify the derivations but constrain their direct applicability to full LM architectures.
  • 9 Limitations: The experiments use decoder LMs with sequential architectures, so other LM types may exhibit different gradient behavior.This is an explicit architectural scope boundary.

10 Ethics and Impact Statement

The paper reports no societal consequences requiring specific emphasis, while noting that its methods could be used for language-model editing and safer-model development.

  • The authors identify no societal consequences that they feel require specific highlighting.
  • Future research could use the developed methods to edit language models.
  • The authors hope such editing will support better and safer models rather than harmful content.

B Why Decomposed Gradient Analysis Makes Sense

The paper argues that gradient matrices can be understood through compact spanning sets rather than individual neurons, revealing how forward and backward information is written into MLP weights.

  • High-norm FF2 gradient neurons align with the final VJP δ_n and project the editing target “Paris”, whereas low-norm neurons remain close to the zero vector.The GPT2-medium example edits “Lionel Messi plays for” from “Barcelona” to “Paris”.
  • Positive and negative activations both add “Paris” in the same direction because the negative learning rate reverses the update for positive versus negative activations.
  • FF1 gradient neurons align with their input spanning set x_i, which represents the embeddings injected into FF1 weights during editing.
  • Spanning-set analysis simplifies thousands of gradient neurons into representative forward inputs x_i or backward VJPs δ_i, reducing computational cost and analysis time.The authors report convergence between individual-neuron Logit Lens analysis and spanning-set interpretations.
  • FF2 VJPs represent backward-pass information complementary to FF1’s forward-pass inputs, enabling vocabulary projections of information stored in gradients.
  • The same editing pattern appears across MLP modules: editing concentrates around the first quarter of layers for the last subject token and around the middle for the last token, while most other VJP norms are near zero.The study did not examine attention modules, and it reports no correlation between x_i and δ_i norms.

E.2 The Ranks of FF1 and the Models’ Original Answer

The analysis tracks how target and original-answer tokens appear in gradient and forward-pass projections across layers, while highlighting Logit Lens sensitivity and the scope of the editing experiments.

  • The reported experiments use 100 edits with one backward pass per edit across GPT2 and Llama2-7B, with prompts and targets from CounterFact.
  • Across GPT2 and Llama2-7B, FF2 gradient projections rank the target token among the most improbable tokens through most layers, with some degradation in early layers.The authors associate the early-layer drop with limitations of Logit Lens projections.
  • The original model answer receives relatively low Logit Lens rank in later layers, suggesting that updates reduce its output probability more smoothly than they increase the target probability.
  • FF1 gradients write forward-pass inputs x_i into model weights, producing gradual intermediate prediction patterns for correctly answered CounterFact prompts.Correctly answered prompts account for approximately 84% of instances across GPT2-medium, GPT2-xl, and Llama.
  • The target-token ranks in FF1’s x_i projections reflect the intermediate predictions that models produce at each layer.
  • Logit Lens is sensitive to low-norm vectors, whose projections can resemble the zero vector; normalizing VJPs increases the visibility of target tokens.The paper presents normalized Logit Lens examples for GPT2-medium and GPT2-xl.

G.3 Results

Forward pass shift achieves strong knowledge-editing results with minimal impact on generic text generation, but its neighborhood editing remains weaker because similar prompts can be altered unintentionally.

  • After filtering methods that reduce n-gram entropy below 620, the method achieves the best edited-prompt Accuracy, tying only with ROME.
  • On paraphrased prompts, forward pass shift performs comparably but falls behind ROME.
  • Neighborhood editing remains a limitation because the edited model can change prompts that were not intended to change.This may result from shared FF2 activation patterns across similar prompts.
  • Forward pass shift achieves successful knowledge editing while requiring only a single forward pass, making it simpler than more complex methods.The method is described as the simplest in algorithmic complexity.
  • The method has minimal impact on generic text generation, addressing a central challenge of precise fine-tuning.
Loading 2402.12865v1…