Large Language Models (LLMs) have revolutionized natural language processing, achieving impressive performance across various tasks. However, understanding how individual words in a prompt influence the model’s output remains a challenge. This blog post introduces a novel framework for quantifying the influence of individual words in a prompt on the responses generated by LLMs. Our approach leverages the Average Marginal Effect (AME), a metric that measures the expected marginal contribution of adding a particular word to a randomly sampled subset of other words.
The AME is a scalable and principled measure of how individual elements, such as training data points or prompt words, contribute to a model’s behavior. In the context of LLMs, the AME quantifies the influence of individual words on the model’s output. When the sampling distribution is uniform, the AME aligns with the well-known Shapley value.
For a given element $n$, the AME is defined as: $ \text{AME}{n} = E{S^{n}\sim L^{n}}[U(S^{n}+{n})-U(S^{n})$ Here, $L^{n}$ is a distribution over subsets $S^{n}$ that exclude $n$, and $U(S)$ is a utility function that evaluates an outcome of interest (e.g., token-level probabilities) for a subset $S$ of the original set.
We propose two algorithms to estimate sparse AMEs at the word level:
This design aims to reveal word-level contributions by partially masking the original prompt $D_{\text{org}}$. Words are retained with probability $p$, and omitted words are replaced with a placeholder symbol. The LLM is then queried with this incomplete prompt, and the utility $U(D_{p})$ is measured. This approach provides insights into which words have a substantial influence on the model’s output.
This design includes an additional step: filling in the missing words with candidates drawn from a baseline distribution before querying the LLM. This approach simulates a more “complete” context, making the causal effect of including or excluding a particular word clearer.
To build the baseline distribution, we generate multiple incomplete prompts by sampling $p$ and masking words accordingly. For each incomplete prompt, we query the LLM to produce completions for the missing words. Over multiple trials, these completions form an empirical distribution of likely substitutes for each blank position.
After collecting $M$ sampled prompts and their corresponding utility values, we construct data matrices for AME estimation. Let $\mathbf{X}$ be an $M\times N$ matrix, where $N$ is the number of words in $D_{\text{org}}$, and $\mathbf{Y}$ be an $M$-dimensional vector of utility values. AME estimation proceeds through sparse linear regression, as described by Lin et al. (2022), achieving query complexity on the order of $(O(k\log N)$.
In this example, the LLM output is a single word, and we aim to find the attribution of the context in headline classification. We used the Uncompleted-Prompt Design with $M=200$ samples for estimating AME.
Local Mayor launches Initiative to Enhance Urban Public Transport.
Classify the headline into one of the following categories: Technology, Politics, Sports, Art, or Others.
Politics
Below, we show how each token contributes to the classification decision. Higher positive scores indicate stronger influence toward “Politics”.
| Token | Index | Coefficient |
|---|---|---|
| Local | 1 | 0.00 |
| mayor | 2 | 5.83 |
| launches | 3 | 0.00 |
| Initiative | 4 | -1.93 |
| to | 5 | 0.00 |
| enhance | 6 | 0.00 |
| urban | 7 | 2.76 |
| public | 8 | 4.06 |
| transport | 9 | 6.49 |
Interpretation: The words “mayor” and “transport” are key drivers toward classifying this headline under Politics, while “urban” and “public” also contribute positively.
Our framework provides a scalable, interpretable, and causally grounded means of detecting which parts of the prompt drive token-level generation in LLM outputs. By illuminating the underlying mechanisms that connect input words to generated text, our work aids in diagnosing biases, assessing data fidelity, and enhancing overall trust in large-scale language models.