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

TL;DR

This work addresses the interpretability of transformer representations by modeling contextualized embeddings as sparse, non-negative combinations of latent 'transformer factors' learned via dictionary learning. A single factor dictionary Φ is learned across all layers using non-negative sparse coding, enabling x ≈ Φ α with α ≥ 0, and visualized through top-activation sequences and LIME-based token saliency. Empirically, the approach reveals a hierarchy of patterns: low-level word sense disambiguation, mid-level sentence or phrase patterns, and high-level long-range dependencies, with some factors aligning with linguistic expectations and others yielding novel insights. The method provides a complementary, visualization-based alternative to probing tasks for understanding how transformers organize information across layers, and it includes an interactive resource to explore the discovered factors.

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

Figures (10)

• Figure 1: Building block (layer) of transformer
• Figure 2: Importance score (IS) across all layers for two different transformer factors. (a) This figure shows a typical IS curve of a transformer factor corresponding to low-level information. (b) This figure shows a typical IS curve of a transformer factor corresponds to mid-level information.
• Figure 3: Visualization of a low-level transformer factor, $\Phi_{:,30}$ at different layers. (a), (b) and (c) are the top-activated words and contexts for $\Phi_{:,30}$ in layer-$0$, $2$ and $4$ respectively. We can see that at layer-$0$, this transformer factor corresponds to word vectors that encode the word "left" with different senses. In layer-2, a majority of the top activated words "left" correspond to a single sense, "leaving, exiting." In layer 4, all of the top-activated words "left" have corresponded to the same sense, "leaving, exiting." Due to space limitations, we invite the readers to use our website to see more of those disambiguation effects.
• Figure 4: (a) Average activation of $\Phi_{:,30}$ for word vector "left" across different layers. (b) Instead of averaging, we plot the activation of all "left" with different contexts in layer-$0$, $2$, and $4$. Random noise is added to the y-axis to prevent overplotting. The activation of $\Phi_{:,30}$ for two different word senses of "left" is blended together in layer-$0$. They disentangle to a great extent in layer-$2$ and nearly separable in layer-$4$ by this single dimension.
• 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$.