Exploring Neural Language Models via Analysis of Local and Global Self-Attention Spaces (2021)

Abstract: Large pretrained language models using the transformer neural network architecture are becoming a dominant methodology for many natural language processing tasks, such as question answering, text classification, word sense disambiguation, text completion and machine translation. Commonly comprising hundreds of millions of parameters, these models offer state-of-the-art performance, but at the expense of interpretability. The attention mechanism is the main component of transformer networks. We present AttViz, a method for exploration of self-attention in transformer networks, which can help in explanation and debugging of the trained models by showing associations between text tokens in an input sequence. We show that existing deep learning pipelines can be explored with AttViz, which offers novel visualizations of the attention heads and their aggregations. We implemented the proposed methods in an online toolkit and an offline library. Using examples from news analysis, we demonstrate how AttViz can be used to inspect and potentially better understand what a model has learned.

## introduction currently the most successful machine learning approaches for text-related tasks predominantly use large language models. they are implemented with transformer neural network architecture @xcite , extensively pretrained on large text corpora to capture context-dependent meanings of individual tokens @xcite @xcite . even though training of such neural networks with hundreds of millions of the online visualization ( http://attviz.ijs.si ; first part of the paper) offers direct exploration of token attention across the space of input documents; its purpose is anomaly detection and general inspection of the attention space (of trained models). the offline part of attviz (second part of the paper) is a python library that offers computationally more demanding statistical analyses, ranging from visualization of key tokens for each attention head, comparison of the attention head properties via fuji integrals, and inspection of the attention distribution per-token basis. as neural networks require numerical input, words are first transformed into a high dimensional numeric vector space, in a process called embedding that aims to preserve similarities and relations between words. visualizations of embedding spaces is becoming ubiquitous in contemporary natural language processing. for example, google's online embedding projector foot_0 offers numerous visualizations for technically non-savvy users, by projecting word vectors to low dimensional (humanunderstandable) spaces. while visualization of embedding spaces is already accessible, visualization of internal workings of complex transformer neural networks (e.g.,their self-attention mechanism) is a challenging task. the works of @xcite and @xcite attempt to unveil the workings of black-box attention layers and offer an interface for human researches to learn and inspect their models. @xcite visualize the attention space by coloring it, and @xcite visualize the self-attention with examples from a sentiment analysis. in this work, we present attviz, an online system that focuses exclusively on self-attention and introduces two novel ways of visualizing this prop-erty. the tool serves as an additional tool in the toolbox of a language model researcher, offering exploration of the learned models with minimal effort. attviz can interactively aggregate the attention vectors and offers simultaneous exploration of the output probability space, as well as the attention space. a schematic overview of the proposed work is shown in figure 1 , and the main contributions are summarised as follows: the remainder of the paper is structured as follows. in section 2, we discuss works related to the proposed attviz approach. in section 3, we present the key ideas and technical implementation of the online part of the attviz system, including a use case on news classification. in section 4, we discuss the capabilities of the attviz library, available in an offline mode, and showcase its use on additional two datasets. in section 5 we discuss capabilities and limitations of attviz, present conclusions, and propose ideas for future work. ## background and related work on attention visualization neural language models are becoming the prevailing methodology for solving various text-related tasks, from entity recognition to classification. visualization of the attention mechanism that is the key component of such models has recently emerged as an active research area due to an increased popularity of attention based methods in natural language processing. recent deep neural network language models such as bert @xcite , xlnet @xcite , and roberta @xcite consist of multiple attention heads-separate weight spaces each associated with the input sequence in a unique way. these transformer language models consist of multiple attention matrices, all contributing to the final prediction. visualising the attention weights from each of the attention matrices is an important component in understanding and interpreting these models. the attention mechanism, which originated in the neural machine translation, lends itself naturally to visualisation. @xcite used heat maps to display the attention weights between input and output text. this visualisation technique was first applied in machine translation but found its use in many other tasks. @xcite visualized an input sentence and the output abstractive summary, while @xcite showed an association between an input document and a textual entailment hypothesis on the output. in these heat map visualisations, a matrix is used to represent the token-token pairs and color intensity illustrates attention weights. this provides a summary of the attention patterns describing how they map the input to the output. for classification tasks, a similar visualisation approach can be used to display the attention weights between the classified document and the predicted label @xcite . here, the visualisation of attention often displays the input document with the attention weights superimposed onto individual words. the superimposed attention weights are represented similarly to heat map visualisations, using the color saturation to encode attention value. the neat-vision tool foot_1 encodes attention weights associated with input text in this manner. similarly, the text attention heatmap visualization (tahv foot_2 ) which is included in the ncrf++ toolkit @xcite can be used to generate weighted sequences which are visualised using superimposed attention scores. the purpose of the proposed attviz is to unveil the attention layer space to human explorers in an intuitive manner. the tool emphasizes selfattention, that is, the diagonal of the token-token attention matrix which possibly corresponds to the relevance of individual tokens. using different encoding techniques, attention weights across the layers and attention heads can be explored dynamically to investigate the interactions between the model and the input data. the attviz tool differs from other tools in that if focuses on self-attention, thus allowing visualization of (attention-annotated) input token sequences to be carried out directly. ## attviz: an online toolkit for visualization of self-attention attviz is an online visualization tool that can visualize neural language models from the pytorchtransformers library foot_3 -one of the most widely used resources for natural language modeling. the idea behind attviz is that it is simple to use and lightweight, therefore it does not offer computationally expensive (online) neural model training, but facilitates the exploration of trained models. along with attviz, we provide a set of python scripts that take as an input a trained neural language model and output a json file to be used by the attviz visualisation tool. a common pipeline for using attviz is outlined in figure 1 . first, a transformerbased trained neural network model is chosen to obtain predictions on a desired set of instances (documents or some other texts). the predictions are converted into the json format suitable for use with the attviz tool, along with the attention space of the language model. the json file is loaded into the attviz tool (on the user's machine, i.e. on the client side), where its visualization and exploration is possible. in sections 3.1 and 3.3, we present the proposed self-attention visualizations, followed by an example of their use on the news classification task in section 3.4. ## attviz library: statistical analysis of the attention space in section 3 we presented how the online version of attviz can be used for direct analysis of model output (in the json format). albeit suitable for quick inspections, the online system has its limitations such as poor support for computationally more intensive types of analysis (in terms of waiting times), and the lack of customized visualization tools accessible in the python ecosystem. to address these aspects, we developed attviz library that offers more detailed analysis of a given neural language model's properties. the library operates on the same json structures as the online version and is compatible with the initial user input. we demonstrate the analytical capabilities of our visualization tools on three datasets. the bbc news classification was already presented in section 3.4. ## discussion and conclusions as attviz is an online and offline toolkit for attention exploration, we discuss possible concerns regarding its use, namely: privacy, memory and performance overheads, and coverage. privacy is a potential concern for most web-based systems. as currently attviz does not employ any anonymization strategy, private processing of the input data is not guaranteed. while we intend to address this issue in furture work, a private installation of the tool can be done to get around this current limitation. attviz uses the users' computing capabilities, which means that large data sets may cause memory overheads when a large number of instances is loaded (typically several million). such situa-tions are difficult to address with attviz and similar web-based tool, but users can filter instances before using them in attviz and explore a subset of the data (e.g., only (in)correctly predicted instances, or certain time slot of instances). finally, attviz is focused on the exploration of self-attention. this is not the only important aspect of a transformer neural network, but it is the one, where visualisation techniques have not yet been sufficiently explored. similarly to the work of @xcite , we plan to further explore potentially interesting relations emerging from the attention matrices.