How to Dissect a M uppet: The Structure of Transformer Embedding Spaces (2022)

Abstract: Pretrained embeddings based on the Transformer architecture have taken the NLP community by storm. We show that they can mathematically be reframed as a sum of vector factors and showcase how to use this reframing to study the impact of each component. We provide evidence that multi-head attentions and feed-forwards are not equally useful in all downstream applications, as well as a quantitative overview of the effects of finetuning on the overall embedding space. This approach allows us to draw connections to a wide range of previous studies, from vector space anisotropy to attention weights.

## introduction the transformer architecture @xcite has taken the nlp community by storm. based on the attention mechanism @xcite , it was shown to outperform recurrent architectures on a wide variety of tasks. another step was taken with pretrained language models derived from this architecture (bert, @xcite : they now embody the default approach to a vast swath of nlp applications. success breeds scrutiny; likewise the popularity of these models has fostered research in explainable nlp interested in the behavior and explainability of pretrained language models @xcite . in this paper, we develop a novel decomposition of transformer output embeddings. our approach consists in quantifying the contribution of each network submodule to the output contextual embedding, and grouping those into four terms: (i) what relates to the input for a given position, (ii) what pertains to feed-forward submodules, @xcite what corresponds to multi-head attention, and @xcite what is due to vector biases. this allows us to investigate transformer embeddings without relying on attention weights or treating the entire model as a black box, as is most often done in the literature. the usefulness of our method is demonstrated on bert: our case study yields enlightening connections to state-of-the-art work on transformer explainability, evidence that multi-head attentions and feed-forwards are not equally useful in all downstream applications, as well as an overview of the effects of finetuning on the embedding space. we also provide a simple and intuitive measurement of the importance of any term in this decomposition with respect to the whole embedding. we will provide insights on the transformer architecture in section 2, and showcase how these insights can translate into experimental investigations in sections 3 to 6. we will conclude with connections to other relevant works in section 7 and discuss future perspectives in section 8. ## additive structure in transformers we show that the transformer embedding e t for a token t is as a sum of four terms: ## the mlm objective an interesting follow-up question concerns which of the four terms allow us to retrieve the target word-piece. we consider two approaches: (a) using the actual projection learned by the non-fine-tuned bert model, or (b) learning a simple categorical regression for a specific term. we randomly select 15% of the word-pieces in our europarl sample. as in the work of @xcite , 80% of these items are masked, 10% are replaced by a random word-piece, and 10% are left as is. selected embeddings are then split between train (80%), validation (10%), and test (10%). results are displayed in the default projection matrix is also highly dependent on the normalization offsets c t and the ff terms f t being added together: removing this c t term from any experiment using f t is highly detrimental to the accuracy. on the other hand, combining the two produces the highest accuracy scores. our logistic regressions show that most of this performance can be imputed to the f t term. learning a projection from the f t term already yields an accuracy of almost 54%. on the other hand, a regression learned from c t only has a limited performance of 9.72% on 9 we thank an anonymous reviewer for pointing out that the bert model ties input and output embeddings; we leave investigating the implications of this fact for future work. in all, do ffs make the model more effective? the f t term is necessary to achieve the highest accuracy on the training objective of bert. on its own, it doesn't achieve the highest performances: for that we also need to add the mha outputs h t . however, the performances we can associate to f t on its own are higher than what we observe for h t , suggesting that ffs make the transformer architecture more effective on the mlm objective. this result connects with the work of @xcite who argue that ffs update the distribution over the full vocabulary, hence it makes sense that f t would be most useful to the mlm task. ## lexical contents and wsd we now turn to look at how the vector spaces are organized, and which term yields the most linguistically appropriate space. we rely on wsd, as distinct senses should yield different representations. we consider an intrinsic knn-based setup and an extrinsic probe-based setup. the former is inspired from @xcite : we assign to a target the most common label in its neighborhood. we restrict neighborhoods to words with the same annotated lemma and use the 5 nearest overall, the same picture emerges from the knn setup and all 5 runs of the classifier setup. the f t term does not yield the highest performances in our experiment-instead, the h t term systematically dominates. in single term models, h t is ranked first and f t second. as for sums of two terms, the setups ranked 1 st , 2 nd , and 3 rd are those that include h t ; setups ranked 3 rd to 5 th , those that include f t . even more surprisingly, when summing three of the terms, the highest ranked setup is the one where we exclude f t , and the lowest corresponds to excluding h t . removing f t systematically yields better performances than using the full embedding. this suggests that f t is not necessarily helpful to the final representation for wsd. this contrast with what we observed for mlm, where h t was found to be less useful then f t . one argument that could be made here would be to posit that the predictions derived from the different sums of terms are intrinsically different, hence a purely quantitative ranking might not capture this important distinction. to verify whether this holds, we can look at the proportion of predictions that agree for any two models. because our intent is to see what can be retrieved from specific subterms of the embedding, we focus solely on the most efficient classifiers across runs. this is summarized in figure 4 : an individual cell will detail the proportion of the assigned labels shared by the models for that row and that column. in short, we see that model predictions tend to a high degree of overlap. for both knn and classifier setups, the three models that appear to make the most distinct predictions turn out to be computed from the i t term, the c t term or their sum: that is, the models that struggle to perform better than the mfs baseline and are derived from static representations. we consider the task of named entity recognition, using the wnut 2016 shared task dataset @xcite . we contrast the performance of the non-finetuned bert model to that of the aforementioned variant fine-tuned on the conll 2003 ner dataset using shallow probes. we can also look at whether the highest scoring classifiers across runs classifiers produce different outputs. given the high class imbalance of the dataset at hand, we macro-average the prediction overlaps by label. the result is shown in figure 5 ; the upper triangle details the behavior of the untuned model, and the lower triangle details that of the ner-fine-tuned model. in this round of experiments, we see much more distinctly that the i t model, the c t model, and the i t + c t model behave markedly different from the rest, with c t yielding the most distinct predictions. as for the ner-fine-tuned model (lower triangle), aside from the aforementioned static representations, most predictions display a degree of overlap much higher than what we observe for the non-finetuned model: both ffs and mhas are skewed towards producing outputs more adapted to ner tasks. ## relevant work the derivation we provide in section 2 ties in well with other studies setting out to explain how transformers embedding spaces are structured @xcite @xcite and more broadly how they behave @xcite . for instance, lower layers tend to yield higher performance on surface tasks (e.g., predicting the presence of a word, @xcite or static benchmarks (e.g., analogy, @xcite ): this ties in with the vanishing prominence of i t across layers. likewise, probe-based approaches to unearth a linear structure matching with the syntactic structure of the input sentence @xcite can be construed as relying on the explicit linear dependence that we highlight here. another connection is with studies on embedding space anisotropy @xcite : our derivation provides a means of circumscribing which neural components are likely to cause it. also relevant is the study on sparsifying transformer representations of @xcite : the linearly dependent nature of transformer embeddings has some implications when it comes to dictionary coding. also relevant are the works focusing on the interpretation of specific transformer components, and feed-forward sublayers in particular @xcite @xcite . lastly, our approach provides some quantitative argument for the validity of attention-based studies @xcite @xcite and expands on earlier works looking beyond attention weights @xcite . ## conclusions and future work in this paper, we stress how transformer embeddings can be decomposed linearly to describe the impact of each network component. we showcased how this additive structure can be used to investigate transformers. our approach suggests a less central place for attention-based studies: if multi-head attention only accounts for 30% of embeddings, can we possibly explain what transformers do by looking solely at these submodules? the crux of our methodology lies in that we decompose the output embedding by submodule instead of layer or head. these approaches are not mutually exclusive (cf. section 3), hence our approach can easily be combined with other probing protocols, providing the means to narrow in on specific network components. the experiments we have conducted in sections 3 to 6 were designed so as to showcase whether our decomposition in equation ( 1 ) could yield useful results-or, as we put it earlier in section 2.3, whether this approach could be conducive to future research. we were able to use the proposed approach to draw insightful connections. the noticeable anisotropy of contextual embeddings can be connected to the prominent trace of the biases in the output embedding: as model biases make up an important part of the whole embedding, they push it towards a specific sub-region of the embedding. the diminishing importance of i t links back to earlier results on word-type semantic benchmarks. we also report novel findings, showcasing how some submodules outputs may be detrimental in specific scenarios: the output trace of ff modules was found to be extremely useful for mlm, whereas the h t term was found to be crucial for wsd. our methodology also allows for an overview of the impact of finetuning (cf. section 6): it skews components towards more task-specific outputs, and its effect are especially noticeable in upper layers (figures 3d and 3e ). analyses in sections 3 to 6 demonstrate the immediate insight that our transformer decomposition can help achieve. this work therefore opens a number of research perspectives, of which we name three. first, as mentioned in section 2.3, our approach can be extended further to more thoroughly disentangle computations. second, while we focused here more on feed-forward and multi-head attention components, extracting the static component embeddings from i t would allow for a principled comparison of contextual and static distributional semantics models. last but not least, because our analysis highlights the different relative importance of transformer components in different tasks, it can be used to help choose the most appropriate tools for further interpretation of trained models among the wealth of alternatives.