T oken S hapley: Token Level Context Attribution with Shapley Value (2025)

Abstract: Large language models (LLMs) demonstrate strong capabilities in in-context learning, but verifying the correctness of their generated responses remains a challenge. Prior work has explored attribution at the sentence level, but these methods fall short when users seek attribution for specific keywords within the response, such as numbers, years, or names. To address this limitation, we propose TokenShapley, a novel token-level attribution method that combines Shapley value-based data attribution with KNN-based retrieval techniques inspired by recent advances in KNN-augmented LLMs. By leveraging a precomputed datastore for contextual retrieval and computing Shapley values to quantify token importance, TokenShapley provides a fine-grained data attribution approach. Extensive evaluations on four benchmarks show that TokenShapley outperforms state-of-the-art baselines in token-level attribution, achieving a 11–23% improvement in accuracy.

## introduction in context learning (icl) @xcite has emerged as a powerful capability of large language models (llms) @xcite @xcite . llms can perform a wide range of tasks (for example, question answering @xcite , solving math problems) by learning from the context information users provide. icl is particularly valuable in scenarios where data collection and retraining are costly or impractical (for example, medical diagnosis @xcite ). however, responses generated by large language models (llms) are not inherently grounded in truth, as these models can sometimes hallucinate @xcite , producing factually incorrect or misleading information. this can lead to serious risks, particularly in high-stakes applications (e.g., healthcare diagnosis). as such, fact-checking here the response contains multiple information. doing sentencelevel attribution may provide incomplete evidence. a better way is to provide token-level attribution. @xcite and context attribution @xcite is crucial to ensure that llm responses are both interpretable and verifiable. consider an llm that recommends medication based on a medical database (figure 1 ). a user may want to get evidence from key tokens, such as "migraine frequency" and "liver toxicity", that capture comprehensive evidence of both benefits and risks, which might be overlooked with a single attributed sentence in the context. by pinpointing the tokens behind a response and providing supporting evidence, users can more easily assess its reliability. approaches to attribution in llms generally fall into three categories @xcite . 1) direct generated attribution @xcite . these methods prompt llms to generate self-attributions to improve response accuracy. however, researchers have found that llms often struggle to provide reliable evidence, particularly for domain-specific questions @xcite . 2) post-retrieval answering @xcite . these approaches involve additional model updates or training to improve response accuracy. however, the requirement for direct access to the model's parameters can introduce additional computational overhead. 3) post-generation at-tribution @xcite @xcite . these methods utilize search engines, document retrieval systems, or surrogate models to retrieve evidence based on input questions and generated answers. they primarily attribute information at the sentence level, which lacks fine-grained attribution when a sentence contains multiple pieces of information. see figure 1 for an example. to address the limitations of existing attribution methods, we aim to propose a new attribution method that 1) provides accurate evidence. 2) eliminates the need for additional tuning with llms. 3) enables fine-grained attribution. this motivates us to adopt the shapley value @xcite , which fairly attributes contributions by distributing credit among input components based on their marginal impact. shapley values have been widely used for training data attribution @xcite . in supervised learning tasks, when a k-nearest neighbors (knn) model is used for attribution, shapley values can be computed in polynomial time @xcite . however, applying knn for attribution in llms remains challenging due to the vague notion of "labels" among tokens. we observe that knn-lm @xcite enhances llm generation by leveraging a knn model, which stores (key, value) pairs for retrieval, where the key represents the prefix of a given token and the value is the token itself. building on this idea, we introduce token-shapley, a novel token-level attribution method that employs a knn model specifically for attribution. the flowchart of tokenshapley is shown in figure 2 , with further details provided in section 4. ## related work we provide additional related work not discussed in the previous section. ## problem setup and background in this section, we will start with the problem setup of context attribution. the notation used in this paper is listed in ## algorithm the goal of context attribution is to identify the most contributive sources within the context d that affect the model's response r. specifically, we aim to assign scores to each token in the context z @xmath0, reflecting its contribution to a given token in the response ẑt ∈ r. recall that, in order to run the knn-shapley value algorithm which gives exact shapley values in polynomial run time, we need to design a knn model in language model generation. in section 4.1 we follow the design of knn-lm to create a datastore with (key, value) pairs for the knn model. then in section 4.2, we design the utility function v(•) for efficient knn-shapley value. in section 4.3, we briefly review the core technique used in @xcite . lastly, we show how to accumulate shapley values for attribution from tokens to tokens. algorithm 1 shows the pseudocode for the proposed tokenshapley algorithm. ## experiments in this section, we compare tokenshapley with several baselines across four datasets. we use a single a100 gpu with 80 gb memory for experiments. for parameters in tokenshapley, we set @xmath1. ## discussion on computational cost our method involves two main computational components: (1) offline datastore construction, and (2) runtime knn-shapley inference. datastore construction. this step is performed once and offline. for each token in the context dataset, we cache its prefix embedding, resulting in a datastore whose size is linear in the number of context tokens. for example, constructing the datastore for the naturalquestions dataset requires approximately 100kb of memory per example (com-puted as embedding dimension 768 × number of tokens). importantly, we only store one wikipedia page per question, rather than the full wikipedia corpus, which greatly reduces the memory footprint. runtime inference. at inference time, the knn-shapley computation retrieves top-k neighbors from the datastore to assess token-level attribution. when @xmath2, the computational complexity is o(n), where n is the number of tokens in the input context. in general, for @xmath3, the complexity becomes o(n 2 ), though n remains small in practice (e.g., <512 tokens). we further reduce overhead by precomputing and storing datastore embeddings. given the small per-example memory cost and the offline nature of the datastore construction, the overall computational and memory overhead of our method is minimal. in addition, the modest datastore size allows it to be stored efficiently in memory or on low-cost cloud storage, making our approach practical for real-world applications. ## conclusion in this work, we introduced tokenshapley, a novel token-level attribution method that integrates shapley value-based data attribution with knn-based retrieval, leveraging a precomputed datastore for efficient contextual retrieval. by computing shapley values to quantify token importance, tokenshapley enables a more precise and interpretable attribution of language model outputs.