Feature Attribution Methods for LLMs

Understanding why a Large Language Model generates a specific output is fundamental to ensuring its safety. When an LLM produces harmful, biased, or unexpected text, simply knowing that it happened isn't enough. We need tools to diagnose the cause, tracing the problematic output back to specific parts of the input or internal model states. Feature attribution methods provide a way to assign importance scores to input features, typically tokens or their embeddings, indicating how much each feature contributed to a particular prediction or behavior. This allows us to pinpoint potential triggers for unsafe responses, debug alignment failures, and better understand the model's decision-making process, directly supporting the goal of maintaining safety throughout the model's operational life.

Feature attribution aims to answer the question: "Which parts of the input most influenced this specific output?" For text models, the "input features" are usually the individual tokens provided in the prompt or context. The "output" could be the probability of the next token, the likelihood of generating a harmful continuation, a classification score (e.g., toxicity), or even the activation of specific internal neurons.

Several families of methods exist, each with its strengths and weaknesses:

Gradient-based Methods

These techniques leverage the gradients of the model's output with respect to its input embeddings. The intuition is straightforward: if a small change in an input token's embedding leads to a large change in the output, that token is considered influential.

• Saliency Maps: The simplest approach involves calculating the gradient of the target output (e.g., the logit of a specific token or a toxicity score) with respect to the input embeddings. The absolute value or square of the gradient magnitude for each token embedding is taken as its importance score. While easy to compute, raw gradients can be noisy and suffer from saturation effects, where the gradient becomes uninformatively small even if the feature is important.

• Integrated Gradients (IG): This method offers an alternative by addressing some limitations of simple saliency maps. IG computes importance by accumulating gradients along a straight path from a baseline input (often a zero-embedding vector representing "no input") to the actual input. It satisfies important axioms like completeness, meaning the sum of attributions across all input features equals the difference between the model's output for the input and the baseline. The core idea is captured by:

  $$IG_i(x) = (x_i - x'_i) \times \int_{\alpha=0}^{1} \frac{\partial F(x' + \alpha(x - x'))}{\partial x_i} d\alpha$$

  Here, $x$ is the input embedding vector, $x'$ is the baseline embedding vector, $F$ is the model's output function (e.g., producing a score), $x_i$ is the $i$-th dimension of the embedding, and the integral averages the gradients along the path from $x'$ to $x$. In practice, this integral is approximated numerically. Choosing an appropriate baseline ($x'$) is significant for interpretation. Common choices include zero vectors, padding token embeddings, or averaged embeddings.

Attention-based Methods

Transformers rely heavily on attention mechanisms. It's tempting to directly use the attention weights assigned by the model's attention heads as feature importance scores. If token A attends strongly to token B in a given layer, perhaps token B is important for token A's representation?

• Direct Attention Weights: Visualizing attention patterns can reveal how information flows within the model. High attention weights between tokens might indicate a relationship the model considers relevant.

• Caveats: While interpretable and readily available, attention weights reflect internal information flow dynamics, which don't always equate directly to the final output's attribution. A token might receive high attention in early layers but have its influence modified or diminished by subsequent computations (LayerNorm, feed-forward networks, aggregation across heads/layers). Research has shown that attention weights can sometimes be misleading as direct indicators of feature importance for the final prediction compared to gradient or perturbation methods. However, they remain valuable for understanding intermediate processing steps.

Perturbation-based Methods

These methods assess feature importance by systematically modifying or removing parts of the input and observing the effect on the model's output.

• Occlusion/Masking: A straightforward approach is to mask or replace input tokens (e.g., with a padding token or a generic [MASK] token) one by one or in groups, then measure the change in the output probability or score. Tokens whose removal causes a large drop in the probability of a specific (perhaps undesirable) output are considered important for generating that output. This can be computationally expensive, requiring multiple forward passes through the model.

• LIME (Local Interpretable Model-agnostic Explanations): While model-agnostic, LIME can be adapted. It works by creating perturbed versions of the input instance (e.g., by removing words from a sentence) in the vicinity of the original input, getting model predictions for these neighbors, and then fitting a simple, interpretable model (like a weighted linear model) to these predictions. The weights of the simple model then serve as explanations for the local behavior of the complex LLM. Its applicability to very large models and long sequences can be limited by sampling efficiency.

SHAP (SHapley Additive exPlanations)

Based on Shapley values from cooperative game theory, SHAP provides a unified framework for feature attribution. It assigns an importance value to each feature that represents its marginal contribution to the output, averaged across all possible subsets (coalitions) of features.

• Properties: SHAP values have desirable theoretical properties like local accuracy (the sum of feature contributions equals the difference between the prediction and the average prediction), missingness (features missing from a subset have zero contribution), and consistency (if a model changes so a feature's contribution increases or stays the same regardless of other features, its SHAP value won't decrease). "* Computation: Calculating exact Shapley values is computationally prohibitive for most models. Approximations like KernelSHAP (which uses a LIME-like sampling approach with specific weighting) or methods tailored for deep learning models are often used. Applying SHAP effectively to large transformer architectures remains an area of active development due to computational demands and the complexity of model internals."

Applying Attribution for LLM Safety

Feature attribution becomes particularly useful when investigating safety-related incidents or behaviors:

1. Identifying Harmful Triggers: Suppose a user prompt like "Tell me about the history of group X" results in a biased or stereotypical response. Attribution methods can highlight whether the model's negative output is primarily driven by "group X" or perhaps by other subtle framing words in the prompt. High attribution scores on specific tokens can guide efforts to refine data filtering or alignment fine-tuning.

2. Debugging Alignment Failures: An instruction-following model might ignore a safety constraint in a complex prompt. Attribution can show if the model predominantly focused on the instruction part while assigning low importance to the constraint tokens, suggesting a need to improve its ability to handle multi-part instructions during alignment.

3. Analyzing Bias: We can examine the attribution scores for outputs related to different demographic groups mentioned in the input. For instance, if a sentiment analysis task consistently yields lower scores for texts discussing certain groups, attribution might reveal that tokens representing those groups disproportionately contribute to the negative sentiment prediction.

4. Understanding Refusals (or Lack Thereof): When an LLM correctly refuses a harmful request, attribution can show which input tokens triggered the safety mechanism (e.g., keywords related to violence or illegal acts). Conversely, if it fails to refuse, attribution might show low importance assigned to those same sensitive tokens, indicating a blind spot in the safety training.

Consider a simplified scenario where an LLM assigns a high "harmfulness" score to the input "How can I create fake news about elections?". We can use attribution to see which words contributed most to this score.

Simplified attribution scores for tokens in the input "How can I create fake news about elections?". Higher scores indicate greater contribution to the model predicting a high harmfulness score. The words "fake", "news", and "elections" show the highest attribution.

Challenges and Approaches

While powerful, feature attribution methods applied to LLMs come with caveats:

• Computational Cost: Methods like Integrated Gradients, perturbation analysis, and SHAP often require multiple forward passes or gradient computations, making them slow and expensive for large models.
• Baseline Choice: The interpretation of methods like IG and SHAP depends critically on the chosen baseline input. A zero baseline might yield different attributions than a baseline using padding tokens.
• Tokenization: Attribution operates at the token level. Since words can be split into subword tokens (e.g., "attribution" -> "att", "ribution"), interpreting the importance of these fragments can be less intuitive than word-level importance.
• Interaction Effects: Most methods assign scores to individual tokens, potentially missing complex, non-linear interactions between multiple tokens that collectively drive an output.
• Faithfulness vs. Plausibility: An explanation is faithful if it accurately reflects the model's internal reasoning. It's plausible if it makes sense to a human observer. Attention weights might offer plausible but not always faithful explanations. Methods like IG aim for higher faithfulness but can sometimes produce less intuitive results.

In summary, feature attribution techniques like Integrated Gradients, attention analysis, and perturbation methods are valuable tools in the LLM safety toolkit. They provide granular insights into which parts of an input drive specific model behaviors, enabling more targeted debugging of safety failures, analysis of bias, and verification of alignment mechanisms. However, they should be applied with an understanding of their computational costs and interpretation, often complementing other interpretability approaches and evaluation frameworks for a comprehensive safety assessment.