Transformer visualization via dictionary learning: contextualized embedding as a linear superposition of transformer factors (2021)

Abstract: Transformer networks have revolutionized NLP representation learning since they were introduced. Though a great effort has been made to explain the representation in transformers, it is widely recognized that our understanding is not sufficient. One important reason is that there lack enough visualization tools for detailed analysis. In this paper, we propose to use dictionary learning to open up these ‘black boxes’ as linear superpositions of transformer factors. Through visualization, we demonstrate the hierarchical semantic structures captured by the transformer factors, e.g., word-level polysemy disambiguation, sentence-level pattern formation, and long-range dependency. While some of these patterns confirm the conventional prior linguistic knowledge, the rest are relatively unexpected, which may provide new insights. We hope this visualization tool can bring further knowledge and a better understanding of how transformer networks work. The code is available at: https://github.com/zeyuyun1/TransformerVis .

## introduction though the transformer networks @xcite have achieved great success, our understanding of how they work is still fairly limited. this has triggered increasing efforts to visualize and analyze these "black boxes". besides a direct visualization of the attention weights, most of the current efforts to interpret transformer models involve "probing tasks". they are achieved by attaching a light-weighted auxiliary classifier at the output of the target transformer layer. then only the auxiliary classifier is trained for wellknown nlp tasks like part-of-speech (pos) tagging, named-entity recognition (ner) tagging, syntactic dependency, etc. @xcite and @xcite show transformer models have excellent performance in those probing tasks. these results indicate that transformer models have learned the language representation related to the probing tasks. though the probing tasks are great tools for interpreting language models, their limitation is explained in @xcite . we summarize the limitation into three major points: efforts are made to remove those limitations and make probing tasks more diverse. for instance, @xcite proposes "structural probe", which is a much more intricate probing task. @xcite proposes to generate specific probing tasks automatically. non-probing methods are also explored to relieve the last two limitations. for example, @xcite visualizes embedding from bert using umap and shows that the embeddings of the same word under different contexts are separated into different clusters. @xcite analyzes the similarity between embeddings of the same word in different contexts. both of these works show transformers provide a context-specific representation. @xcite ; @xcite ; @xcite demonstrate how to use dictionary learning to explain, improve, and visualize the uncontextualized word embedding representations. in this work, we propose to use dictionary learning to alleviate the limitations of the other transformer interpretation techniques. our results show that dictionary learning provides a powerful visualization tool, leading to some surprising new knowledge. we view the latent representation of words in a transformer as contextualized word embedding. similarly, we hypothesize that a contextualized word embedding vector can also be factorized as a sparse linear superposition of a set of elementary elements, which we call transformer factors. the exact definition will be presented later in this section. due to the skip connections in each of the transformer blocks, we hypothesize that the representation in any layer would be a superposition of the hierarchical representations in all of the lower layers. as a result, the output of a particular transformer block would be the sum of all of the modifications along the way. indeed, we verify this intuition with the experiments. based on the above observation, we propose to learn a single dictionary for the contextualized word vectors from different layers' output. to learn a dictionary of transformer factors with non-negative sparse coding. given a set of tokenized text sequences, we collect the contextualized embedding of every word using a transformer model. we define the set of all word embedding vectors from lth layer of transformer model as x (l) . furthermore, we collect the embeddings across all layers into a single set l) . by our hypothesis, we assume each embedding vector @xmath0: ## experiments and discoveries we use a 12-layer pre-trained bert model (pre; @xcite and freeze the weights. since we learn a single dictionary of transformer factors for all of the layers in the transformer, we show that these transformer factors correspond to different levels of semantic or syntactic patterns. the patterns can be roughly divided into three categories: word-level disambiguation, sentence-level pattern formation, and long-range dependency. in the following, we provide detailed visualization for each pattern category. due to the space limit, only a small amount of the factors are demonstrated in the paper. to alleviate the "cherry-picking" bias, we also build a website for the interested readers to play with these results. low-level: word-level polysemy disambiguation. while the input embedding of a token contains polysemy, we find transformer factors with early is curve peaks usually correspond to a specific word-level meaning. by visualizing the top activation sequences, we can see how word-level disambiguation is gradually developed in a transformer. we show how the disambiguation effect develops progressively through each layer in figure 3 . in figure 3 , the top 5 activated words and their contexts for transformer factor φ :,30 in different layers are listed. the top activated words in layer 0 contain the word "left" varying senses, which is being mostly disambiguated in layer 2 albeit not completely. in layer 4, the word "left" is fully disambiguated since the top-activated word contains only "left" with the word sense "leaving, exiting." we also show more examples of those types of transformer factors in further, we can quantify the quality of the disambiguation ability of the transformer model. in the example above, since the top 1000 activated words top 3 activated words and their contexts explanation φ:,2 and contexts are "left" with only the word sense "leave, exiting", we can assume "left" when used as a verb, triggers higher activation in φ :,30 than "left" used as other sense of speech. we can verify this hypothesis using a human-annotated corpus: brown corpus @xcite . in this corpus, each word is annotated with its corresponding part-of-speech. we collect all the sentences contains the word "left" annotated as a verb in one set and sentences contains "left" annotated as other part-of-speech. as shown in figure 4a , in layer 0, the average activation of φ :,30 for the word "left" marked as a verb is no different from "left" as other senses. however, at layer 2, "left" marked as a verb triggers a higher activation of φ :,30 . in layer 4, this difference further increases, indicating disambiguation develops progressively across layers. in fact, we plot the activation of "left" marked as verb and the activation of other "left" in figure 4b . in layer 4, they are nearly linearly separable by this (a) layer 4 (b) layer 6 (c) layer 8 figure 5: visualization of a mid-level transformer factor. (a), (b), (c) are the top 5 activated words and contexts for this transformer factor in layer-4, 6, and 8 respectively. again, the position of the word vector is marked blue. please notice that sometimes only a part of a word is marked blue. this is due to that bert uses word-piece tokenizer instead of whole word tokenizer. this transformer factor corresponds to the pattern of "consecutive adjective". as shown in the figure, this feature starts to develop at layer-4 and fully develops at layer-8. single feature. since each word "left" corresponds to an activation value, we can perform a logistic regression classification to differentiate those two types of "left". from the result shown in figure 4a , it is pretty fascinating to see that the disambiguation ability of just φ :,30 is better than the other two classifiers trained with supervised data. this result confirms that disambiguation is indeed done in the early part of pre-trained transformer model and we are able to detect it via dictionary learning. mid level: sentence-level pattern formation. we find most of the transformer factors, with an is curve peak after layer 6, capture mid-level or highlevel semantic meanings. in particular, the midlevel ones correspond to semantic patterns like phrases and sentences pattern. we first show two detailed examples of mid-level transformer factors. figure 5 shows a transformer factor that detects the pattern of consecutive usage of adjectives. this pattern starts to emerge at layer 4, develops at layer 6, and becomes quite reliable at layer 8. figure 6 shows a transformer factor, which corresponds to a pretty unexpected pattern: "unit exchange", e.g., 56 inches (140 cm). although this exact pattern only starts to appear at layer 8, the sub-structures that make this pattern, e.g., parenthesis and numbers, appear to trigger this factor in layers 4 and 6. thus this transformer factor is also φ:,13 places in us, followings the convention "city, state" 51.5 91.5 91.0 77. 5 while some mid-level transformer factors verify common semantic or syntactic patterns, there are also many surprising mid-level transformer factors. we list a few in for each listed transformer factor, we analyze the top 200 activating words and their contexts in each layer. we record the percentage of those words and contexts that correspond to the factors' semantic pattern in ## discussion dictionary learning has been successfully used to visualize the classical word embeddings @xcite . in this paper, we propose to use this simple method to visualize the representation learned in transformer networks to supplement the implicit "probing-tasks" methods. our results show that the learned transformer factors are relatively reliable and can even provide many surprising insights into the linguistic structures. this simple tool can open up the transformer networks and show the hierarchical semantic or syntactic representation learned at different stages. in short, we find word-level disambiguation, sentence-level pattern formation, and long-range dependency. the idea of a neural network learns low-level features in early layers, and abstract concepts in the later stages are very similar to the visualization in @xcite . dictionary learning can be a convenient tool to help visualize a broad category of neural networks with skip connections, like @xcite , vit models @xcite , etc. for more interested readers, we provide an interactive website foot_0 for the readers to gain some further insights.