Probabilistic Attribution For Large Language Models

Technical Summary: Probabilistic Attribution For Large Language Models

Problem Statement
The generative nature of Large Language Models (LLMs) relies on computing conditional probabilities to sample response tokens based on previous context. While these models are increasingly deployed in high-stakes domains, their "black box" nature complicates the assessment of output validity, uncertainty, and the specific influence of input tokens on generated responses. Existing Explainable AI (XAI) methods often fall into two categories: model-specific approaches (e.g., gradient-based saliency maps) that require access to internal architecture, and model-agnostic approaches (e.g., SHAP, LIME) that rely on input-output perturbation but may lack stability or causal faithfulness. There is a need for a method that interprets LLM behavior through the lens of statistical first principles, independent of the underlying computational structure, to better understand token sensitivity, model stability, and training convergence.

Methodology
The authors situate LLM inference within the mathematical theory of stochastic processes. In this framework, an LLM is viewed as defining a probability distribution over sequences of tokens, where the probability measure is constructed from the model's training data. The core of the methodology is the Attribution Score (AS), a model-agnostic metric designed to quantify the importance of a specific token in a prompt given a generated response.

The AS is derived by inverting the next-token log-probabilities using Bayes' rule. Specifically, for a prompt $p$ and response $r$, the attribution score $A_\mu$ for the $\mu$-th token in the prompt is defined as the log-ratio of two probabilities:

1. The probability of the response $r$ given the full prompt $p$.
2. The probability of the response $r$ given the prompt with the $\mu$-th token marginalized away (i.e., $p$ where $p_\mu$ is allowed to take any value in the vocabulary).

Mathematically, $A_\mu = \log \frac{Pr(R=r | P=p)}{Pr(R=r | P - P_\mu = p - p_\mu)}$.

Since LLMs do not directly output probabilities for non-contiguous or marginalized sequences, the authors derive a method to reconstruct the denominator using only the standard conditional probabilities emitted by the model (Equations 4–9). This involves summing over all possible vocabulary tokens for the marginalized position, weighted by the probability of the remaining context.

Complementing the AS, the authors introduce Contextual Entropy measures:

• Prompt-only Contextual Entropy ($S^{(P)}$): The entropy of the token distribution at a specific position conditioned only on the remaining prompt tokens.
• Prompt+Response Contextual Entropy ($S^{(P+R)}$): The entropy conditioned on both the remaining prompt and the generated response.
• KL Divergence: The divergence between the token distributions under the two contexts, used to detect anomalies where the response significantly alters the model's uncertainty about a prompt token.

Additionally, Replacement Entropy is computed by measuring the frequency distribution of exact response matches when prompt tokens are replaced by the top 90% candidate tokens, providing an empirical measure of response stability.

Key Contributions
The paper outlines four primary contributions:

1. Model-Agnostic Attribution: The AS is applicable to text-generation tasks across any LLM without requiring access to internal weights or gradients, relying solely on emitted log-probabilities.
2. Parameter-Free Design: The method requires no hyperparameters (unlike LIME or ReAgent), enhancing reproducibility and objectivity.
3. Stability and Convergence Metrics: The contextual and replacement entropy measures serve as indicators of model stability and the degree to which a model has converged to its training data distribution.
4. Built-in Sensitivity: Because the method is grounded in statistical first principles, the attribution score inherently evaluates the correctness of the explainable results as part of its core methodology.

Experimental Results
The authors evaluated 8 LLMs (ranging from 137M to 1.7B parameters, including GPT-2, GPT-Neo, Llama, and Qwen variants) across 7 prompts involving various tasks (e.g., reading comprehension, fact retrieval, and deceptive queries).

• Token Sensitivity: High AS values corresponded to tokens that were critical for the specific response, while negative AS values indicated tokens the model would have preferred to replace given the response context.
• Model Convergence: Analysis of contextual entropy revealed that older models (GPT-2, GPT-Neo) exhibited higher entropy at near-zero AS values compared to modern instruction-tuned models (e.g., Llama 3.1). The authors interpret high entropy at near-zero AS as "noise tokens"—tokens that contribute little to the response probability yet remain unconstrained by context—suggesting poorer training convergence in older models.
• Anomaly Detection: The framework successfully identified specific anomalies, such as tokens where the model preferred a rephrasing of the prompt (e.g., preferring "in" over "?" in a specific question) or where the response context unexpectedly broadened the token distribution (increasing entropy).
• Evaluation Against XAI Baselines: The AS was compared against Input X Gradient, Gradient Shapley, LIME, Occlusion, and ReAgent using metrics like Infidelity, Deletion/NAOPC, Comprehensiveness, and Sufficiency. AS performed competitively, often matching or exceeding the performance of other methods, particularly noting its robustness across different models without architecture-specific tuning.
• Computational Cost: The method is computationally intensive, requiring GPU hours proportional to the vocabulary size due to the summation over all possible replacement tokens. However, the authors note potential for optimization by restricting summations to top-probability tokens.

Significance and Claims
The paper claims that by framing LLMs as stochastic processes, it provides a rigorous, model-agnostic framework for interpretability that avoids assumptions about internal "knowledge representation." The significance lies in the ability to:

• Quantify the causal influence of prompt tokens on responses without accessing model internals.
• Diagnose model training quality (convergence) through entropy analysis, identifying "noise tokens" that indicate poor training.
• Guide users to uncertain or unstable parts of a generation, thereby improving interpretability.
• Serve as a parameter-free baseline for tuning hyperparameters in other, faster XAI methods.

The authors conclude that while the approach incurs higher computational costs, it offers a deeper, statistically grounded understanding of LLM behavior, enabling the identification of training deficiencies and the engineering of prompts to produce more coherent responses.