PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space

Mor Geva, Avi Caciularu, Kevin Ro Wang, Yoav Goldberg

arXiv:2203.14680v3cs.CL

TL;DR

The paper addresses limited understanding of how transformer LMs construct predictions by reverse-engineering FFN-layer updates in vocabulary space. It decomposes these updates into value-vector sub-updates, finding interpretable concept promotion and practical gains in toxicity control and inference efficiency.

  • Problem

    The internal process by which transformer-based language models construct predictions remains insufficiently understood, particularly for FFN layers.

  • Method

    The paper views token representations as evolving vocabulary distributions and decomposes each FFN output into weighted sub-updates associated with individual value vectors.

  • Results

    The value-vector sub-updates often encode human-interpretable concepts, promote tokens in the output distribution, reduce GPT2 toxicity by almost 50%, and save 20% of computation on average.

  • Takeaways & Limitations

    The analysis provides a more detailed account of internal prediction construction and supports research directions in interpretability, prediction control, and efficiency.

  • Takeaways & Limitations

    The analysis focuses on individual value vectors and decoder-based autoregressive LMs, so its applicability to compositional value-vector meanings and other transformer objectives remains to be established.

Abstract

from arXiv · show

Transformer-based language models (LMs) are at the core of modern NLP, but their internal prediction construction process is opaque and largely not understood. In this work, we make a substantial step towards unveiling this underlying prediction process, by reverse-engineering the operation of the feed-forward network (FFN) layers, one of the building blocks of transformer models. We view the token representation as a changing distribution over the vocabulary, and the output from each FFN layer as an additive update to that distribution. Then, we analyze the FFN updates in the vocabulary space, showing that each update can be decomposed to sub-updates corresponding to single FFN parameter vectors, each promoting concepts that are often human-interpretable. We then leverage these findings for controlling LM predictions, where we reduce the toxicity of GPT2 by almost 50%, and for improving computation efficiency with a simple early exit rule, saving 20% of computation on average.

1 Introduction

The paper reverse-engineers how FFN layers contribute to transformer LM predictions by interpreting token representations as vocabulary distributions and decomposing FFN updates into concept-related sub-updates. These findings support both prediction control and computational savings.

  • Core approach: FFN layers apply additive updates to token representations that can be interpreted as changes to distributions over the output vocabulary.The representation and its updates are analyzed in vocabulary space.
  • Core approach: Each FFN update decomposes into sub-updates corresponding to columns of the second FFN matrix, which often encode human-interpretable concepts.Examples include “breakfast” and “pronouns.”
  • Prediction construction: FFN updates primarily promote tokens rather than eliminate them, so highly ranked output tokens are those strongly pushed by sub-updates.
  • Practical utility: Increasing the weight of only 10 GPT2 sub-updates reduces toxicity in generations by almost 50%.
  • Practical utility: Dominant sub-updates provide a signal for predicting early exits, saving 20% of computation on average.

2 Token Representations as Evolving Distributions Over the Vocabulary

The model’s contextualized token representation evolves across transformer layers and can be translated at each stage into a vocabulary distribution. FFN outputs therefore act as additive changes to that evolving distribution, motivating their decomposition into smaller interpretable updates.

  • Vocabulary distributions: Transformer LMs obtain predictions by projecting the final-layer representation to an embedding matrix and applying softmax over the vocabulary.
  • Evolving representations: Each input token receives a contextualized representation that is updated throughout the transformer layers.The analysis focuses on the FFN updates applied to a single token.
  • Evolving representations: The evolving representation can be viewed as an information stream, with the final representation determining the output probability distribution.
  • Vocabulary distributions: Applying the final prediction projection to an intermediate representation yields a distribution over the output vocabulary at that layer.
  • FFN updates: By linearity, the projected FFN output forms an additive update in vocabulary space, although the undivided projection is not itself interpretable.The paper addresses this by decomposing the update into smaller sub-updates whose vocabulary projections often express human-interpretable concepts.

3 The FFN Output as a Collection of Updates to the Output Distribution

An FFN output is represented as a collection of weighted value-vector sub-updates. In vocabulary space, each value vector has an input-independent token ranking and an input-dependent coefficient that determines its contribution to a particular update.

  • FFN decomposition: FFN layers are modeled as two linear transformations with a point-wise activation, using parameter matrices interpreted as keys and values.The key transformation produces coefficients that weight the value vectors.
  • FFN decomposition: Each FFN update is a collection of sub-updates, with each sub-update corresponding to a weighted value vector in the FFN output.
  • FFN decomposition: A transformer with L = 10 and dm = 3000 contains 30,000 value vectors, but only a few sub-updates typically receive high weights for each token.
  • Vocabulary-space effects: A value-vector sub-update changes token probabilities according to each token embedding’s dot product with that sub-update.Positive dot products increase a token’s probability, while negative dot products decrease it.
  • Vocabulary-space effects: The projection of a value vector produces a static vocabulary score that ranks tokens independently of the model input.
  • Vocabulary-space effects: The coefficients assigned to value vectors are fixed for a given input and compare their contributions within that update.
  • Research questions: The analysis asks what information sub-updates encode and how FFN updates construct the output probability distribution.

4 Sub-Updates Encode Concepts in the Vocabulary Space

The paper projects FFN sub-updates into vocabulary space and finds that their top-scoring tokens often reveal human-interpretable concepts. Decomposing updates into sub-updates provides a more fine-grained interface for interpreting FFN outputs.

  • Concept annotation: Experts annotate recurring patterns among each value vector’s top-30 tokens as semantic, syntactic, or name concepts.Patterns must occur in at least four tokens before being described and classified.
  • Real versus random vectors: 55.1% vs. 22.7% in WIKILM and 37% vs. 16% in GPT2 of top-tokens were associated with concepts for real versus random value vectors.Average concepts per vector exceeded 1 for real vectors but was approximately 0.5 for random vectors.
  • Updates versus sub-updates: 19.7% vs. 55.1% in WIKILM and 11.8% vs. 36.7% in GPT2 of tokens were covered by whole FFN updates versus sub-updates after excluding stopwords and punctuation.The comparison supports decomposing FFN outputs for fine-grained interpretation.
  • Concept prevalence: 40%-70% of WIKILM and 20%-65% of GPT2 top-tokens across layers were associated with well-defined concepts.Most annotations were semantic; each value vector expressed 1.5 concepts on average in WIKILM and 1.1 in GPT2.
  • Interpretation: Across all layers, value-vector projections represented general concepts rather than prioritizing only specific next tokens.This extends earlier observations about patterned next-token distributions in upper-layer value vectors.
  • Limitation: Human annotation likely underestimates concept frequency because tokens lack context and may require world knowledge or linguistic background.Additional context and token-related information are proposed as future improvements to the annotation protocol.

5 FFN Updates Promote Tokens in the Output Distribution

The paper tests how FFN sub-updates alter the vocabulary distribution by examining dominant contributions to promoted and eliminated candidates. Across layers, the evidence supports token promotion as the primary mechanism, with specialized extreme updates in easy predictions.

  • Mechanism: FFN updates can increase, decrease, or leave unchanged a token’s probability according to the vocabulary projection of each sub-update.The analysis focuses on the ten sub-updates with the largest representation contributions.
  • Promoted versus eliminated candidates: Tokens promoted to the top receive higher maximum scores than eliminated tokens: 1.2 → 0.5 in WIKILM and 8.5 → 4.0 in GPT2.Promoted tokens are pushed strongly by a few dominant sub-updates.
  • Promoted versus eliminated candidates: Eliminated top candidates receive near-zero mean scores from both dominant and random sub-updates, suggesting they are not directly eliminated.For promoted tokens, maximum scores substantially exceed minimum scores: 1.2 vs. −0.8 in WIKILM and 8.5 vs. −4.9 in GPT2.
  • Conclusion: The findings support a promotion mechanism in which top-candidate tokens are pushed by dominant sub-updates.The random-sub-update baseline has dramatically lower score magnitudes, supporting the relevance of the selected dominant sub-updates.
  • Across layers: Before the last few layers, maximum and minimum scores generally surround non-negative means, with prominent maximum-score peaks in layers 3-5 of GPT2 and 4-11 of WIKILM.Strong negative minimum scores emerge in GPT2 layers 23-24 and WIKILM layers 14-16.
  • Extreme sub-updates: Only 1.7% of GPT2 and 1.1% of WIKILM value vectors form homogeneous clusters responsible for extreme sub-updates.These clusters promote rare tokens in upper layers or common tokens such as stopwords across layers, often for inputs of three or fewer tokens or easily inferred targets.

6 Applications

The paper applies FFN sub-update analysis to toxic-language suppression and self-supervised early exiting, showing practical control and efficiency gains.

  • 6.1 Zero-Shot Toxic Language Suppression: The suppression method identifies value vectors promoting safe concepts, then increases their coefficients to 3.Vectors were selected through manual inspection or Perspective API toxicity grading.
  • 6.1 Zero-Shot Toxic Language Suppression: Activating only 10 GPT2 value vectors reduced toxic completions by 47%, outperforming Self-Debiasing at 37% and WORDFILTER at 20%.The evaluation used the challenging REALTOXICPROMPTS subset and six toxicity attributes.
  • 6.1 Zero-Shot Toxic Language Suppression: The method increased perplexity more than Self-Debiasing, although the increase remained relatively small.
  • 6.2 Self-Supervised Early Exit Prediction: The early-exit method uses dominant FFN sub-updates to predict saturation without training an external model.It was evaluated on WIKILM using nearest-neighbour matching over clusters of dominant sub-updates.
  • 6.2 Self-Supervised Early Exit Prediction: 94.1% accuracy and 20% average computation savings were achieved without changing the final model prediction.Accuracy measures correct exits yielding the final model prediction, while efficiency measures saved layers for those examples.

7 Related Work

Prior work studied transformer representations, probed models with target tasks, and analyzed FFN knowledge; this paper instead interprets FFN outputs and their role in prediction construction.

  • Earlier interpretability work mainly analyzed hidden-representation evolution or probed models using target tasks.
  • Recent studies showed that FFN layers store knowledge in specific neurons that can be located and edited.
  • This paper differs by focusing on FFN outputs and their contribution to constructing predictions rather than only stored knowledge.
  • The vocabulary-space view relates to work interpreting parameter groups in discrete vocabulary space and representations as information streams.

8 Conclusions

The paper analyzes FFN outputs as value-vector updates that promote interpretable vocabulary concepts, providing a finer-grained view of prediction construction and practical directions for control and efficiency.

  • FFN outputs are modeled as linear combinations of parameter vectors called values, whose updates modify token representations.
  • Value vectors often encode human-interpretable concepts that are promoted in the output distribution.
  • The analysis provides a more detailed understanding of internal prediction construction and suggests research directions at the level of individual vectors.

Limitations

The study's conclusions are bounded by its focus on FFN layers, individual value vectors, decoder-based autoregressive LMs, and a manually inspected concept-annotation process.

  • The study examines FFN operation without fully analyzing its interplay with other transformer components such as attention heads.
  • Individual value vectors may be compositional, so combinations could produce meanings not captured by single-vector analysis.The authors retain the individual-vector focus because combinations are exponentially numerous and less interpretable in their analysis.
  • The findings were developed for decoder-based autoregressive LMs, so their validity for encoder-only models and different training objectives remains to be established.
  • Manual concept interpretation may depend on tokenization, and scaling the annotation protocol could provide a more refined concept map.

Ethics Statement

The proposed controls can reduce toxicity and computation, but the authors warn that neither should be deployed in real-world critical systems without further work, caution, and monitoring.

  • The toxic-language suppression method reduces the probability of toxic generation rather than eliminating it.
  • Changing FFN sub-update weights can steer predictions in particular directions, but the approach may also be abused.
  • The authors caution against integrating these methods into critical systems without further work, caution, and monitoring.

A Appendix

The appendix details how FFN sub-updates are interpreted and weighted, how their dominant components support prediction control and early exit, and where annotation and deployment limitations remain.

  • Normalization analysis: 64.5% of the top-30 tokens overlap on average between normalized and unnormalized projections, compared with approximately 0% for random vectors.This suggests that layer normalization generally preserves the concepts identified through value-vector projection.
  • Concept identification: Sub-updates are interpreted by projecting value vectors into vocabulary space and identifying repeated semantic or syntactic patterns among the top-scoring 30 tokens.Patterns were annotated by NLP graduate students using task instructions, with additional lexical descriptions for uncommon WIKILM tokens.
  • Sub-update contribution: The contribution of a sub-update is its absolute coefficient relative to the sum of absolute coefficients across all sub-updates.The top-10 dominant sub-updates are evaluated by summing these relative contributions; absolute values accommodate negative coefficients produced by activations such as GeLU.
  • Sub-update contribution: Although dominant sub-updates cover only 0.24% of value vectors, they typically contribute around 5% of the FFN output and exceed random sub-updates across layers.Negative coefficients can also be among GPT2’s dominant sub-updates, allowing reversal of the scores induced by value vectors.
  • Early exit: 34.15% of GPT2 examples require full computation, compared with 15.22% for WIKILM, motivating early-exit experiments on WIKILM.The appendix attributes WIKILM’s higher saving potential partly to more saturation events per layer.
Loading 2203.14680v3…