TokenSHAP: Interpreting Large Language Models with Monte Carlo Shapley Value Estimation

Interpretability in machine learning has gained significant attention as models become increasingly complex. Methods for explaining AI systems can be broadly categorized into two approaches: black box methods and white box methods [1].

Black box methods, such as LIME [5] and SHAP [6], have emerged as popular approaches for explaining predictions across various ML models without requiring access to the model’s internal architecture or parameters. LIME provides local approximations of model behavior by perturbing input data, while SHAP unifies several feature attribution methods under the Shapley value framework. These methods are particularly valuable when working with proprietary or complex models where internal access is limited or impractical [2].

White box methods, on the other hand, require knowledge of and access to the model’s internal structure. These include techniques like gradient-based saliency maps [7] and layer-wise relevance propagation (LRP) [8]. While these methods can provide more detailed insights into the model’s decision-making process, they are limited to scenarios where the model architecture is fully accessible and understood [3].

Recent advancements include counterfactual explanations [9], which explore how altering inputs changes model predictions. While these methods offer valuable insights for tabular and image data, they face challenges when applied to the sequential and contextual nature of language, highlighting the need for specialized NLP interpretability techniques [4].

In the NLP domain, attention visualization techniques [10] have gained popularity, offering insights into which parts of the input a model focuses on. However, these visualizations often lack quantitative rigor. More sophisticated methods like Integrated Gradients [11] and Layer-wise Relevance Propagation (LRP) [8] provide continuous importance scores for input tokens but can struggle with gradient saturation and non-linearity in deep models.

Probing tasks [12] have also been employed to examine the representations learned by language models, revealing the types of linguistic information encoded at different layers. However, these methods do not directly interpret how inputs lead to specific outputs in inference tasks.

Shapley values, originating from cooperative game theory, have emerged as a powerful tool for feature importance estimation in machine learning. Lundberg and Lee’s SHAP method [6] unified several feature attribution techniques under the Shapley value framework, ensuring consistency and local accuracy. However, the computational intensity of exact Shapley value calculation has led to approximations like KernelSHAP and TreeSHAP, which are primarily designed for fixed-length feature vectors.

Applying Shapley values to NLP tasks presents unique challenges due to the combinatorial explosion of possible token subsets in variable-length text. Recent work by Sundararajan et al. [11] introduced TracIn to track the influence of training data points on predictions, but it doesn’t provide granular token-level insights for individual predictions.

TokenSHAP attributes importance to individual tokens or substrings in an input prompt by estimating their Shapley values. The Shapley value for a token represents its average marginal contribution to the model’s output across all possible combinations of tokens. This approach provides a rigorous framework for understanding how each part of the input influences the final response of large language models (LLMs).

Given an input prompt $x=(x_{1},...,x_{n})$, where $x_{i}$ represents individual tokens or substrings, we consider all possible subsets $S\subseteq N$, where $N=\{1,...,n\}$. The exponential number of subsets ($2^{n}$) makes exact computation impractical, so we employ Monte Carlo sampling to estimate Shapley values efficiently. This sampling approach balances the need for computational feasibility with the accuracy of Shapley value estimations.

We adapt the Monte Carlo sampling approach proposed by Castro et al. [16] to the context of text inputs. For each token $x_{i}$, we estimate its Shapley value $\phi_{i}$ as follows:

1. 1.

   Randomly select a subset size $k\in\{0,...,n-1\}$

2. 2.

   Randomly select a subset $S$ of size $k$ that does not include $i$

3. 3.

   Compute the marginal contribution: $v(S\cup\{i\})-v(S)$

4. 4.

   Repeat steps 1-3 $M$ times and average the results

This estimator is unbiased and converges to the true Shapley value as $M$ increases, ensuring the robustness and reliability of the importance scores generated. Importantly, in our adaptation, we methodically ensure that each token $x_{i}$ is excluded precisely once in the subsets considered for each calculation. This consistent exclusion of each token is crucial, as it provides a stable calculation base, ensuring that the estimations are not only unbiased but also systematically verifiable across different textual inputs. By doing so, we maintain methodological integrity, allowing for a robust basis of comparison and significance testing among tokens’ contributions to the model output.

Algorithm 1 provides a high-level overview of the TokenSHAP method. It iterates over each sample, performing Monte Carlo sampling to estimate its Shapley value. The algorithm efficiently computes marginal contributions and averages them to produce the final Shapley values for all tokens in the input prompt.

We define the value function $v(S)$ as the cosine similarity between the TF-IDF vectors of the model’s response to the subset $S$ and the response to the full prompt. Formally:

where $r(S)$ is the model’s response to the subset $S$, and $r(N)$ is the response to the full prompt. This formulation allows us to measure how closely the response to a subset resembles the response to the entire input, providing a quantitative basis for attributing importance to individual tokens.

For each sampled subset $S$, we query the LLM to generate a response. The prompt for a subset is constructed by concatenating the tokens or substrings corresponding to the indices in $S$. This step ensures that the model’s behavior is consistently evaluated across varying subsets of the input.

The estimated Shapley value for token $x_{i}$ is computed as:

where $S_{m}$ are the randomly sampled subsets, and $M$ is the number of Monte Carlo iterations. This aggregation of marginal contributions over multiple samples ensures a robust and reliable estimation of token importance.

We present the results using a color-coded visualization of the input text. The color intensity represents the magnitude of the Shapley value for each token or substring, with a diverging color map (e.g., coolwarm) to distinguish positive and negative values. This visualization aids in intuitively understanding the model’s decision-making process by highlighting the most influential parts of the input.

By providing a clear and quantitative analysis of token importance, TokenSHAP enhances the interpretability of LLMs, offering insights that are critical for improving model transparency, trustworthiness, and overall performance.

This experiment evaluates the ability of different interpretability methods to accurately assign low importance to randomly injected words within prompts. The goal is to test each method’s sensitivity and precision in identifying extraneous words that should not significantly impact model decisions.

TokenSHAP offers several advantages for interpreting LLM outputs:

1. 1.

   Quantitative Measure: It provides a rigorous, quantitative assessment of token importance, utilizing the Shapley value framework to quantify the contribution of each token to the model’s output in a consistent and objective manner.

2. 2.

   Context-awareness: The method captures the interdependence between tokens, reflecting how the model processes the entire input. This contextual sensitivity is essential for accurately interpreting the sophisticated dynamics of LLMs.

3. 3.

   Model-agnostic: TokenSHAP can be applied to any LLM without requiring access to its internal architecture, making it a versatile tool for users working with proprietary or black-box models. This positions TokenSHAP as a powerful black box method in the landscape of explainable AI, contrasting with white box methods that require detailed knowledge of model internals.

4. 4.

   Granularity: The approach allows for analysis at both token and substring levels, offering significant flexibility and enabling detailed exploration of how linguistic constructs larger than single tokens influence the model’s decisions.

As a black box method, TokenSHAP complements existing white box approaches in the field of explainable AI. While white box methods like gradient-based saliency maps offer insights tied directly to model parameters, TokenSHAP’s model-agnostic nature allows it to be applied more broadly, including to proprietary models or in situations where internal model access is restricted. This versatility makes TokenSHAP particularly valuable in real-world scenarios where model transparency may be limited due to commercial or security considerations.

1. 1.

   Computational Cost: Despite the efficiency gains from Monte Carlo sampling, TokenSHAP remains more computationally intensive than simpler interpretability methods, due to the need for repeated model evaluations.

2. 2.

   Sensitivity to Sampling: The stochastic nature of Monte Carlo sampling introduces variability in the importance scores, which may slightly vary between runs, affecting reproducibility in sensitive applications.

3. 3.

   Assumption of Additivity: The theoretical foundation of Shapley values assumes that contributions from individual tokens can be additively combined, which may not always be accurate in cases where complex interactions and non-linear dynamics dominate.

1. 1.

   Exploring Alternative Value Functions: Future research could include developing more sophisticated value functions that better capture nuanced aspects of semantic similarity and contextual alignment. Usage of LLM can also be considered for this task.

2. 2.

   Investigating Shapley Value Stability: Further studies are needed to assess the stability of Shapley values across different LLM architectures and input sizes, to understand their robustness and generalizability.

3. 3.

   Developing Interactive Tools: There is a substantial opportunity to create interactive, user-friendly tools that allow practitioners to dynamically explore token importance, enhancing the accessibility and practical utility of TokenSHAP.

4. 4.

   Extending to Multi-turn Conversations: Applying TokenSHAP to multi-turn conversational contexts could provide insights into how contextual understanding evolves in dialogue systems.

5. 5.

   Bias Analysis: Utilizing TokenSHAP for systematic identification and analysis of potential biases in LLM outputs could contribute to the development of more ethical and fair AI systems.