Language Models Implement Simple W ord2 V ec-style Vector Arithmetic (2024)

Abstract: A primary criticism towards language models (LMs) is their inscrutability. This paper presents evidence that, despite their size and complexity, LMs sometimes exploit a simple vector arithmetic style mechanism to solve some relational tasks using regularities encoded in the hidden space of the model (e.g., Poland:Warsaw::China:Beijing). We investigate a range of language model sizes (from 124M parameters to 176B parameters) in an in-context learning setting, and find that for a variety of tasks (involving capital cities, uppercasing, and past-tensing) a key part of the mechanism reduces to a simple additive update typically applied by the feedforward (FFN) networks. We further show that this mechanism is specific to tasks that require retrieval from pretraining memory, rather than retrieval from local context. Our results contribute to a growing body of work on the interpretability of LMs, and offer reason to be optimistic that, despite the massive and non-linear nature of the models, the strategies they ultimately use to solve tasks can sometimes reduce to familiar and even intuitive algorithms.

## introduction the growing capabilities of large language models (llms) have led to an equally growing interest in understanding how such models work under the hood. such understanding is critical for ensuring that llms are reliable and trustworthy once deployed. recent work in interpretability has contributed to this understanding by reverseengineering the data structures and algorithms that are implicitly encoded in the model's weights, e.g., by identifying detailed circuits @xcite @xcite or by identifying mechanisms for factual storage and retrieval which support intervention and editing @xcite @xcite . here, we contribute to this growing body of work by analyzing how llms recall information during in-context learning. modern llms are based on a complex transformer architecture @xcite which produces contextualized word embeddings @xcite connected via multiple non-linearities. despite this, we find that llms implement a basic vector-addition mechanism qualitatively similar to relational information encoded in their static word embeddings predecessors @xcite . we also find that for non-injective relations that static embeddings typically fail to encode @xcite , lms do not use the identified mechanism (appendix g). we study this phenomenon across nine tasks, but focus on three in the main paper: recalling capital cities, uppercasing tokens, and past-tensing verbs. our key findings are: taken together, our results offer new insights about one component of the complex algorithms that underlie in-context learning. the mechanism's simplicity raises the possibility that other apparently complicated behaviors may be supported by a sequence of simple operations under the hood. moreover, our results suggest a distinct processing signature and hint at a method for intervention. these ideas could support future work on detecting and preventing unwanted behavior by llms at runtime. ## methods in decoder-only transformer language models @xcite , a sentence is processed one word at a time, from left to right. in this paper, we focus on the transformations that the next-token prediction undergoes in order to predict the answer to some task. at each layer, an attention module and feed-forward network (ffn) module apply subsequent additive updates to this representation. consider the ffn update at layer i, where x i is the current next-token representation. the update applied by the ffn here is calculated as ## stages of processing in predicting the next token first, we use the early decoding method in order to investigate how the processing proceeds over the course of a forward pass to the model. each task requires the model to infer some relation to recall some fact, e.g., retrieving the capital of poland. in these experiments, we see several discrete stages of processing that the next token undergoes before in figure 2 we illustrate an example of the stages we observe across models. for the first several layers, we see no movement on the words of interest. then, during argument formation, the model first represents the argument to the desired relation in the residual stream. this means that the top token in the vocabulary distribution at some intermediate layer(s) is the subject the question inquires about (e.g., the x, in get_capital(x)). during function application we find that the model abruptly switches from the argument to the output of the function (the y, in get_capital(x) = y). we find that function application is typically applied by the ffn update at that layer to the residual stream. this is done by adding the output vector ⃗ o of the ffn to the residual stream representation, thus transforming it with an additive update. we study these ⃗ o vectors in detail in section 4. finally, the model enters saturation 3 , where the model recognizes it has solved the next token, and ceases updating the token representation for the remaining layers. the trend can be characterized by an x-shaped pattern of the argument and final output tokens when plotting the ranks of the argument(x) and output (y) tokens. we refer to this behavior as argument-function processing. figure 3 shows that this same processing signature can be observed consistently across tasks and models. moreover, it appears to become more prominent as the models increase in size. interestingly, despite large differences in number of layers and overall size, models tend to undergo this process at similar points proportionally in the model. ## implementation of context-independent functions in ffn updates the above results on processing signature suggest that the models "apply" a function about 2/3rds of the way through the network with the addition of an 3 saturation events are described in geva et al. (2022a) where detection of such events is used to "early-exit" out of the forward pass ffn update. here, we investigate the mechanism via which that function is applied more closely. specifically, focusing on gpt2-medium foot_2 , we show that we can force the encoded function to be applied to new arguments in new contexts by isolating the responsible ffn output vector and then dropping into a forward pass on a new input. ## the role of ffns in out-of-context retrieval so far, we have shown that ffn output vectors can encode functions that transfer across contexts. here, we investigate the role of this mechanism when we control whether the answer occurs in context. the tasks we study previously require recalling a token that does not appear in the given context (abstractive tasks). in this section we show that mid-higher layer ffns are crucial for this process. when the answer to the question does appear in context (extractive tasks), we find that ablating a subset of ffns has a comparatively minor effect on performance, indicating that they are relatively modular and there is a learned division of labor within the model. this observation holds across the decoder-only lms tested in this paper. this breakdown is consistent with previous work finding that ffns store facts learned from pre-training @xcite and attention heads copy from the previous context @xcite . ## related work attributing roles to components in pretrained lms is a widely studied topic. in particular, the attention layers @xcite @xcite and in the ffn modules, which are frequently associated with factual recall and knowledge storage @xcite . how language models store and use knowledge has been studied more generally as well @xcite @xcite @xcite @xcite as well as in static embeddings @xcite . recent work in mechanistic interpretability aims to fully reverse engineer how lms perform some behaviors @xcite . our work builds on the finding that ffn layers promote concepts in the vocabulary space @xcite by breaking down the process the model uses to do this in context; @xcite perform ablation studies to test the importance of attention and ffn layers on in-context learning tasks. other work analyze information flow within an lm to study how representations are built through the layers, finding discrete processing stages @xcite . we also follow this approach, but our analysis focuses on interpreting how models use individual updates within the forward pass, rather than probing for information encoded within some representation. @xcite show that vector arithmetic can be performed with the weights of finetuned models to compose tasks, similar to how ⃗ o vectors can induce functions in the activation space of the model. ## discussion & conclusion a core challenge in interpreting neural networks is determining whether the information attributed to certain model components is actually used for that purpose during inference @xcite . while previous work has implicated ffns in recalling factual associations @xcite , we show through intervention experiments that we can manipulate the information flowing through the model according to these stages. this process provides a simple explanation for the internal subprocesses used by lms and our findings invite future work aimed at understanding why, and under what conditions, lms learn to use this mechanism when they are capable of solving such tasks using, e.g., adhoc memorization. the mechanism we identify bears similarities to linguistic regularities that allow for vector arithmetic analogies in static word embeddings @xcite suggesting at least a qualitative similarity between large complex contextual models and these simpler static models. @xcite show that not all relations can be encoded with vector arithmetic analogies, specifically, relations that are not one-to-one (e.g., mapping a country to its official language). in appendix g we find evidence that lms exhibit similar success and failure cases by analyzing six additional tasks. we provide our most detailed investigation on gpt2-medium, which clearly illustrates the phenomenon. our experiments on stages of processing with gpt-j suggest that the same phenomena is in play, although (as discussed in section 4 and appendix a), the procedures we derive for interventions on gpt2-medium do not transfer perfectly. specifically, we can strongly reproduce the intervention results on uppercasing for gpt-j; results on the other two tasks are positive but with overall weaker effects. as we understand these processes more deeply, a priority in future work must be to generalize specific findings to model-agnostic phenomena. that said, in this work and other similar efforts, a single positive example as a proof of concept is often sufficient to advance understanding and spur future work that improves robustness across models. contemporaneous work @xcite has studied a different mechanism for factual recall in lms, but it is unclear how and when these mechanisms interact. eventually, if we can understand how models break down complex problems into simple and predictable subprocesses, we can help more readily audit their behavior. interpreting the processing signatures of model behaviors might offer an avenue via which to evaluate and intervene at runtime in order to prevent unwanted behavior.