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

By Mor Geva, Avi Caciularu, Ke Wang, Kevin Wang, Yoav Goldberg. 2022

595 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Materials Chemistry, Natural Language Processing, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Materials Chemistry

In Natural Language Processing

In Model Architectures

Citations in the corpus

Language Generation Models Can Cause Harm: So What Can We Do About It? An Actionable Survey (2022, background): “Culling such parameters has been shown to reduce toxicity (Geva et al., 2022).”
Do Llamas Work in English? On the Latent Language of Multilingual Transformers (2024, methodology): “The model can use these directions as extra degrees of freedom for building rich feature representations from its raw inputs (Yosinski et al., 2014, 2015; Geva et al., 2022), which could be seen as forming an abstract “concept space”.”
Dissecting Recall of Factual Associations in Auto-Regressive Language Models (2023, methodology): “This is done through additional interventions to the MHSA and MLP sublayers and projections to the vocabulary (Dar et al., 2022; Geva et al., 2022b; Nostalgebraist, 2020).”
Eliciting Latent Predictions from Transformers with the Tuned Lens (2023, background): “But recent work has pointed out that setting activations to zero may inadvertently push a network out of distribution, making the experimental results hard to interpret (Wang et al., 2022; Chan et al., 2022).”
Finding Neurons in a Haystack: Case Studies with Sparse Probing (2023, background): “Although research on reverse-engineering specific neurons in LLMs is limited, [48] proposed—and later refined [49]—the hypothesis that feed-forward layers function as key-value memories, responding to specific input features (keys) and updating the output vocabulary distribution accordingly [50].”
Editing Large Language Models: Problems, Methods, and Opportunities (2023, unknown): “However, erasing the knowledge or information stored in LLMs (Belrose et al., 2023; Geva et al., 2022; Ishibashi and Shi-modaira, 2023) is also an important direction for investigating.”
Mechanistic Interpretability for AI Safety - A Review (2024, background): “This is possible because transformers tend to build their predictions across layers iteratively (Geva et al., 2022).”
The Linear Representation Hypothesis and the Geometry of Large Language Models (2023, background): “…interpreting (probing) (e.g., Alain & Ben-gio, 2017; Kim et al., 2018; nostalgebraist, 2020; Rogers et al., 2021; Belinkov, 2022; Li et al., 2022; Geva et al., 2022; Nanda et al., 2023) and controlling (steering) (e.g., Wang et al., 2023; Turner et al., 2023; Merullo et al., 2023; Trager et al.,…”
Explainability for Large Language Models: A Survey (2023, background): “Another study transfers the feedforward layer as a sub-update vector which is interpreted as a small set of human-interpretable concepts [Geva et al. 2022].”
Confident Adaptive Language Modeling (2022, background): “This measure tries to identify early saturation events of the hidden-state [28].”
The Reversal Curse: LLMs trained on "A is B" fail to learn "B is A" (2023, background): “When a model is updated on “ A is B ”, this gradient update may slightly alter the representation of A such that it contains information about B (e.g. in the middle MLP layers as per Geva et al. (2022; 2023)).”
Analyzing Transformers in Embedding Space (2022, methodology): “[Elhage et al., 2021] and [Geva et al., 2022b] observed that the residual stream is often barely updated in the last layers, and thus the final prediction is determined in early layers and the hidden state is mostly passed through the later layers.”