You shouldn’t miss this if you want to know how the context of your prompt affects in your Preferred LLM response.

Causal Prompt Attribution in Large Language Models (LLMs)

Introduction

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 Average Marginal Effect (AME)

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.

Definition of AME

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.

Algorithms for Sparse AME Estimation

We propose two algorithms to estimate sparse AMEs at the word level:

1. Uncompleted-Prompt Design

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.

2. Completed-Prompt Design

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.

Baseline Construction

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.

Computing AMEs via Regression

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)$.

Experimental Details

Headline Classification Example with Token-Level Attribution

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.

Context

Local Mayor launches Initiative to Enhance Urban Public Transport.

Query

Classify the headline into one of the following categories: Technology, Politics, Sports, Art, or Others.

Generated Response

Politics

Context Attribution

Below, we show how each token contributes to the classification decision. Higher positive scores indicate stronger influence toward “Politics”.

Interpretation: The words “mayor” and “transport” are key drivers toward classifying this headline under Politics, while “urban” and “public” also contribute positively.

Conclusion

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.

References

• Cohen-Wang, B., Shah, H., Georgiev, K., & Madry, A. (2024). Contextcite: Attributing model generation to context. arXiv preprint arXiv:2409.00729.
• Ghorbani, A., & Zou, J. (2019). Data shapley: Equitable valuation of data for machine learning. In International conference on machine learning (pp. 2242-2251). PMLR.
• Hammoudeh, Z., & Lowd, D. (2024). Training data influence analysis and estimation: a survey. Machine Learning, 113(5), 2351-2403.
• Jia, R., Dao, D., Wang, B., Hubis, F. A., Hynes, N., Gurel, N. M., Li, B., Zhang, C., Song, D., & Spanos, C. J. (2019). Towards efficient data valuation based on the shapley value. In The 22nd International Conference on Artificial Intelligence and Statistics (pp. 1167-1176). PMLR.
• Lin, J., Zhang, A., Lecuyer, M., Li, J., Panda, A., & Sen, S. (2022). Measuring the effect of training data on deep learning predictions via randomized experiments. ArXiv, abs/2206.10013.