LLM Explainability via Attributive Masking Learning (2024)

Abstract: In this paper, we introduce Attributive Masking Learning (AML), a method designed for explaining language model predictions by learning input masks. AML trains an attribution model to identify influential tokens in the input for a given language model’s prediction. The central concept of AML is to train an auxiliary attribution model to simultaneously 1) mask as much input data as possible while ensuring that the language model’s prediction closely aligns with its prediction on the original input, and 2) ensure a significant change in the model’s prediction when applying the inverse (complement) of the same mask to the input. This dual-masking approach further enables the optimization of the explanation w.r.t. the metric of interest. We demonstrate the effectiveness of AML on both encoder-based and decoder-based language models, showcasing its superiority over a variety of state-of-the-art explanation methods on multiple benchmarks.

## introduction ai is transforming every facet of modern life, sparking innovation and driving progress across diverse fields. from groundbreaking advancements in computer vision @xcite @xcite @xcite and natural language processing (nlp) @xcite @xcite @xcite @xcite @xcite to audio synthesis @xcite @xcite and recommender systems @xcite @xcite @xcite @xcite @xcite . in the field of nlp, the transformer architecture @xcite , has given rise to increasingly sophisticated language models (lms) @xcite @xcite @xcite . these models, characterized by their growing size and complexity, have become pivotal components in numerous applications. consequently, the demand for a comprehensive understanding of the decision-making processes employed by lms has surged. however, the inherent complexity of these models often obscures the transparency of their predictions, necessitating the development of explainability methods to reveal the underlying factors in the input influencing their outputs. as a result, a large variety of explanation methods were developed @xcite @xcite @xcite @xcite . in tandem with the evolution of explanation methods, various explanation metrics and benchmarks were proposed to quantitatively assess the efficacy of explainability methods @xcite . however, consensus within the literature regarding the ultimate explanation metric remains elusive, with each metric providing unique insights into different facets of the explanation quality. in this work, we introduce the attributive masking learning (aml) framework that enables explanation via the generation of an attribution map tailored to the specific metric at hand. aml involves training an auxiliary attribution model to identify influential tokens in the input for a given lm's prediction. this attribution model is specifically trained to achieve two goals simultaneously: 1) masking as much input data as possible while ensuring the lm's predictions on the masked and original inputs remain closely aligned, and 2) inducing a significant change in the model's prediction when the inverse of the same mask is applied to the input. these two objectives bear resemblance to faithfulness metrics @xcite . the aml optimization is self-supervised and consisting of a two-phase process: a mandatory offline aml pretraining phase (paml), followed by an optional online aml finetuning phase (faml). in the paml phase, the attribution model is pretrained on a training set of examples and their respective predictions (generated by the lm to be explained). the pretraining process aims to enable the model to generate meaningful attribution maps by optimizing the aml dual-mask loss while monitoring the metric of interest. the paml phase occurs once and offline. at the end of the paml phase, an attribution map can be generated for any input through a single forward pass of the input and its prediction via the pretrained attribution model. for further refinement of the attribution map, an optional subsequent aml finetuning (faml) step can be employed in an online manner by continuously finetuning the pretrained attribution model on a specific instance w.r.t. the metric of interest, hence producing metric-adapted explanations. we present a comprehensive evaluation encompassing 13 explanation methods, 5 lms (including the latest llama2 and mistral models @xcite ; albert q. jiang, 2023)), 4 datasets, and 5 explanation metrics. our results demonstrate the effectiveness of aml. specifically, paml exhibits competitive performance with a notably fast runtime, while faml surpasses all state-of-the-art approaches, achieving significant improvements in accuracy across multiple benchmarks. in summary, our contributions include: 1) we introduce aml method, facilitating a novel dual-mask optimization of the attribution map w.r.t. the specific metric of interest. 2) we present a comprehensive investigation encompassing a wide range of datasets, explanation methods, and evaluation metrics, applied to various llm architectures. our findings show that aml set a new state-of-the-art in llm explainability. ## related work explainability research has made significant strides over the past decade, introducing a wide range of methods across various application domains @xcite @xcite @xcite @xcite @xcite @xcite . gradientbased methods rely on the gradients of the model's prediction score concerning the input embeddings. basic gradient methods, such as vanilla gradients @xcite and gradientxinput @xcite , operate on the principle of gradient computation. integrated gradients (ig) @xcite is a pathintegration method, computing gradients on interpolated points along a straight line between the data and an uninformative baseline. approaches like deeplift @xcite and gradi-entshap (lundberg and lee, 2017) can be viewed as approximations of ig. a recent addition, sequential integrated gradients (sig) @xcite , offers a path-integration method that addresses concerns about altered sentence meaning by computing the importance of each word while keeping other words fixed and creating interpolations between the baseline and the word of interest. relevance-decomposition methods break down the network's representation into vectors, each influencing the model's prediction differently. globenc @xcite and alti @xcite incorporate local-level decomposition, aggregating resulting vector norms using rollout @xcite to construct global-level explanations. the recent introduction of decompx @xcite showcases state-of-the-art results by constructing decomposed token representations and sequentially propagating them through the model without interlayer mixing. perturbation-based methods, including wellknown techniques like lime @xcite and shap (lundberg and lee, 2017), introduce perturbations to individual inputs or neurons and observe their impact on subsequent neurons in the network. solvex @xcite , on the other hand, is a search method that applies beamsearch as an approximation to an exhaustive search on the ranking of tokens to find the best explanation per metric. however, a major limitation of solvex is its inefficiency due to the application of naive beam-search, resulting in significantly slower runtime compared to methods that require a single or a small number of passes through the model, such as vanilla gradients and ig. in contrast to the aforementioned works, aml utilizes an auxiliary attribution model that learns to preserve and alter the prediction of the model to be explained by masking out less and more influential elements in the input, respectively and simultaneously. this unique process facilitates the creation of metric-adapted explanations through a single pass (paml) or, optionally, a small number of additional passes (faml) through the model during inference. ## attributive masking learning in this work, our focus is on classification tasks. therefore, we assume a model f that takes an input x v and produces an output f (x v ) ∈ [0, 1] c , where the majority of recent lm architectures, also known as large language models (llms) are decoder-based, performing a completion task by generating a sequence of tokens in the output @xcite . while these llms can also be finetuned by employing classification heads, a more common approach is classification via completion: the llm is prompted with the specific task either in a few-shot or zero-shot and instructed to output a token representing the correct class. for example, in sentiment analysis, if the relevant classes are defined to the llm (via prompt) as 'pos' and 'neg', one can compute the logit scores for the tokens 'pos' and 'neg' and apply softmax to obtain the probabilities of the token associated with each individual class. we note that this approach can be implemented either before or after finetuning the llm on the relevant dataset in a completion manner. in this work, we conduct experiments in various settings, including finetuning with classification heads on top of the llm as well as classification via completion (decoder-based models). aml enables the attribution of the prediction f (x v ) to individual tokens in x v by training an attribution model g θ (parameterized by θ). this model learns to mask as much data as possible from x v while preserving important information, ensuring that the predictions on the masked input and the inversely masked input foot_0 match and diverge from the prediction on the original input f (x v ), respectively. the attribution model g θ takes the input x v and information foot_1 regarding f (x v ) denoted by y v , and produces an attribution map the masking operation is carried out in a tokenwise manner and can be implemented as either the replacement or omission of the token to be masked in various forms, including, but not limited to: 1) replacement by designated tokens, e.g., @xmath0, respectively. 2) replacement by a randomly drawn token. 3) replacement by another token predicted based on the surrounding context of the token to be masked. 4) replacement by a new token embedding representing the mask token, which is learned as part of the optimization process. 5) omission of the token to be masked from the original sequence. in this work, we have opted for option 1, while leaving the exploration of other options for future research foot_2 . therefore, we assume the existence of a unique embedding @xmath1, such that every token representation to be masked is either replaced by m (hard masking) with some masking probability or combined with m (soft masking) according to the masking probability. in this study, we opt for soft masking, where each token representation in the masked input x ′ v is obtained by a convex combination with the mask token m according to the attribution map a v , and therefore, m is implemented as follows: ## experimental setup and results datasets we evaluate the explanation methods on 4 different datasets, offering a comprehensive evaluation across a wide spectrum of classification tasks and text lengths: sst2 @xcite , rotten tomatoes (rtn) @xcite , and imdb @xcite are binary sentiment classification in short, mediumsized, and long texts, respectively. emotion recognition (emr) @xcite : emotion classification (6 classes: sadness, joy, love, anger, fear, and surprise) predominantly in short texts. to ensure transparency and reproducibility, we processed datasets using the hugging-face library @xcite , employing a procedure similar to that outlined in @xcite . furthermore, our evaluation includes the annotated movie reviews datasets sourced from the eraser benchmark @xcite , where ground-truth annotations are provided by human annotators. these datasets are used for assessing the concordance between the generated explanations and human-annotated explanations. the code for data processing is available in our github repository. the results of this evaluation are presented in the appendix. ## conclusion in this paper, we introduced aml -a selfsupervised optimization framework for explaining lms. by introducing an auxiliary attribution model, aml navigates the complexity of explaining model predictions through a dual-masking approach. notably, aml presents a flexible loss facilitating the creation of a metric-adapted explanation. our extensive evaluation, encompassing 13 explanation methods, 4 datasets, 5 lms, and 5 explanation metrics, consistently demonstrates the superior performance of aml compared to other state-of-the-art methods.