Token-wise Decomposition of Autoregressive Language Model Hidden States for Analyzing Model Predictions (2023)

Abstract: While there is much recent interest in studying why Transformer-based large language models make predictions the way they do, the complex computations performed within each layer have made their behavior somewhat opaque. To mitigate this opacity, this work presents a linear decomposition of final hidden states from autoregressive language models based on each initial input token, which is exact for virtually all contemporary Transformer architectures. This decomposition allows the definition of probability distributions that ablate the contribution of specific input tokens, which can be used to analyze their influence on model probabilities over a sequence of upcoming words with only one forward pass from the model. Using the change in next-word probability as a measure of importance, this work first examines which context words make the biggest contribution to language model predictions. Regression experiments suggest that Transformer-based language models rely primarily on collocational associations, followed by linguistic factors such as syntactic dependencies and coreference relationships in making next-word predictions. Additionally, analyses using these measures to predict syntactic dependencies and coreferent mention spans show that collocational association and repetitions of the same token largely explain the language models’ predictions on these tasks.

## introduction much of contemporary natural language processing (nlp) is driven by transformer-based large language models, which are trained to make predictions about words in their context by aggregating representations through their self-attention mechanism. the breakthrough in many nlp tasks these models have achieved has led to active research into interpreting their predictions and probing the knowledge embodied by these models @xcite @xcite . one line of such research focuses on quantifying the importance of each input token to the models' final output, but due to the complexity of the computations performed within the transformer layers, analysis has been limited to studying the self-attention mechanism and the feedforward neural network independently @xcite @xcite @xcite or has relied on e.g. gradient-based attribution methods @xcite that yield measures that are not interpretable in terms of output model probabilities. to address these limitations, this work presents a linear decomposition of final language model hidden states into the sum of final output representations of each initial input token and a cumulative bias term, which is schematized in figure 1 . this work focuses on decomposing autoregressive language models, in which the final hidden states are used to calculate a probability distribution over the next token. the decomposition allows the definition of probability distributions that ablate the contribution of specific input tokens, which can be used to study their impact on next-token probabilities with only one forward pass from the model. this decomposition is exact if the activation function of the feedforward neural network is differentiable almost everywhere, foot_0 and therefore it does not require perturbing the original computations of the language model (e.g. by using approximations) to gauge the influence of input tokens for virtually all contemporary transformer architectures. additionally, this work defines an intuitive importance measure for each context token based on the change in next-token log probability, which does not correlate strongly with layer-wise attention weights or gradient norms. since this measure is defined in terms of log probabilities, they can also be summed to quantify importance in predicting an arbitrary sequence of tokens according to the chain rule of conditional probabilities. using the proposed decomposition and associated importance measure, this work characterizes which kinds of context words autoregressive language models leverage most in order to make nextword predictions. results from stepwise regression analyses suggest that transformer-based language models rely mainly on collocational associations, followed by linguistic factors such as syntactic dependencies and coreference relationships. followup analyses using these importance measures to predict syntactic dependencies and coreferent mention spans additionally show that collocational association and repetitions of the same token largely explain the language models' predictions on these tasks. ## background: transformer decoder of autoregressive language models transformer-based autoregressive language models (e.g. @xcite @xcite use a variant of the multi-layer transformer decoder @xcite . each decoder layer consists of a masked self-attention block and a feedforward neural network, which together calculate a vector x l,@xmath0: n l,α , where α ∈ {in, out}, foot_3 is a vector-wise layer normalization operation @xcite that first standardizes the vector and subsequently conducts elementwise transformations using trainable parameters c l,α , b l,α ∈ r d : ## token-wise decomposition of language model hidden states this section provides a mathematical definition of the token-wise decomposition of language model hidden states, which allows the quantification of the contribution of each input token to the conditional probability of the next token. ## experiment 1: correlation with other importance measures this work first compares the decomposition-based ∆lp defined in section 3.2 with other measures of importance that have been used in the literature to examine the degree to which ∆lp may be redundant with them. to this end, pearson correlation coefficients were calculated between the proposed ∆lp and attention weights and gradient norms at a token level. ## experiment 2: characterizing high-importance context words having established that ∆lp provides a novel method to quantify the importance of each context token to language model predictions, the second experiment conducts a series of regression analyses to characterize high-importance context words (i.e. words with high ∆lp values) and shed light on which kinds of context words language models leverage most in order to make predictions about the next word. ## experiment 3: syntactic dependency and coreference prediction using ∆lp the previous experiment revealed that compared to measures of collocational association, syntactic dependency and coreference relationships were not as strong predictors of ∆lp. experiment 3 further examines the connection between high-importance context words and syntactic dependency and coreference relationships by using ∆lp to predict them independently and analyzing the extent to which each relationship type aligns with ∆lp. ## discussion and conclusion this work advances recent efforts to interpret the predictions of transformer-based large language models. to this end, a linear decomposition of final language model hidden states into the sum of final output representations of each initial input token and a cumulative bias term was presented. this decomposition is exact as long as the activation function of the feedforward neural network is differentiable almost everywhere, and therefore it is applicable to virtually all transformer-based architectures. additionally, this decomposition does not require perturbing any intermediate computations nor re-running the language model to examine the impact of each input token. the decomposition in turn allows the definition of probability distributions that ablate the influence of input tokens, which was used to define the importance measure ∆lp that quantifies the change in nexttoken log probability. the first experiment in this work demonstrated that ∆lp does not capture a redundant quantity from importance measures that have been used in previous work to examine language model predictions such as layer-wise attention weights or gradient norms. subsequently, based on the proposed ∆lp, a stepwise regression analysis was conducted to shed light on the characteristics of context words that autoregressive language models rely on most in order to make next-word predictions. the regression results show that transformer-based language models mainly leverage context words that form strong collocational associations with the predicted word, followed by context words that are in syntactic dependencies and coreference relationships with the predicted word. the high reliance on collocational associations is consistent with the mathematical analysis of transformers that a layer of selfattention effectively functions as a lookup table that tracks bigram statistics of the input data @xcite , as well as empirical observations that transformer-based autoregressive language models have a propensity to 'memorize' sequences from the training data @xcite . finally, as a follow-up analysis, ∆lp was used to predict syntactic dependencies and coreferent mentions to further examine their relationship to highimportance context words. the precision scores on both tasks revealed a large discrepancy in performance according to the type of syntactic dependency relations and entity mentions. on syntactic dependency prediction, ∆lp corresponded closer to words in relations with high collocational association such as compounds and adjectival modifiers, providing further support for its importance in a language model's next-word prediction. moreover, on coreferent antecedent selection, ∆lp more accurately identified previous mentions of proper nouns and common nouns that were already observed verbatim in context. this is consistent with the tendency of transformer-based language models to predict identical tokens from its context @xcite , which seems to be enabled by dedicated 'induction heads' @xcite that learn such in-context copying behavior. taken together, these results suggest that collocational association and verbatim repetitions strongly drive the predictions of transformer-based autoregressive language models. as such, the connection drawn between a large language model's computations and linguistic phenomena such as syntactic dependencies and coreference observed in previous work (e.g. manning et al., 2020) may underlyingly be the effects of these factors.