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

Abstract: 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.

## introduction how do transformer-based language models (lms) construct predictions? we study this question through the lens of the feed-forward network (ffn) layers, one of the core components in transformers @xcite . recent work showed that these layers play an important role in lms, acting as memories that encode factual and linguistic knowledge @xcite @xcite . in this work, we investigate how outputs from the ffn layers are utilized internally to build predictions. we begin by making two observations with respect to the representation of a single token in the input, depicted in fig. 1 . first, each ffn layer she ordered a pancake few coffee pancake few coffee ffn layer fruit, apples, snack, vitamins, berries, oats, yogurt, tea, … (breakfast) we then decompose the ffn update ( §3), interpreting it as a collection of sub-updates, each corresponding to a column in the second ffn matrix (fig. 1 , c ) that scales the token probabilities in the output distribution. through a series of experiments, we find that (a) sub-update vectors across the entire network often encode a small-set of human-interpretable well-defined concepts, e.g. "breakfast" or "pronouns" ( §4, fig. 1 , d ), and (b) ffn updates rely primarily on token promotion (rather than elimination), namely, tokens in the top of the output distribution are those pushed strong enough by sub-updates ( §5). overall, these findings allow fine-grained interpretation of the ffn operation, providing better understanding of the prediction construction process in lms. beyond interpretation, our findings also have practical utility. in §6.1, we show how we can intervene in the prediction process, in order to manipulate the output distribution in a direction of our choice. specifically, we show that increasing the weight of only 10 sub-updates in gpt2 reduces toxicity in its generations by almost 50%. also, in §6.2, we show that dominant sub-updates provide a useful signal for predicting an early exit point, saving 20% of the computation on average. in conclusion, we investigate the mechanism in which ffn layers update the inner representations of transformer-based lms. we propose that the ffn output can be viewed as a collection of updates that promote concrete concepts in the vocabulary space, and that these concepts are often interpretable for humans. our findings shed light on the prediction construction process in modern lms, suggesting promising research directions for interpretability, control, and efficiency. ## token representations as evolving distributions over the vocabulary modern lms @xcite @xcite are transformer models primarily trained to predict the next token probability for a given input. such lms are composed of intertwined multi-head self-attention (mhsa) layers and ffn layers @xcite , with residual connections @xcite between each pair of consecutive layers. the lm prediction is obtained by projecting the output vector from the final layer to an embedding matrix @xmath0, with a hidden dimension d, to get a distribution over a vocabulary v (after softmax). given a sequence @xmath1, ..., w t ⟩ of input tokens, the model creates a contextualized representation x @xmath2, that is being updated throughout the layers. in this work, we analyze the updates applied by the ffn layers and how they construct the model prediction. concretely, each ffn layer ℓ = 1, ..., l processes x ℓ i and produces an output o ℓ i , which is then added to x ℓ i to yield an updated representation xℓ i : the updated representation xℓ i then goes through a mhsa layer, foot_0 yielding the input x ℓ+1 i for the next ffn layer. the evolving representation in this process (i.e. x ℓ @xmath3, ∀ℓ) can be viewed as an information stream that is being processed and updated by the layers @xcite . the output probability distribution is obtained from the final representation of the token, i.e., to analyze the ffn updates, we read from the representation at any layer a distribution over the output vocabulary, by applying the same projection as in eq. 1 @xcite : importantly, by linearity: in the rest of the paper, we focus on ffn updates to the representation of a single token in the sequence, and omit the token index for brevity, i.e. x ℓ := x ℓ i and p ℓ := p ℓ i . ## the ffn output as a collection of updates to the output distribution we now decompose the ffn output, and interpret it as a set of sub-updates in the vocabulary space. ## sub-updates encode concepts in the vocabulary space we evaluate whether projection to the vocabulary provides a meaningful way to "read" ffn updates, and the extent to which sub-updates are interpretable based on their projections. to this end, we manually inspect the top-scoring tokens by value vectors and check if they express interpretable concepts. concretely, we consider two representative lms (details below), and for each vector v ℓ i compute a ranking over the vocabulary by sorting the projection r ℓ i ( §3). then, we try to detect patterns in the top-scoring tokens of each value vector. concepts annotation task. we let experts (nlp graduate students) annotate concepts by identifying common patterns among the top-30 scoring tokens of each value vector. for a set of tokens, the annotation protocol includes three steps of: (a) identifying patterns that occur in at least 4 tokens, (b) describing each recognized pattern, and (c) classifying each pattern as either "semantic" (e.g., mammals), "syntactic" (e.g., past-tense verbs), or "names". the last class was added only for wikilm (see below), following the observation that a large portion of the model's vocabulary consists of names. further details, including the complete instructions and a fully annotated example can be found in app. a.2. concept sub-update top-scoring tokens gpt2 v 3 1018 measurement semantic kg, percent, spread, total, yards, pounds, hours v 8 1900 wh-relativizers syntactic which, whose, which, whom, where, who, wherein v 11 2601 food and drinks semantic drinks, coffee, tea, soda, burgers, bar, sushi wikilm v 1 1 pronouns syntactic her, she, their, her, she, they, their, they, his v 6 3025 adverbs syntactic largely, rapidly, effectively, previously, normally v 13 3516 groups of people semantic policymakers, geneticists, ancestries, ohioans ## ffn updates promote tokens in the output distribution we showed that sub-updates often encode interpretable concepts ( §4), but how do these concepts construct the output distribution? in this section, we show that sub-updates systematically configure the prediction via promotion of candidate tokens. ## applications we leverage our findings for controlled text generation ( §6.1) and computation efficiency ( §6.2). ## related work the lack of interpretability of modern lms has led to a wide interest in understanding their prediction construction process. previous works mostly focused on analyzing the evolution of hidden representations across layers @xcite , and probing the model with target tasks @xcite @xcite . in contrast, our approach aims to interpret the model parameters and their utilization in the prediction process. more recently, a surge of works have investigated the knowledge captured by the ffn layers @xcite @xcite @xcite . these works show that the ffn layers store various types of knowledge, which can be located in specific neurons and edited. unlike these works, we focus on the ffn outputs and their contribution in the prediction construction process. last, our interpretation of ffn outputs as updates to the output distribution relates to recent works that interpreted groups of lm parameters in the discrete vocabulary space @xcite , or viewed the representation as an information stream @xcite . ## conclusions understanding the inner workings of transformers is valuable for explainability to end-users, for debugging predictions, for eliminating undesirable behavior, and for understanding the strengths and limitations of nlp models. the ffn is an understudied core component of transformer-based lms, which we focus on in this work. we study the ffn output as a linear combination of parameter vectors, termed values, and the mechanism by which these vectors update the token representations. we show that value vectors often encode human-interpretable concepts and that these concepts are promoted in the output distribution. our analysis of transformer-based lms provides a more detailed understanding of their internal prediction process, and suggests new research directions for interpretability, control, and efficiency, at the level of individual vectors.