Analyzing Transformers in Embedding Space (2023)

Abstract: Understanding Transformer-based models has attracted significant attention, as they lie at the heart of recent technological advances across machine learning. While most interpretability methods rely on running models over inputs, recent work has shown that a zero-pass approach, where parameters are interpreted directly without a forward/backward pass is feasible for some Transformer parameters, and for two-layer attention networks. In this work, we present a theoretical analysis where all parameters of a trained Transformer are interpreted by projecting them into the embedding space, that is, the space of vocabulary items they operate on. We derive a simple theoretical framework to support our arguments and provide ample evidence for its validity. First, an empirical analysis showing that parameters of both pretrained and fine-tuned models can be interpreted in embedding space. Second, we present two applications of our framework: (a) aligning the parameters of different models that share a vocabulary, and (b) constructing a classifier without training by “translating” the parameters of a fine-tuned classifier to parameters of a different model that was only pretrained. Overall, our findings open the door to interpretation methods that, at least in part, abstract away from model specifics and operate in the embedding space only.

## introduction transformer-based models @xcite currently dominate natural language processing @xcite @xcite as well as many other fields of machine learning @xcite @xcite . consequently, understanding their inner workings has been a topic of great interest. typically, work on interpreting transformers relies on feeding inputs to the model and analyzing the resulting activations @xcite @xcite . thus, interpretation involves an expensive forward, and sometimes also a backward pass, over multiple inputs. moreover, such interpretation methods are conditioned on the input and are not guaranteed to generalize to all inputs. in the evolving literature on static interpretation, i.e., without forward or backward passes, @xcite showed that the value vectors of the transformer feed-forward module (the second layer of the feed-forward network) can be interpreted by projecting them into the embedding space, i.e., multiplying them by the embedding matrix to obtain a representation over vocabulary items. foot_0 @xcite have shown that in a 2-layer attention network, weight matrices can be interpreted in the embedding space as well. unfortunately, their innovative technique could not be extended any further. in this work, we extend and unify the theory and findings of @xcite and @xcite . we present a zero-pass, input-independent framework to understand the behavior of transformers. concretely, we interpret all weights of a pretrained language model (lm) in embedding space, including both keys and values of the feed-forward module ( @xcite considered just ff values) as well as all attention parameters ( @xcite analyzed simplified architectures up to two layers of attention with no mlps). our framework relies on a simple observation. since @xcite have shown that one can project hidden states to the embedding space via the embedding matrix, we intuit this can be extended to other parts of the model by projecting to the embedding space and then projecting back by multiplying with a right-inverse of the embedding matrix. thus, we can recast inner products in the model as inner products in embedding space. viewing inner products this way, we can interpret such products as interactions between pairs of vocabulary items. this applies to (a) interactions between attention queries and keys as well as to (b) interactions between attention value vectors and the parameters that project them at the output of the attention module. taking this perspective to the extreme, one can view transformers as operating implicitly in the embedding space. this entails the existence of a single linear space that depends only on the tokenizer, ea eb layer 18 head 1 ('women', ' marie') (' actresses', ' marie') ('women', ' anne') ('women', ' anne') ('woman', ' marie') ('women', ' marie') w i model b l1 l2 l3 l4 l5 l1 l2 l3 l4 l5 vo e e w qk e e t t e v e k embedding space compute dot-product in embedding space between all ff key pairs* from all layers k j k parameter interpretation parameter alignment zero-shot stitching i j model a pack them into a similarity matrix average absolute similarity over layer pairs l1 l2 l3 l4 l5 matching layers top k top k top k layer 10 dim 8 tai #jin jin makoto #etsu shin hai top k layer 21 dim 400 #july july #february #january #feb november #october layer 17 head 6 (' legally', ' legal') (' legal', ' sentencing') (' legal', ' arbitration') (' boycot', ' boycott') (' legal', ' criminal') (' legal', ' judicial') (' legal', ' rulings') in which parameters of different transformers can be compared. thus, one can use the embedding space to compare and transfer information across different models that share a tokenizer. we provide extensive empirical evidence for the validity of our framework, focusing mainly on gpt-2 medium @xcite . we use gpt-2 for two reasons. first, we do this for concreteness, as this paper is mainly focused on introducing the new framework and not on analyzing its predictions. second, and more crucially, unlike many other architectures (such as bert @xcite , roberta @xcite , and t5 @xcite ), the gpt family has a linear language modeling head (lm head) -which is simply the output embedding matrix. all the other architectures' lm heads are two layer networks that contain non-linearities before the output embedding matrix. our framework requires a linear language modeling head to work. that being said, we believe in practice this will not be a major obstacle, and we indeed see in the experiments that model alignment works well for bert in spite of the theoretical difficulties. we leave the non-linearities in the lm head for future work. on the interpretation front (fig. 1 , left), we provide qualitative and quantitative evidence that transformer parameters can be interpreted in embedding space. we also show that when fine-tuning gpt-2 on a sentiment analysis task (over movie reviews), projecting changes in parameters into embedding space yields words that characterize sentiment towards movies. second (fig. 1 , center), we show that given two distinct instances of bert pretrained from different random seeds @xcite , we can align layers of the two instances by casting their weights into the embedding space. we find that indeed layer i of the first instance aligns well to layer i of the second instance, showing the different bert instances converge to a semantically similar solution. last (fig. 1 , right), we take a model fine-tuned on a sentiment analysis task and "transfer" the learned weights to a different model that was only pretrained by going through the embedding spaces of the two models. we show that in 30% of the cases, this procedure, termed stitching, results in a classifier that reaches an impressive accuracy of 70% on the imdb benchmark @xcite without any training. overall, our findings suggest that analyzing transformers in embedding space is valuable both as an interpretability tool and as a way to relate different models that share a vocabulary and that it opens the door to interpretation methods that operate in embedding space only. our code is available at https://github.com/guyd1995/embedding-space . ## background we now present the main components of the transformer @xcite relevant to our analysis. we discuss the residual stream view of transformers, and recapitulate a view of the attention layer parameters as interaction matrices w vo and w qk @xcite . similar to them, we exclude biases and layer normalization from our analysis. ## parameter projection in this section, we propose that transformer parameters can be projected into embedding space for interpretation purposes. we empirically support our framework's predictions in §4- §5. given a matrix @xmath0, we can project it into embedding space by multiplying by the embedding matrix e as â = @xmath1. let e ′ be a right-inverse of e, that is, ee ′ = @xmath2. 3 we can reconstruct the original matrix with e ′ as @xmath3. we will use this simple identity to reinterpret the model's operation in embedding space. to simplify our analysis we ignore layernorm and biases. this has been justified in prior work @xcite . briefly, layernorm can be ignored because normalization changes only magnitudes and not the direction of the update. at the end of this section, we discuss why in practice we choose to use e ′ = e t instead of a seemingly more appropriate right inverse, such as the pseudo-inverse @xcite @xcite . in this section, we derive our framework and summarize its predictions in since the output of the attention module is added to the residual stream, we can assume according to the residual stream view that it is meaningful to project it to the embedding space, similar to ff values. thus, we expect the sequence of n e-dimensional vectors similarly, the matrix w i qk can be viewed as a bilinear map ( eq. 2.3 ). to interpret it in embedding space, we perform the following operation with e ′ : therefore, the interaction between tokens at different positions is determined by an e×e matrix that expresses ff module @xcite showed that ff value vectors @xmath4, i.e., for a ff value vector @xmath5, @xmath6.1). in vectorized form, the rows of v @xmath7. on the other hand, the keys k of the ff layer are multiplied on the left by the output of the attention module, which are the queries of the ff layer. denoting the output of the attention module by q, we can write this product as qk @xmath8. because q is a hidden state, we assume according to the residual stream view that q is interpretable in embedding space. when multiplying q by ke ′t , we are capturing the interaction in embedding space between each query and key, and thus expect ke ′t to be interpretable in embedding space as well. overall, ff keys and values are intimately connected -the i-th key controls the coefficient of the i-th value, so we expect their interpretation to be related. while not central to this work, we empirically show that keyvalue pairs in the ff module are similar in embedding space in appendix b.1. ## interpretability experiments in this section, we provide empirical evidence for the viability of our approach as a tool for interpreting transformer parameters. for our experiments, we use huggingface transformers ( @xcite ; license: apache-2.0). ## aligning models in embedding space the assumption transformers operate in embedding space leads to an exciting possibility -we can relate different models to one another so long as they share the vocabulary and tokenizer. in §5.1, we show that we can align the layers of bert models trained with different random seeds. in §5.2, we show the embedding space can be leveraged to "stitch" the parameters of a fine-tuned model to a model that was not fine-tuned. ## related work interpreting transformers is a broad area of research that has attracted much attention in recent years. a large body of work has focused on analyzing hidden representations, mostly through probing @xcite @xcite . @xcite used statistical tools to analyze the evolution of hidden representations throughout layers. recently, @xcite proposed to decompose the hidden representations into the contributions of different transformer components. unlike these works, we interpret parameters rather than the hidden representations. another substantial effort has been to interpret specific network components. previous work analyzed single neurons @xcite , attention heads @xcite , and feedforward values @xcite @xcite . while these works mostly rely on input-dependent neuron activations, we inspect "static" model parameters, and provide a comprehensive view of all transformer components. our work is most related to efforts to interpret specific groups of transformer parameters. @xcite ] made observations about the interpretability of weights of neural networks. @xcite analyzed 2-layer attention networks. we extend their analysis to multi-layer pre-trained transformer models. @xcite interpreted feedforward values in embedding space. we coalesce these lines of work and offer a unified interpretation framework for transformers in embedding space. importantly, our framework allows us to view parameters from different models as residents of a canonical embedding space, where they can be compared in model-agnostic fashion. this has interesting implications. we demonstrate two consequences of this observation (model alignment and stitching) and argue future work can yield many more use cases. ## limitations our work has a few limitations that we care to highlight. first, it focuses on interpreting models through the vocabulary lens. while we have shown evidence for this, it does not preclude other factors from being involved. second, we used e ′ = e t , but future research may find variants of e that improve performance. additionally, most of the work focused on gpt-2. this is due to shortcomings in the current state of our framework, as well as for clear presentation. we believe nonlinearities in language modeling are resolvable, as is indicated in the experiment with bert. in terms of potential bias in the framework, some parameters might consider terms related to each due to stereotypes learned from the corpus.