More Identifiable yet Equally Performant Transformers for Text Classification (2021)

Abstract: Interpretability is an important aspect of the trustworthiness of a model’s predictions. Transformer’s predictions are widely explained by the attention weights, i.e., a probability distribution generated at its self-attention unit (head). Current empirical studies provide shreds of evidence that attention weights are not explanations by proving that they are not unique. A recent study showed theoretical justifications to this observation by proving the non-identifiability of attention weights. For a given input to a head and its output, if the attention weights generated in it are unique, we call the weights identifiable. In this work, we provide deeper theoretical analysis and empirical observations on the identifiability of attention weights. Ignored in the previous works, we find the attention weights are more identifiable than we currently perceive by uncovering the hidden role of the key vector. However, the weights are still prone to be non-unique attentions that make them unfit for interpretation. To tackle this issue, we provide a variant of the encoder layer that decouples the relationship between key and value vector and provides identifiable weights up to the desired length of the input. We prove the applicability of such variations by providing empirical justifications on varied text classification tasks. The implementations are available at https://github.com/declare-lab/identifiable-transformers .

## introduction widely adopted transformer architecture @xcite has obviated the need for sequential processing of the input that is enforced in traditional recurrent neural networks (rnn). as a result, compared to a single-layered lstm or rnn model, a single-layered transformer model is computationally more efficient, reflecting in a relatively shorter training time @xcite . this advantage encourages the training of deep transformer-based language models on largescale datasets. their learning on large corpora has already attained state-of-the-art (sota) performances in many downstream natural language processing (nlp) tasks. a large number of sota machine learning systems even beyond nlp @xcite are inspired by the building blocks of transformer that is multi-head self-attention @xcite . a model employing an attention-based mechanism generates a probability distribution @xmath0, . . . , a n } over the n input units @xmath1, . . . , z n }. the idea is to perform a weighted sum of inputs, denoted by n @xmath2, to produce a more context-involved output. the attention vector, a, are commonly interpreted as scores signifying the relative importance of input units. however, counter-intuitively, it is recently observed that the weights generated in the model do not provide meaningful explanations @xcite . attention weights are (structurally) identifiable if we can uniquely determine them from the output of the attention unit @xcite . identifiability of the attention weights is critical to the model's prediction to be interpretable and replicable. if the weights are not unique, explanatory insights from them might be misleading. the self -attention transforms an input sequence of vectors @xmath3, . . . , z n } to a contextualized output sequence @xmath4, . . . , y n }, where y @xmath5,i) z i . the scalar a (k,i) captures how much of the i th token contributes to the contextualization of k th token. a transformer layer consists of multiple heads, where each head performs selfattention computations, we break the head computations in two phases: the identifiability in transformer has been recently studied by @xcite which provides theoretical claims that under mild conditions of input length, attention weights are not unique to the head's output. essentially their proof was dedicated to the analysis of the computations in the second phase, i.e., token contextualization. however, the theoretical analysis ignored the crucial first phase where the attention weights are generated. intrinsic to their analysis, the attention identifiability can be studied by studying only the second phase of head computations. however, even if we find another set of weights from the second phase, it depends on the first phase if those weights can be generated as the part of key-query multiplication. in this work, we probe the identifiability of attention weights in transformer from a perspective that was ignored in @xcite . we explore the previously overlooked first phase of selfattention for its contribution to the identifiability in transformer. during our analysis of the first phase, we uncover the critical constraint imposed by the size of the key vector 1 d k . the flow of analysis can be described as • we first show that the attention weights are identifiable for the input sequence length d s no longer than the size of value vector d v ( §3.1) @xcite 2 . thus, the size of key vector plays an important role in the identifiability of the self-attention head. the pieces of evidence suggest that the current analysis in @xcite ignored the crucial constraints from the first phase in their analysis. to resolve the unidentifiability problem, we propose two simple solutions ( §4). for the regular setting of the transformer encoder where d v depends on the number of attention heads and token embedding dimension, we propose to reduce d k . this may lead to more identifiable attention weights. alternatively, as a more concrete solution, we propose to set d v equal to token embedding dimension while adding head outputs as opposed to the regular approach of concatenation @xcite . embedding dimension can be tuned according to the sequence length up to which identifiability is desired. we evaluate the performance of the proposed variants on varied text classification tasks comprising of ten datasets ( §5). in this paper, our goal is to provide concrete theoretical analysis, experimental observations, and possible simple solutions to identifiability of attention weights in transformer. the idea behind identifiable variants of the transformer is-the harder it is to obtain alternative attention weights, the likelier is they are identifiable, which is a desirable property of the architecture. thus, our contribution are as follows: ## identifiability of attention the output of an attention head h is the product of a and t (eq. ( 2 )). formally, we define identifiability of attention in a head: definition 3.1. for an attention head's output h, attention weights a are identifiable if there exists a unique solution of a @xmath6. the above definition can be reformulated as definition 3.2. a is unidentifiable if there exist an ã, ( ã = 0), such that (a + ã) is obtainable from phase-1 of head computations and satisfy ## solutions to identifiability based on the identifiability analysis in §3, we propose basic solutions to make transformer's attention weights identifiable. decoupling d k . contrary to the regular transformer setting where d @xmath7, a simple approach is to decrease the value of d k that is the size of the key and query vector. it will reduce the possible solutions of ã by putting harder constraints on the rank of attention logits, i.e., a l in eq. ( 9 ). however, theoretically, d k decides the upper bound on dimensions of the space to which token embeddings are projected before the dot product. higher the upper bound, more degree of freedom to choose the subspace dimensions as compared to the lower d k variants. thus, there is a plausible trade-off when choosing between d k induced identifiability and the upper bound on the dimension of projected space. ## experimental setup setting up the encoder. we normalize the text by lower casing, removing special characters, etc. 9 for each task, we construct separate 1-gram vocabulary (u ) and initialize a trainable randomly sampled token embedding (u × d e ) from n (0, 1). similarly, we randomly initialize a (d s-max × d e ) positional embedding. the encoder ( §2.2) takes input a sequence of token vectors (d s × d e ) with added positional vectors. the input is then projected to key and query vector of size d @xmath8, 2, 4, 8, 16, 32, 64, 128, 256}. for the regular transformer setting, we fix the number of heads h to 8 and the size of value vector d @xmath9. for each token at the input, the outputs of attention heads are concatenated to generate a d e -sized vector. for the identifiable variant of the transformer encoder, d @xmath10, this is equal to d s-max to keep it identifiable up to the maximum permissible number of tokens. the outputs of all the heads are then added. each token's contextualized representations (added head outputs) are then passed through the feed-forward network ( §2.2). for classification, we use the encoder layer's output for the first token and pass it through a linear classification layer. in datasets with more than two classes, the classifier output is softmaxed. in the case of snli, we use the shared encoder for both premise and hypothesis; the output of their first tokens is then concatenated just before the final classification layer. we use adam optimizer, with learning @xmath11.001, to minimize the cross-entropy loss between the target and predicted label. for all the experiments, we keep the batch size as 256 and train for 20 epochs. we report the test accuracy obtained at the epoch with the best validation accuracy. numerical rank. to generate the numerical rank plot on imdb dataset as shown in fig. 2 , we train a separate transformer encoder-based classifier. for a particular d s value, we sample 100 reviews from the dataset with token @xmath12. the clipping will ensure the number of tokens is d s before feeding it to the encoder. the numerical rank is calculated for t's obtained from the first head of the encoder. ## results and discussion for the identifiable variant, similar to §3.1, we plot the numerical rank of t with input sequence length as shown in fig. 6 . unlike fig. 2 , where dim ln(t) linearly increases after d @xmath13, we find the dimension is zero for a larger d s (∼ 380). the zero dimensional (left) null space of t con-firms there exist no nontrivial solution to the constraint constraint-r2, i.e., ã = {0}. thus, the attention weights a are identifiable for a larger range of length of the input sequence. it is important that the identifiability of attention weights should not come at the cost of reduced performance of the model. to investigate this issue, we compare the performance of the identifiable transformer encoder against its regular settings ( §6) on varied text classification tasks. for the regular setting, as discussed in §4 as one of the solutions, the transformer can be made identifiable by decreasing the size of the key vector d k . the rows of the another solution to identifiability is to increase d v to d e and add the heads' outputs. this setting corresponds to the add rows in the to the regular settings. for d @xmath17, as a general observation, we find the performance of add does not drop as drastically as con with an increase in d k . this could be due to the larger size of value vector leading to the more number of parameters in add that compensate for the significant reduction in the model's accuracy. on the large-scale datasets, we observe that add performs slightly better than con. intuitively, as shown in fig. 1 , we can increase the size of value vector to increase the dimension of the space on which each token is projected. a higher dimensional subspace can contain more semantic information to perform the specific task. even though the theoretical analysis shows the possibility of a full row rank of t and identifiable attention weights, the t obtained from a trained model might not contain all the rows linearly independent as d s increases. we can explain this from the semantic similarities between words cooccurring together @xcite . the similarity is captured as the semantic relationship, such as dot product, between vectors in a linear space. as the number of tokens in a sentence, i.e., d s increases, it becomes more likely to obtain a token vector from the linear combination of other tokens. ## conclusion this work probed transformer for identifiability of self-attention, i.e., the attention weights can be uniquely identified from the head's output. with theoretical analysis and supporting empirical evidence, we were able to identify the limitations of the existing study by @xcite . we found the study largely ignored the constraint coming from the first phase of self-attention in the encoder, i.e., the size of the key vector. later, we proved how we can utilize d k to make the attention weights more identifiable. to give a more concrete solution, we propose encoder variants that are more identifiable, theoretically as well as experimentally, for a large range of input sequence lengths. the identifiable variants do not show any performance drop when experiments are done on varied text classification tasks. future works may analyse the critical impact of identifiability on the explainability and interpretability of the transformer.