VISIT : Visualizing and Interpreting the Semantic Information Flow of Transformers (2023)

Abstract: Recent advances in interpretability suggest we can project weights and hidden states of transformer-based language models (LMs) to their vocabulary, a transformation that makes them more human interpretable. In this paper, we investigate LM attention heads and memory values, the vectors the models dynamically create and recall while processing a given input. By analyzing the tokens they represent through this projection, we identify patterns in the information flow inside the attention mechanism. Based on our discoveries, we create a tool to visualize a forward pass of Generative Pre-trained Transformers (GPTs) as an interactive flow graph, with nodes representing neurons or hidden states and edges representing the interactions between them. Our visualization simplifies huge amounts of data into easy-to-read plots that can reflect the models’ internal processing, uncovering the contribution of each component to the models’ final prediction. Our visualization also unveils new insights about the role of layer norms as semantic filters that influence the models’ output, and about neurons that are always activated during forward passes and act as regularization vectors.

## introduction wouldn't it be useful to have something similar to an x-ray for transformers language models? recent work in interpretability found that hidden-states (hss), intermediate activations in a neural network, can reflect the "thought" process of transformer language models by projecting them to the vocabulary space using the same transformation that is applied to the model's final hs, a method known as the "logit lens" @xcite . for instance, the work of @xcite shows how the fully-connected blocks of transformer lms add information to the model's residual stream, the backbone route of information, promoting tokens that eventually make it to the final predictions. subsequent work by @xcite shows that projections of activated neurons, the static weights of the models' matrices, are correlated in their meaning to the projections of their block's outputs. this line of work suggests we can stop reading vectors (hss or neurons) as just numbers; rather, we can read them as words, to better understand what models "think" before making a prediction. these studies mostly interpret static components of the models or are limited to specific case studies that require resources or expertise. to address the gap in accessibility of the mechanisms behind transformers, some studies create tools to examine how lms operate, mostly by plotting tables of data on the most activated weights across generations or via plots that show the effect of the input or specific weights on a generation @xcite . yet, such tools do not present the role of each of the lm's components to get the full picture of the process. in this paper, we analyze another type of lms' components via the logit lens: the attention module's dynamic memory @xcite , the values (hs) the module recalls from previous inputs. we describe the semantic information flow inside the attention module, from input through keys and values to attention output, discovering patterns by which notions are passed between the lm's components into its final prediction. based on our discoveries, we model gpts as flow-graphs and create a dynamic tool showing the information flow in these models (for example, figure 1 ). the graphs simplify detection of the effect that single to small sets of neurons have on the prediction during forward passes. we use this tool to analyze gpt-2 @xcite in three case studies: (1) we reflect the mechanistic analysis of @xcite on indirect object identification each node represents a small group of neurons or hs, which are labeled by the top token of their projection to the vocabulary space. the plot should be read from left to right and includes the attention block: ln (the node at the end of the first purple edge), query, memory keys and values (with greenish edges) and the topmost activated attention output neurons (in blue and red), followed by the mlp: ln (in purple), first and second matrices' most activated neurons (in blue and red). the dark edges in the upper parts of the plot are the residuals of each sub-block. in a simple way; (2) we analyze the role of layer norm layers, finding they act as semantic filters; and (3) we discover neurons that are always activated, related to but distinct from rogue dimensions @xcite , which we term regularization neurons. ## tracing the semantics behind the attention's output in this section, we trace the components which create the semantics of the attention block's output, by comparing vectors at different places along the computation graph. in all the following experiments, we project hs into the vocabulary using the logit lens to get a ranking of all the tokens, then pick the top-k tokens according to their ranking. we measure the common top tokens of two vectors (x 1 and x 2 ) via their intersection score i k @xcite : throughout this section, we used counterfact @xcite , a dataset that contains factual statements, such as the prompt "the capital of norway is" and the correct answer "oslo". we generate 100 prompts randomly selected from coun-terfact using gpt-2-medium, which we verify the model answers correctly. we collect the hss from the model's last forward-passes (the passes that plot the answers) and calculate i @xmath0. foot_2 ## modeling the information flow as a flow-graph as in most neural networks, information processing in an autoregressive lm can be viewed as a flow graph. the input is a single sentence (a sequence of tokens), the final output is the probability of the next word, with intermediate nodes and edges. geva et al. (2022b, 2021 ) focused on information flow in and across the mlp blocks, while our analysis in section 3 focused on the information flow in the attention block. in this section, we describe how to construct a readable and succinct graph for the full network, down to the level of individual neurons. our graph is built on collected hs from a single forward pass: it uses single and small sets of hss as nodes, while edges are the interactions between nodes during the forward pass. one option for constructing a flow graph is to follow the network's full computation graph. common tools do this at the scale of matrices @xcite , coarser than the neuronal scale we seek. they usually produce huge, almost unreadable, graphs that lack information on which values are passed between matrices and their effect. similarly, if we were to connect all possible nodes figure 6 : while 6a shows quantitative results, the graph model in figure 6b shows qualitative information that is otherwise difficult to notice: the attention's ln is the first place in the model where the attention input's most probable token is " mary" (1). the negative name mover head from layer 10, represented by the blue cell in the table's 10-th row, is visualized in the graph with a red pop-up showing it assigns the token " mary" its lowest possible ranking, meaning its role is to reduce the probability of this token (2). the output of the attention block is the token " john" but its second most probable output is " mary" with around 2% chance (3). however, when added to the residual, together they predict almost 93% chance for " mary" (4). to present the semantic information flow, we assign each node with its most probable projected token and the ranking it gives to the model's final prediction, according to the logit lens. each node is colored based on its ranking, thereby emphasizing the correlation between the node's meaning and the final prediction. additionally, we utilize the width of the edges to reflect the scores used for pruning. figures 1 and 4 show static examples on one sentence, the first for a single transformer block's graph and the second with an annotated explanation on the attention sub-blocks's sub-graph. ## example of use and immediate discoveries the flow-graph model is especially beneficial for qualitative examinations of lms to enhance research and make new discoveries. in this section, we demonstrate this with several case studies. ## related work derived from the original logit lens (nostalgebraist, 2020), several studies analyze the role of each component in lms using token projection @xcite . in the last few months, new studies suggest trainable transformation for projecting hs @xcite , promising to better project hs in the earlier layers of lms (which currently seems to have less alignment with the final output than later ones). other work took a more mechanistic approach in identifying the role of different weights, mostly by removing weights or changing either weights or activations, and examining how the final prediction of the altered model is affected @xcite @xcite . there has been much work analyzing the attention mechanism from various perspectives, like trying to assign linguistic meaning to attention scores, questioning their role as explanations or quantify its flow @xcite . see @xcite for an overview. our work is different from feature attribution methods @xcite , which focus on identifying the tokens in the input that exert a greater influence on the model's prediction. some studies visualise the inner computation in lms. for example, the work of @xcite tries to look into the inner representation of model by visualizing the logit lens projection of the hss between blocks and on the mlp weights. other tools that focused on the attention described the connection between input tokens @xcite but did not explore the internals of the attention module. there are general tools for visualizing deep learning models, like @xcite , but they only describe the flow of information between matrices, not between neurons. @xcite visualize hidden states and attention in recurrent neural network models, allowing for interaction and counterfactual exploration. ## conclusion in this work, we used token projection methods to trace the information flow in transformer-based lms. we have analyzed in detail the computation in the attention module from the perspective of intermediate semantics the model processes, and assessed the interactions between the attention memory values and attention output, and their effect on the residual stream and final output. based on the insights resulting from our analysis, we created a new tool for visualizing this information flow in lms. we conducted several case studies for the usability of our new tool, for instance revealing new insights about the role of the layer norm. we also confirmed the validity of our approach and showed how it can easily support other kinds of analyses. our tool and code will be made publicly available, in hope to support similar interpretations of various auto-regressive transformer models.