ContextCite: Attributing Model Generation to Context

Authors: benjamin  cohen-wang
harshay  shah
kristian  georgiev
aleksander m ądry
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:06:37

Abstract: How do language models use information provided as context when generating a response?
Can we infer whether a particular generated statement is actually grounded in the context, a misinterpretation, or fabricated?
To help answer these questions, we introduce the problem of *context attribution*: pinpointing the parts of the context (if any) that *led* a model to generate a particular statement.
We then present ContextCite, a simple and scalable method for context attribution that can be applied on top of any existing language model.
Finally, we showcase the utility of ContextCite through three applications:
(1) helping verify generated statements
(2) improving response quality by pruning the context and
(3) detecting poisoning attacks.
We provide code for ContextCite at https://github.com/MadryLab/context-cite.

Main paper: 
 introduction: 
 Suppose that we would like to use a language model to learn about recent news. We would first need to provide it with relevant articles as context 2 . We would then expect the language model to interact with this context to answer questions. Upon seeing a generated response, we might ask: is everything accurate? Did the model misinterpret any of the context or fabricate anything? Is the response actually grounded in the provided context? Answering these questions manually could be tedious-we would need to first read the articles ourselves and then verify the statements. To automate this process, prior work has focused on teaching models to generate citations: references to parts of the context that support a response [1] [2] [3] [4] [5] . They typically do so by explicitly training or prompting language models to produce citations.In this work, we explore a different type of citation: instead of teaching a language model to cite its sources, can we directly identify the pieces of information that it actually uses? Specifically, we ask:Can we pinpoint the parts of the context (if any) that led to a particular generated statement?We refer to this problem as context attribution. Suppose, for example, that a language model misinterprets a piece of information and generates an inaccurate statement. In this case, context attribution would surface the misinterpreted part of the context. On the other hand, suppose that a language model uses knowledge that it learned from pre-training to generate a statement, rather than the context. In this case, context attribution would indicate this by not attributing the statement to any part of the context.Unlike citations generated by language models, which can be difficult to validate [6, 7] , in principle it is easy to evaluate the efficacy of context attributions. Specifically, if a part of the context actually led to a particular generated response, then removing it should substantially affect this response. solar_eclipse_2024.pdf:…To witness this incredible total solar eclipse, you will need to be within the 115-mile-wide path of totality.[1] The path arches from Mexico to Texas to Maine...<solar_eclipse_2024.pdf> I live in Boston, MA. When and where should I go to see the eclipse?Since Boston is not on the path of totality, you'll only see a partial eclipse. You can travel to Maine, which is on the path of totality [1] …Figure 1 : CONTEXTCITE. Our context attribution method, CONTEXTCITE, traces any specified generated statement back to the parts of the context that are responsible for it.Formalizing context attribution (Section 2). We begin this work by formalizing the task of context attribution. Specifically, a context attribution method assigns a score to each part of the context indicating the degree to which it is responsible for a given generated statement. We provide metrics for evaluating these scores, guided by the intuition that removing high-scoring parts of the context should have a greater effect than removing low-scoring parts of the context.Performing context attribution with CONTEXTCITE (Sections 3 and 4). Next, we present CONTEXTCITE, a simple and scalable method for context attribution that can be applied on top of any existing language model (see Figure 1 ). CONTEXTCITE learns a surrogate model that approximates how a language model's response is affected by including or excluding each part of the context. This methodology closely follows prior work on attributing model behavior to features [8] [9] [10] and training examples [11, 12] . In the context attribution setting, we find that it is possible to learn a linear surrogate model that (1) faithfully models the language model's behavior and (2) can be efficiently estimated using a small number of additional inference passes. The weights of this surrogate model can be directly treated as attribution scores. We benchmark CONTEXTCITE against various baselines on a diverse set of generation tasks and find that it is indeed effective at identifying the parts of the context responsible for a given generated response.Applying context attribution (Section 5). Finally, we showcase the utility of CONTEXTCITE through three applications:1. Helping verify generated statements (Section 5.1): We hypothesize that if attributed sources do not also support a generated statement, then it is less likely to be accurate. We find that using CONTEXTCITE sources can greatly improve a language model's ability to verify the correctness of its own statements.2. Improving response quality by pruning the context (Section 5.2): Language models often struggle to correctly use individual pieces of information within long contexts [13, 14] . We use CON-TEXTCITE to select only the information that is most relevant for a given query, and then use this "pruned" context to regenerate the response. We find that doing so improves question answering performance on multiple benchmarks.3. Detecting context poisoning attacks (Section 5.3): Language models are vulnerable to context poisoning attacks: adversarial modifications to the context that can control the model's response to a given query [15] [16] [17] [18] [19] . We illustrate that CONTEXTCITE can consistently identify such attacks. 
problem statement: 
 In this section, we will introduce the problem of context attribution (Section 2.1) and define metrics for evaluating context attribution methods (Section 2.2). To start, we will consider attributing an entire generated response-we will discuss attributing specific statements in Section 2.3.Setup. Suppose that we use a language model to generate a response to a particular query given a context. Specifically, let p LM be an autoregressive language model: a model that defines a probability distribution over the next token given a sequence of preceding tokens. We write p LM (t i | t 1 , . . . , t i-1 ) to denote the probability of the next token being t i given the preceding tokens t 1 , . . . , t i-1 . Next, let C be a context consisting of tokens c 1 , . . . , c |C| and Q be a query consisting of tokens q 1 , . . . , q |Q| . We generate a response R consisting of tokens r 1 , . . . , r |R| by sampling from the model conditioned on the context and query. More formally, we generate the i th token r i of the response as follows:r i ∼ p LM (• | c 1 , .. . , c |C| , q 1 , . . . , q |Q| , r 1 , . . . , r i-1 )3 .We write p LM (R | C, Q) to denote the probability of generating the entire response R-the product of the probabilities of generating the individual response tokens-given the tokens of a context C and the tokens of a query Q.The goal of context attribution is to attribute a generated response back to specific parts of the context. We refer to these "parts of the context" as sources. Each source is just a subset of the tokens in the context; for example, each source might be a document, paragraph, sentence, or even a word. The choice of granularity depends on the application-in this work, we primarily focus on sentences as sources and use an off-the-shelf sentence tokenizer to partition the context into sources4 .A context attribution method τ accepts a list of d sources s 1 , . . . , s d and assigns a score to each source indicating its "importance" to the response. We formalize this task in the following definition: Definition 2.1 (Context attribution). Suppose that we are given a context C with sources s 1 , . . . , s d ∈ S (where S is the set of possible sources), a query Q, a language model p LM and a generated response R. A context attribution method τ (s 1 , . . . , s d ) is a function τ : S d → R d that assigns a score to each of the d sources. Each score is intended to signify the "importance" of the source to generating the response R.What do context attribution scores signify? So far, we have only stated that scores should signify how "important" a source is for generating a particular statement. But what does this actually mean? There are two types of attribution that we might be interested in: contributive and corroborative [20] . Contributive attribution identifies the sources that cause a model to generate a statement. Meanwhile, corroborative attribution identifies sources that support or imply a statement. There are several existing methods for corroborative attribution of language models [1, 2, 4, 5] . These methods typically involve explicitly training or prompting models to produce citations along with each statement they make.In this work, we study contributive context attributions. These attributions give rise to a diverse and distinct set of use cases and applications compared to corroborative attributions (we explore a few in Section 5). To see why, suppose that a model misinterprets a fact in the context and generates an inaccurate statement. A corroborative method might not find any attributions (because nothing in the context supports its statement). On the other hand, a contributive method would identify the fact that the model misinterpreted. We could then use this fact to help verify or correct the model's statement.How might we evaluate the quality of a (contributive) context attribution method? Intuitively, a source's score should reflect the degree to which the response would change if the source were excluded. We introduce two metrics to capture this intuition. The first metric, the top-k log-probability drop, measures the effect of excluding the highest-scoring sources on the probability of generating the original response. The second metric, the linear datamodeling score (LDS) [12] , measures the extent to which attribution scores can predict the effect of excluding a random subset of sources.To formalize these metrics, we first define a context ablation as a modification of the context that excludes certain sources. To exclude sources, we choose to simply remove the corresponding tokens from the context 5 . We write ABLATE(C, v) to denote a context C ablated according to a vector v ∈ {0, 1} d (with zeros specifying the sources to exclude). We are now ready to define the top-k log-probability drop: Definition 2.2 (Top-k log-probability drop). Suppose that we are given a context attribution method τ . Let v top-k (τ ) be an ablation vector that excludes the k highest-scoring sources according to τ . Then the top-k log-probability drop is defined asTop-k-drop(τ ) = log p LM (R | C, Q) original log-probability -log p LM (R | ABLATE(C, v top-k (τ )), Q)log-probability with top-k sources ablated.(1)The top-k log-probability drop is a useful metric for comparing methods for context attribution. In particular, if removing the highest-scoring sources of one attribution method causes a larger drop than removing those of another, then we consider the former method to be identifying sources that are more important (in the contributive sense).For a more fine-grained evaluation, we also consider whether attribution scores can accurately rank the effects of ablating different sets of sources on the log-probability of the response. Concretely, suppose that we sample a few different ablation vectors and compute the sum of the scores corresponding to the sources that are included by each. These summed scores may be viewed as the "predicted effects" of each ablation. We then measure the rank correlation between these predicted effects and the actual resulting probabilities. This metric, known as the linear datamodeling score (LDS), was first introduced by Park et al. [12] to evaluate methods for data attribution. Definition 2.3 (Linear datamodeling score). Suppose that we are given a context attribution method τ . Let v 1 , . . . , v m be m randomly sampled ablation vectors and let f (v 1 ), . . . , f (v m ) be the corresponding probabilities of generating the original response. That is, f(v i ) = p LM (R | ABLATE(C, v i ), Q).Let fτ (v) = ⟨τ (s 1 , . . . , s d ), v⟩ be the sum of the scores (according to τ ) corresponding to sources that are included by ablation vector v, i.e., the "predicted effect" of ablating according to v. Then the linear datamodeling score (LDS) of a context attribution method τ can be defined asEQUATIONwhere ρ is the Spearman rank correlation coefficient [21] .Until now, we have discussed attributing an entire generated response. In practice, we might be interested in attributing a particular statement, e.g., a sentence or phrase. We define a statement to be any contiguous selection of tokens r i , . . . , r j from the response. To extend our setup to attributing specific statements, we let a context attribution method τ accept an additional argument (i, j) specifying the start and end indices of the statement to attribute. Instead of considering the probability of generating the entire original response, we consider the probability of generating the selected statement. Formally, in the definitions above, we replacep LM (R | C, Q) with p LM ( r i , . . . , r j statement to attribute | C, Q, r 1 , . . . , r i-1response so far).We consider the more sophisticated attention-based explanation method of Abnar and Zuidema [38] . Attention rollout seeks to capture the propagated influence of each token on each other token. Specifically, we first average the attention weights of the heads within each layer. Let A ℓ ∈ R n×n denote the average attention weights for the ℓ'th layer, where n is the length of the sequence. Then the propagated attention weights for the ℓ'th layer, which we denote Ãℓ ∈ R n×n , are defined recursively as Ãℓ = A ℓ Ãℓ-1 for ℓ > 1 and Ã1 = A 1 . Attention rollout computes an "influence" of token j on token i by computing the product(A 0 A 1 • • • A L ) ijwhere L is the total number of layers. When the model contains residual connections (as ours do), the average attention weights are replaced with 0.5A ℓ + 0.5I when propagating influences.3. Gradient norm: Following Yin and Neubig [42] , in Section 4 we estimate the attribution score of each source by computing the ℓ 1 -norm of the log-probability gradient of the response with respect to the embeddings of tokens in the source. In Appendix B.3, we also consider the ℓ 2 -norm of these gradients, but find this to be slightly less effective. 
context attribution with contextcite: 
 In the previous section, we established that a context attribution method is effective insofar as it is able to predict the effect of including or excluding certain sources. In other words, given an ablation vector v, a context attribution method should inform how the probability of the original response,f (v) := p LM (R | ABLATE(C, v), Q),changes as a function of v. The design of CONTEXTCITE is driven by the following question: can we find a simple surrogate model f that approximates f well? If so, we could use the surrogate model f to understand how including or excluding subsets of sources would affect the probability of the original response (assuming that f is simple enough). Indeed, surrogate models have previously been used in this way to attribute predictions to training examples [11, 12, 23, 24] , model internals [25, 26] , and input features [8-10]; we discuss connections in detail in Appendix C.1. At a high-level, our approach consists of the following steps: It is also extremely dry (technically a desert), averaging 166 mm (6.5 in) of precipitation per year.-1.71Llama-3-8B In the middle, we list the weights of a linear surrogate model that estimates the logit-scaled probability of the response as a function of the context ablation vector (3); CONTEXTCITE casts these weights as attribution scores. On the right, we plot the surrogate model's predictions against the actual logitscaled probabilities for random context ablations. Two sources appear to be primarily responsible for the response, resulting in four "clusters" corresponding to whether each of these sources is included or excluded. These sources appear to interact linearly-the effect of removing both sources is close to the sum of the effects of removing each source individually. As a result, the linear surrogate model faithfully captures the language model's behavior.Step 1: Sample a "training dataset" of ablation vectors v 1 , . . . , v n and compute f (v i ) for each v i .Step 2: Learn a surrogate model f : {0, 1} d → R that approximates f by training on the pairs(v i , f (v i )).Step 3: Attribute the behavior of the surrogate model f to individual sources.For the surrogate model f to be useful, it should (1) faithfully model f , (2) be efficient to compute, and (3) yield scores attributing its outputs to the individual sources. To satisfy these desiderata, we find the following design choices to be effective:• Predict logit-scaled probabilities: Fitting a regression model to predict probabilities directly might be problematic because probabilities are bounded in [0, 1]. The logit function (σ -1 (p) = log p 1-p ) is a mapping from [0, 1] to (-∞, ∞), making logit-probability a more natural target for regression.• Learn a linear surrogate model: Despite their simplicity, we find that linear surrogate models are often quite faithful. With a linear surrogate model, each weight signifies the effect of ablating a source on the output. As a result, we can directly cast the weights of the surrogate model as attribution scores. We illustrate an example depicting the effectiveness of a linear surrogate model in Figure 2 and provide additional randomly sampled examples in Appendix B.2.• Learn a sparse linear surrogate model: Empirically, we find that a generated statement can often be explained well by just a handful of sources. In particular, Figure 3a shows that the number of sources that are "relevant" to a particular generated statement is often small, even when the context comprises many sources. Motivated by this observation, we induce sparsity in the surrogate model via LASSO [27] . As we illustrate in Figure 3b , this enables learning a faithful linear surrogate model even with a small number of ablations. For example, the surrogate model in Figure 2 uses just 32 ablations even though the context comprises 98 sources (in this case, sentences).• Sample ablation vectors uniformly: To create the surrogate model's training dataset, we sample ablation vectors uniformly from the set of possible subsets of context sources.We summarize the resulting method, CONTEXTCITE, in Algorithm 1. See Figure 2 [28] , a summarization task, and Natural Questions [29], a question answering task, we observe that the number of sources that are "relevant" for a particular statement generated by Llama-3-8B [22] is small, even when the context comprises many sources (Figure 3a ). Therefore, inducing sparsity via LASSO yields an accurate surrogate model with just a few ablations (Figure 3b ). See Appendix A.4 for the exact setup. Output: Attribution scores ŵ ∈ R d 3: f (v) := p LM (R | ABLATE(C, v), Q) ▷ Probability of R when ablating C according to v 4: g(v) := σ -1 (f (v))▷ Logit-scaled version of f 5: for i ∈ {1, . . . , t} do 6: Sample a random ablation vector v i uniformly from {0, 1} d 7: Baselines. We consider three natural baselines adapted from prior work on model explanations. We defer details and additional baselines that we found to be less effective to Appendix A.5.1.y i ← g(v i ) 8: end for 9: ŵ, b ← LASSO({(v i , y i )} n i=1 ,1. Leave-one-out: We consider a leave-one-out baseline that ablates each source individually and compute the log-probability drop of the response as an attribution score. Leave-one-out is an oracle for the top-k log-probability drop metric (1) when k = 1, but may be prohibitively expensive because it requires an inference pass for every source.2. Attention: A line of work on explaining language models leverages attention weights [33] [34] [35] [36] [37] [38] . We use a simple but effective baseline that computes an attribution score for each source by summing the average attention weight of individual tokens in the source across all heads in all layers.3. Gradient norm: Other explanation methods rely on input gradients [39] [40] [41] . Here, following Yin and Neubig [42], we estimate the attribution score of each source by computing the ℓ 1 -norm of the log-probability gradient of the response with respect to the embeddings of tokens in the source.4. Semantic similarity: Finally, we consider attributions based on semantic similarity. We employ a pre-trained sentence embedding model [43] to embed each source and the generated statement.We treat the cosine similarities between these embeddings as attribution scores.Experiment setup. Each example on which we evaluate consists of a context, a query, a language model, and a generated response. As discussed in Section 2.3, rather than attributing the entire response to the context, we consider attributing individual statements in the response to the context. Specifically, given an example, we (1) split the response into sentences using an off-the-shelf tokenizer [44] , and (2) compute attribution scores for each sentence. Then, to evaluate the attribution scores, we measure the top-k log-probability drop for k = {1, 3, 5} (1) and LDS (2) for each sentence separately, and then average performances across sentences. Our experiments perform this evaluation for every combination of context attribution method, dataset, and language model. We evaluate CONTEXTCITE with {32, 64, 128, 256} context ablations.Results. In Figure 4 , we find that CONTEXTCITE consistently outperforms baselines, even when we only use 32 context ablations to compute its surrogate model. While the attention baseline approaches the performance of CONTEXTCITE with Llama-3-8B, it fares quite poorly with Phi-3-mini suggesting that attention is not consistently reliable for context attribution. CONTEXTCITE also attains high LDS across benchmarks and models, indicating that its attributions accurately predict the effects of ablating sources.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA] Justification: We do not have theoretical results. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
gradient times input:: 
 As an additional gradient-based baseline, we also consider taking the dot product of the gradients and the embeddings following Shrikumar et al. [68] in Appendix B.3, but found this to be less effective than the gradient norm.5. Semantic similarity: Finally, we consider attributions based on semantic similarity. We employ a pre-trained sentence embedding model [43] to embed each source and the generated statement.We treat the cosine similarities between these as attribution scores.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: Yes, see appendices. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.• Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
applications of contextcite: 
 In Section 4, we found that CONTEXTCITE is an effective (contributive) context attribution method.In other words, it identifies the sources in the context that cause the model to generate a particular statement. In this section, we present three applications of context attribution: helping verify generated statements (Section 5.1), improving response quality by pruning the context (Section 5.2), and detecting poisoning attacks (Section 5.3).It can be difficult to know when to trust statements generated by language models [45] [46] [47] [48] [49] . In this section, we investigate whether CONTEXTCITE can help language models verify the accuracy of their own generated statements.Approach. Our approach builds on the following intuition: if the sources identified by CONTEXTCITE for a particular statement do not support it, then the statement might be inaccurate. To operationalize this, we (1) use CONTEXTCITE to identify the top-k most relevant sources and (2) provide the same language model with these sources and ask it if we can conclude that the statement is correct. We treat the model's probability of answering "yes" as a verification score.Experiments. We apply our verification pipeline to answers generated by Llama-3-8B for 1, 000 random examples from each of two question answering datasets: HotpotQA [31] and Natural Questions [29] . We provide the language model with the top-k most relevant sources (for a few different values of k) and measure its AUC for predicting whether its generated answer is accurate.As a baseline, we provide the model with the entire context and measure this AUC in the same manner.In Figure 5 , we observe that the verification scores obtained using the top-k sources are substantially higher than those obtained from using the entire context. This suggests that context attribution can be 4a ) and linear datamodeling score (Figure 4b ) of CONTEXTCITE and baselines. We evaluate attributions of responses generated by Llama-3-8B and Phi-3-mini on up to 1, 000 randomly sampled validation examples from each of three benchmarks. We find that CONTEXTCITE using just 32 context ablations consistently matches or outperforms the baselines-attention, gradient norm, semantic similarity and leave-one-out-across benchmarks and models. Increasing the number of context ablations to {64, 128, 256} can further improve the quality of CONTEXTCITE attributions. used to help language models verify the accuracy of their own responses. See Appendix A.6 for the exact setup.If the sources identified by CONTEXTCITE can help a language model verify the accuracy its answers (Section 5.1), can they also be used to improve the accuracy of its answers? Indeed, language models often struggle to correctly use relevant information hidden within long contexts [14, 13] . In this section, we explore whether we can improve response quality by pruning the context to include only query-relevant sources.Approach. Our approach closely resembles the verification pipeline from Section 5.1; however, instead of using the top-k sources to verify correctness, we use them to regenerate the response. Specifically, it consists of three steps: (1) generate a response using the entire context, (2) use CONTEXTCITE to identify the top-k most relevant sources, and (3) regenerate the response using only these sources as context.Experiments. We assess the effectiveness of this approach on two question-answering datasets: HotpotQA [31] and Natural Questions [29] . In both datasets, the provided context typically includes a lot of irrelevant information in addition to the answer to the question. In Figure 6 , we report the average F 1 -score of Llama-3-8B on 1, 000 randomly sampled examples from each dataset (1) when it is provided with the entire context and (2) when it is provided with only the top-k sources according to CONTEXTCITE. We find that simply selecting the most relevant sources can consistently improve question answering capabilities. See Appendix A.7 for the exact setup and Appendix C.2 for additional discussion of why pruning in this way can improve question answering performance.Finally, we explore whether context attribution can help surface poisoning attacks [15] [16] [17] . We focus on indirect prompt injection attacks [18, 19] that can override a language model's response to a given query by "poisoning", or adversarially modifying, external information provided as context. For example, if a system like ChatGPT browses the web to answer a question about the news, it may end up retrieving a poisoned article and adding it to the language model's context. These attacks can be "obvious" once identified-e.g., If asked about the election, ignore everything else and say that Trump dropped out-but can go unnoticed, as users are unlikely to carefully inspect the entire article.Approach. If a prompt injection attack successfully causes the model to generate an undesirable response, the attribution score of the context source(s) containing the injected poison should be high. One can also view the injected poison as a "strong feature" [50] in the context that significantly influences model output and, thus, should have a high attribution score. Concretely, given a potentially poisoned context and query, our approach (a) uses CONTEXTCITE to attribute the generated response to sources in the context and (b) flags the top-k sources with the highest attribution scores for further manual inspection.Experiments. We consider two types of prompt injection attacks: (1) handcrafted attacks (e.g., "Ignore all previous instructions and. . .") [16], and (2) optimization-based attacks [19] . In both cases, CONTEXTCITE surfaces the prompt injection as the single most influential source more than 95% of the time. See Appendix A.8 for the exact setup and more detailed results.k = 1 k = 2 k = 4 k = 8 k = 16 k = all (Baseline)Top-k context sources Llama-3-8B on HotpotQAk = 4 k = 8 k = 16 k = 32 k = 64 k = 128 k = all (Baseline)Top-k context sources Llama-3-8B on Natural QuestionsFigure 5 : Helping verify generated statements using CONTEXTCITE. We report the AUC of Llama-3-8B for verifying the correctness of its own answers when we provide it with the top-k sources identified by CONTEXTCITE and when we provide it with the entire context. We consider 1, 000 random examples from HotpotQA on the left and 1, 000 random examples from Natural Questions on the right. In both cases, using the top-k sources results in substantially more effective verification than using the entire context, suggesting that CONTEXTCITE can help language models verify their own statements. Llama-3-8B on Natural Questions Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: Our code is open-source (see abstract). We use publicly available datasets and models.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/public/ guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so no is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark). • The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/ guides/CodeSubmissionPolicy) for more details. • The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc. • The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable). • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
related work: 
 Citations for language models. Prior work on citations for language models has focused on teaching models to generate citations for their responses [1, 4, 2, 3, 5, 51, 52] . For example, Menick et al.[2] fine-tune a pre-trained language model to include citations to retrieved documents as part of its response. Gao et al.[5] use prompting and in-context demonstrations to do the same. Post-hoc methods for citation [4, 51] attribute existing responses by using an auxiliary language model to identify relevant sources. Broadly, existing methods for generating citations are intended to be corroborative [20] in nature; citations are evaluated on whether they support or imply a generated statement [53, 6, 7, 54] . In contrast, CONTEXTCITE-a contributive attribution method-identifies sources that cause a language model to generate a given response.Explaining language model behavior. Related to context attribution is the (more general) problem of explaining language model behavior. Methods for explaining language models have used attention weights [37, 38] Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results? Answer: [Yes] Justification: Yes, see appendices Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 7. Experiment Statistical Significance Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: Our empirical findings account for statistical significance. Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
conclusion: 
 We introduce the problem of context attribution whose goal is to trace a statement generated by a language model back to the specific parts of the context that caused the model to generate it. Our proposed method, CONTEXTCITE, leverages linear surrogate modeling to accurately attribute statements generated by any language model in a scalable manner. Finally, we present three applications of CONTEXTCITE: (1) helping verify generated statements (2) improving response quality by pruning the context and (3) detecting poisoning attacks. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] Justification: Yes, see appendices. Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 9. Code Of Ethics Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] Justification: Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: Yes, see appendices.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: We do not release data or models that have a high risk for misuse. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: Yes, we cite all the datasets and pre-trained models used in the paper.• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We run all experiments on a cluster of A100 GPUs. We use the scikit-learn [64] implementation of LASSO for CONTEXTCITE, always with the regularization parameter alpha set to 0.01. When splitting the context into sources or splitting a response into statements, we use the off-the-shelf sentence tokenizer from the nltk library [44] . Our implementation of CONTEXTCITE is available at https://github.com/MadryLab/context-cite.The language models we consider in this work are Llama-3-{8/70}B [22], Mistral-7B [65] and Phi-3-mini [32] . We use instruction-tuned variants of these models. We use the implementations of language models from HuggingFace's transformers library [66] . Specifically, we use the following models:• Llama-3-{8/70}B: meta-llama/Meta-Llama-3-{8/70}B-Instruct • Mistral-7B: mistralai/Mistral-7B-Instruct-v0.2 • Phi-3-mini: microsoft/Phi-3-mini-128k-instructWhen generating responses with these models, we use their standard chat templates, treating the prompt formed from the context and query as a user's message.We consider a variety of datasets to evaluate CONTEXTCITE spanning question answering and summarization tasks and different context structures and lengths. We provide details about these datasets and preprocessing steps in this section. Some of the datasets, namely Natural Questions and TyDi QA, contain contexts that are longer than the maximum context window of the models we consider. In particular, Llama-3-8B has the shortest context window of 8, 192 tokens. When evaluating, we filter datasets to include only examples that fit within this context window (with a padding of 512 tokens for the response).CNN DailyMail [28] is a news summarization dataset. The contexts consists of a news article and the query asks the language model to briefly summarize the articles in up to three sentences. We use the following prompt template:Query: Please summarize the article in up to three sentences.Hotpot QA.[31] is a multi-hop question-answering dataset in which the context consists of multiple short documents. Answering the question requires combining information from a subset of these documents-the rest are "distractors" containing information that is only seemingly relevant. We use the following prompt template: A.3.1 Dataset statistics.In Table 1 , we provide the average and maximum numbers of sources in the datasets that we consider. In this section, we discuss how we partition contexts into sources and perform context ablations. For every dataset besides Hotpot QA, we use an off-the-shelf sentence tokenizer [44] to partition the context into sentences. To perform a context ablation, we concatenate all of the included sentences and provide the resulting string to the language as context. The Hotpot QA context consists of multiple documents, each of which includes annotations for individual sentences. Furthermore, the documents have titles, which we include in the prompt (see Appendix A.3). Here, we still treat sentences as sources and include the title of a document as part of the prompt if at least one of the sentences of this document is included.A.4 Learning a sparse linear surrogate modelIn Figure 3 , we illustrate that CONTEXTCITE can learn a faithful surrogate model with a small number of ablations by exploiting underlying sparsity. Specifically, we consider CNN DailyMail and Natural Questions. For 1, 000 randomly sampled validation examples for each dataset, we generate a response with Llama-3-8B using the prompt templates in Appendix A.3. Following the discussion in Section 2.3, we split each response into sentences and consider each of these sentences to be a "statement." For the experiment in Figure 3a , for each statement, we ablate each of the sources individually and consider the source to be relevant if this ablation changes the probability of the statement by a factor of at least δ = 2. For the experiment in Figure 3b , we report the average root mean squared error (RMSE) over these statements for surrogate models trained using different numbers of context ablations. See Appendices A.2 and A.3 for additional details on datasets and models.See Appendices A.1 to A.3 for details on implementation, datasets and models for our evaluations.We provide a detailed list of baselines for context attribution in this section. In addition to the baselines described in Section 4, we consider additional attention-based and gradient-based baselines.We provide evaluation results including these baselines in Appendix B.3.1. Average attention: We compute average attention weights across heads and layers of the model. We compute the sum of these average weights between every token of a source and every token of the generated statement to attribute as an attribution score. This is the attention-based baseline that we present in Figure 4 .In Section 5.1, we explore whether CONTEXTCITE can help language models verify the accuracy of their own generated statements. Specifically, we first use CONTEXTCITE to identify a set of the top-k most relevant sources. We then ask the language model whether we can conclude that the statement is accurate based on these sources. The following are additional details for this experiment:1. Datasets and models. We evaluate this approach on two question-answering datasets: HotpotQA [31] and Natural Questions [29] . For each of these datasets, we evaluate the F 1 score of instructiontuned Llama-3-8B (Figure 6 ) on 1, 000 randomly sampled examples from the validation set.2. Question answering prompt. We modify the prompts outlined for HotpotQA and Natural Questions in Appendix A.3 to request the answer as a short phrase or sentence. This allows us to assess the correctness of the generated answer.<Original prompt>Please answer with a single word or phrase when possible.If the question cannot be answered from the context, say so instead.3. Applying CONTEXTCITE. We compute CONTEXTCITE attributions using 256 calls to the language model.Extracting the top-k most relevant sources. Given the CONTEXTCITE attributions for a context and generated statement, we extract the top-k most relevant sources to verify the generated statement.In this case, sources are sentences. For Hotpot QA, in which the context consists of many short documents, we extract each of the documents containing any of the top-k sentences to provide the language model with a more complete context. For Natural Questions, we simply extract the top-k sentences.5. Verification prompts. To verify the generated answer using the language model and the top-k sources, we first convert the model's answer to the question (which is a word or short phrase) into a self-contained statement. We do so by prompting the language model to combine the question and its answer into a self-contained statement, using the following prompt:Please merge the following question and answer into a single statement. For example, if the question is "What is the capital of France?" and the answer is "Paris", you should say: "The capital of France is Paris. Question: {question} Answer: {answer}We then use the following prompt to ask the language model whether the statement is accurate:Context: {pruned_context}Can we conclude that "{self_contained_answer}"? Please respond with just yes or no.Recall that in Section 5.2, we use CONTEXTCITE to improve the question-answering capabilities of language models by extracting the most query-relevant sources from the context. We do so in three steps: (1) generate a response using the entire context, (2) use CONTEXTCITE to compute attribution scores for sources in the context, and (3) construct a query-specific context using only the top-k sources, which can be used to regenerate a response. The implementation details for constructing the query-specific context are the same as for the verification application outlined in Appendix A.6.In Section 5.3, we consider four different attack setups, which we describe below. We design prompt injections with varied goals: false refusal of queries, misinformation, malicious code execution, change of language for the response, etc. Because this process is laborious and time-consuming, we provide a small dataset consisting of twenty context-query pairs. We provide this dataset in our code release.Qualitatively, one case where CONTEXTCITE fails to surface the prompt injection as the highestscoring source (although the prompt injection is still within the top-3 scores) is when the prompt injection makes a subtle change to the output. For example:Context: "context": "Yukio Akakariyama (Japanese \u8d64\u72e9\u5c71 Response: Yukio Akakariyama is a renowned Japanese professional pool player, particularly famous for his achievements in the 10-ball pool discipline....Here Phi-3-mini's response still heavily draws on the original response, but adds the incorrect 10-ball reference.Optimization-based attacks on Phi-3-mini. We also use the GCG attack introduced in Zou et al. [17] . In this setup, we again consider Wikipedia articles as contexts. Here, instead of focusing on question-answering, we turn our attention to summarization. In particular, the query for each of the context articles is Query: Briefly summarize this article.We then sample a random place within the context article and insert a twenty-character placeholder, which we then optimize with GCG to maximize the likelihood of the model outputtingResponse: I am sorry, I cannot assist with that. 
Given the long contexts, as well as the fact that we insert the adversarial tokens in the middle of the context (and not as a suffix), we observe a very low success rate of these optimization-based attacks. In particular, we report a success rate of just 2%. We then filter only the prompts containing a successful attack, and construct a dataset, which we provide in our code release. Due to the high computational cost of the GCG attack (as well as the low success rate), this dataset is also small in size (22 samples, filtered down from 1000 GCG attempts, each on a random Wikipedia article).Qualitatively, CONTEXTCITE fails to surface the GCG-optimized sentence as the one with the highest attribution score when the attack is not fully successful. For example, rather than outputting the target response, for one of the contexts, Phi-3-mini instead generates Python code to give a summary of the article: We found another failure mode to be noteworthy as well. When using the Wikipedia article about Tupper Lake in New York, CONTEXTCITE finds the sentence Roger Allen LaPorte, Vietnam War protester, immolated himself in front of the United Nations building.Response: '''as the main source leading Phi-3-mini to refuse to summarize the article. Indeed, the model refuses to discuss this sensitive topic even without the GCG-optimized prompt.Optimization-based attacks on Llama3-8B. Finally, we mount the prompt injections attack NeuralExec developed by Pasquini et al. [19] . In short, the attack consists of generating a universal optimized prompt injection which surrounds a "payload" message. The goal of the optimized prompt injection is to maximize the likelihood of the payload message being picked up by the model. One can view the NeuralExec attack as an optimization-based counterpart to the handcrafted attacks we consider [16] .For Llama3-8B, the universal (i.e., independent of the context) prompt injection is OtherFully":[{"<<"View361 digestion.DOCUMENT@js tuberculosis_framework);</ annon interle[PAYLOAD]? [[.There explic <!-Votingusterity where [PAYLOAD] is a placeholder for the "payload" message. We use the test set of the NeuralExec paper to evaluate how well CONTEXTCITE can detect the presence of this prompt injection. The NeuralExec attack is successfully mounted on 91 of the 100 test samples. CONTEXTCITE is able to surface the prompt injection as the most influential source in 90 out of these 91 cases, leading to a (top-1) detection accuracy of 98.9%.In Table 2 , we report aggregated results for all attacks on all LLMs. On the right side of Figure 2 , we show the actual logit-probabilities of different context ablations as well as the logit-probabilities predicted by a linear surrogate model. In that example, the linear surrogate model is quite faithful. In this section, we provide additional randomly sampled examples from CNN DailyMail (see Figure 7 ), Natural Questions (see Figure 8 ), and TyDi QA (see Figure 9 ). We use 256 context ablations to train the surrogate model, and observe that a linear surrogate model is broadly faithful across these benchmarks. Figure 10 : Evaluating CONTEXTCITE on additional models and benchmarks using the top-k log-probability drop metric (1). We compare CONTEXTCITE to additional baselines (ℓ 2 -gradient norm, gradient-times-input, and attention rollout) on three models (Llama-3-8B, Phi-3-mini, Mistral-7B) and two additional benchmarks (TyDi QA and MS-MARCO). Each row corresponds to a different benchmark and each column corresponds to a different model. Across all benchmarks and models, CONTEXTCITE (with just 32 calls) consistently outperforms the baselines on the top-k log-probability drop metric, which measures the effect of ablating the top-k context sources with the highest attribution scores. Similar to our results in Figure 4a , increasing the number of context ablations to {64, 128, 256} can further improve the quality of CONTEXTCITE attributions in this setting as well. Figure 11 : Evaluating CONTEXTCITE on additional models and benchmarks using the linear datamodeling score (2). Like in Figure 10 , we compare CONTEXTCITE to additional baselines (ℓ 2 -gradient norm, gradient-times-input, and attention rollout) on three models (Llama-3-8B, Phi-3-mini, Mistral-7B) and two additional benchmarks (TyDi QA and MS-MARCO). Each row corresponds to a different benchmark and each column corresponds to a different model. Across all benchmarks and models, CONTEXTCITE (with just 32 calls) consistently outperforms the baselines on the linear datamodeling score, which quantifies the extent to which context attributions predict the effect of ablating the context sources on the model response. Similar to our results in Figure 4b , increasing the number of context ablations to {64, 128, 256} further improves the quality of CON-TEXTCITE attributions in this setting as well.Our evaluation suite for CONTEXTCITE in Section 4 consists of models with up to 8 billion parameters.In this section, we conduct a more limited evaluation of CONTEXTCITE for a larger model, Llama-3-70B [22] . We find that CONTEXTCITE is effective even at this larger scale.In Figure Figure 12 , we evaluate CONTEXTCITE for Llama-3-70B on the CNN DailyMail and Hotpot QA benchmarks using the top-k log-probability drop metric (1) and the linear datamodeling score (2). We use the same evaluation setup as in Section 4, but use a subset of the baselines and only use 32 context ablations for CONTEXTCITE due to computational cost. We find that CONTEXTCITE consistently outperforms baselines. Figure 12 : Evaluating word-level context attributions. We report the top-k log-probability drop (Figure 12a ) and linear datamodeling score (Figure 12b ) of CONTEXTCITE and baselines. We evaluate attributions of responses generated by Llama-3-70B on 1, 000 randomly sampled validation examples from each of CNN DailyMail and Hotpot QA.In this section, we provide CONTEXTCITE attributions for Llama-3-70B for randomly selected examples. For each example, we randomly select a sentence from the response to attribute and display the 4 sources with the highest attribution scores. Response (attributed statement is highlighted)Paul Bell, a director of online gambling firm 666Bet, has been arrested as part of a £21m fraud probe, along with six others, and £1m in cash was seized. The firm, which has deals with West Bromwich Albion and Leicester City, had its licence suspended by the UK gambling regulator, and its website is currently offline. Harry Redknapp, who appeared in the firm's adverts, claims to know nothing about the investigation.Score world.' The firm's website is currently offline. Its official Twitter account said the site is under maintenance 'due to unforeseen circumstances'. ...In this work, we primarily focus on sentences on sources for context attribution. In this section, we briefly explore using CONTEXTCITE to perform context attribution with individual words as sources on the DROP benchmark [69] . We find that CONTEXTCITE can provide effective word-level attributions, but may require a larger number of context ablations.In Figure 13 , we evaluate word-level CONTEXTCITE on the DROP benchmark using the top-k log-probability drop metric (1) and the linear datamodeling score (2). We use the same evaluation setup as in Section 4. While CONTEXTCITE matches or outperforms baselines, we find that it attains lower absolute values for the linear datamodeling score. This may be because word-level attributions are less sparse: a given generated statement may depend on many individual words within the context. It may also be because there are much stronger dependencies between words than between sentences, rendering a linear surrogate model less faithful. Figure 13 : Evaluating word-level context attributions. We report the top-k log-probability drop (Figure 13a ) and linear datamodeling score (Figure 13b ) of CONTEXTCITE and baselines. We evaluate attributions of responses generated by Llama-3-8B on 1, 000 randomly sampled validation examples from the DROP benchmark.k = 1 k = 3 k = 5In this section, we provide word-level CONTEXTCITE attributions for Llama-3-8B for randomly selected examples. For each example, we randomly select a sentence from the response to attribute and display the 4 sources with the highest attribution scores.Example #1 of word-level of CONTEXTCITE attributions for Llama-3-8B on DROP Query How many field goals were between 20 and 30 yards?Response (attributed statement is highlighted)According to the text, there were three field goals between 20 and 30 yards:1 To estimate SHAP values, Lundberg and Lee [9] propose Kernel SHAP, a method that uses LIME with a specific choice of similarity kernel that yields SHAP values. Specifically, in order for LIME to estimate SHAP values, they show that the similarity kernel for an ablation vector v should beπ SHAP (v) = d -1 d |v| • |v| • (d -|v|)where d is the number of features and |v| is the number of non-zero elements of the ablation vector v.Using the same setup as in Appendix B.3, we compare the Kernel SHAP estimator (which uses LASSO with samples weighted according to π SHAP ) to the CONTEXTCITE estimator (which uses LASSO with In the in-context learning setting, "training examples" are provided to a model as context before it is queried with a test example. Datamodels have previously been used to study in-context learning [23, 24] . If one thinks of in-context learning as sources, this form of training data attribution is a special case of context attribution. In Section 5.2, we show that providing only the top-k most relevant CONTEXTCITE sources for a language model's original answer to a question can improve the quality of its answer. We would like to note that the sources identified by CONTEXTCITE are those that were used to generate the original response. If the original response is incorrect, it may be surprising that providing only the sources that led to this response can improve the quality of the response.To explain why pruning the context does improve question answering performance, we consider two failure modes associated with answering questions using long contexts:1. The model identifies the wrong sources for the question and answers incorrectly.2. The model identifies the correct sources for the question but misinterprets information because it is distracted by other irrelevant information in the context.Intuitively, pruning the context to include only the originally identified sources can help mitigate the second failure mode but not the former. The fact that pruning the context in this way can improve question answering performance suggests that the second failure mode occurs and that mitigating it can thus improve performance.Most of the computational cost of CONTEXTCITE comes from creating the surrogate model's training dataset. Hence, the efficiency of CONTEXTCITE depends on how many ablations it requires to learn a faithful surrogate model. We find that CONTEXTCITE requires just a small number of context ablations to learn a faithful surrogate model-in our experiments, 32 context ablations suffice. Thus, attributing responses using CONTEXTCITE is 32× more expensive than generating the original response. We note that the inference passes for each of these context ablations can be fully parallelized. Furthermore, because CONTEXTCITE is a post-hoc method that can be applied to any existing response, a user could decide when they would like to pay the additional computational cost of CONTEXTCITE to obtain attributions. When we use CONTEXTCITE to attribute multiple statements in the response, we use the same context ablations and inference calls. In other words, there is a fixed cost to attribute (any part of) a generated response, after which it is very cheap to attribute specific statements.C.3.1 Why do we only need a small number of ablations?We provide a brief justification for why 32 context ablations suffice, even when the context comprises many sources. Since we are solving a linear regression problem, one might expect the number of ablations needed to scale linearly with the number of sources. However; in our sparse linear regression setting, we have full control over the covariates (i.e., the context ablations). In particular, we ablate sources in the context independently and each with probability 1/2. This makes the resulting regression problem "well-behaved." Specifically, this lets us leverage a known result (see Theorems 7.16 and 7.20 of Wainwright [72] ) which tells us that we only need O(k log(d)) context ablations, where d is the total number of sources and k is the number of sources with non-zero relevance to the response. In other words, the number of context ablations we need grows very slowly with the total number of sources. It only grows linearly with the number of sources that the model relies on when generating a particular statement. As we show empirically in Figure 3a , this number of sources is often small.In this section, we discuss a few limitations of CONTEXTCITE.Potential failure modes. Although we find a linear surrogate model to often be faithful empirically (see Figure 2 , Appendix B.2), this may not always be the case. In particular, we hypothesize that the linearity assumption may cease to hold when many sources contain the same information. In this case, a model's response would only be affected by excluding every one of these sources. In practice, to verify the faithfulness of the surrogate model, a user of CONTEXTCITE could hold out a few context ablations to evaluate the surrogate model (e.g., by measuring the LDS). They could then assess whether CONTEXTCITE attributions should be trusted.Another potential failure mode of CONTEXTCITE is attributing generated statements that follow from previous statements. Consider the generated response: "He was born in 1990. He is 34 years old." with context mentioning a person born in 1990. If we attribute the statement "He was born in 1990." we would likely find the relevant part of the context. However, if we attribute the statement "He is 34 years old." we might not identify any attributed sources, despite this statement being grounded in the context. This is because this statement is conditioned on the previous statement. Thus, in this case there is an "indirect" attribution to the context through a preceding statement that would not be identified by the current implementation of CONTEXTCITE.Unintuitive behaviors. A potentially unintuitive behavior of CONTEXTCITE is that it can yield a low attribution score even for a source that supports a statement. This is because CONTEXTCITE provides contributive attributions. Hence, if a language model already knows a piece of information from pre-training and does not rely on the context, CONTEXTCITE would not identify sources. This may lead to unintuitive behaviors for users.Validity of context ablations. In this work, we primarily consider sentences as sources for context attribution and perform context ablations by simply removing these sentences. One potential problem with this type of ablation is dependencies between sentences. For example, consider the sentences: "John lives in Boston. Charlie lives in New York. He sometimes visits San Francisco." In this case, "He" refers to Charlie. However, if we ablate just the sentence about Charlie, "He" will now refer to "John." There may be other ablation methods that more cleanly remove information without changing the meaning of sources because of dependencies.Computational efficiency. As previously discussed, attributing responses using CONTEXTCITE is 32× more expensive than generating the original response. This may be prohibitively expensive for some applications.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.In practice, we may include additional tokens, e.g., to specify the beginning and end of a user's message.We also explore using individual words as sources in Appendix B.5.This is a design choice; we could also, for example, replace excluded sources with a placeholder. 


Decision: This paper investigates attribution in generation. The main contributions of this paper include (1) initially introducing the task of context attribution; (2) proposing the simple and scalable method for context attribution, ContextCite; (3) showcasing two applications of context attribution, respectively verifying factuality and selecting query-relevant information to improve response quality.

This paper has several strengths. First, as all reviewers have agreed, this paper is well-structured and easy to follow. Second, this paper formalizes context attribution for language models, i.e., identifying which parts of the input resulted in generation of a particular output. In addition, reviewers (Thhj, A2Ps, and 4MHt) agreed that the proposed ContextCite is simple and intuitive, with superior performance over baselines. The experiments designed for method verification are comprehensive and convincible, as reviewers Zvan and A2Ps mentioned.

Several weaknesses were identified by reviewers. (1) There are some concerns regarding experiments. For example, reviewer A2Ps mentioned exploration of only limited number of tasks, LLMs, and prompts which needs to be extended; reviewer Zvan mentioned concerns regarding ablations. (2) The method seems indeed costly to run. (3) As reviewers pointed out, several in-depth aspects remain to be further analyzed, including adding theoretical justification, and discussion between context and data attribution.

During the rebuttal period, most of the major concerns seem to be addressed, including adding new experiment results in response to reviewers’ questions and discussing each argument in depth.

Reviews: review: paper_summary: paper_summary: The authors of 17808 formalizes the problem of context attribution for language models. That is, identifying which parts of the input context caused a model to generate a particular output. The authors propose ContextCite, a method that learns a sparse linear surrogate model to approximate how ablating different parts of the context can impact the model's output probability. Authors claim several contributions, including 1) Formalizing the official task of context attribution and evaluation metrics, 2) Developing ContextCite, a simple and scalable attribution method, 3) Demonstrating ContextCite outperforms baselines on attribution quality, and 4) Showing applications in verifying generated statements and improving response quality.
strength_weakness: strength_weakness: - As noted in Appendix C.2, the current ablation approach of removing sentences can break dependencies between sents. This limitation could be explored a bit further in the paper, but it's understandable if there's time constraint. 

- A little more theoretical justification for why the linear surrogate model works well would be better.
questions: questions: - While the efficiency and simplicity of linear surrogate models are acknowledged, did you try more sophisticated surrogate models beyond linear regression in your early exps? If so, how did they look like?
limitations: limitations: The authors provide a good discussion of limitations in Appendix C.2.
paper_summary: paper_summary: In an attempt to understand how language models leverage context information in its generation, this work studies contributive context attribution. Following the recent trend in attribution research, authors propose to use (sparse) linear models to learn the importance of each unit/sentence in the context, and demonstrate its effectiveness using top-k log-probability drop and linear datamodeling score (LDS) tests. Lastly, authors explore two potential applications of contributive context attribution, namely generated statement verification and query-relevant information selection.
strength_weakness: strength_weakness: - While (contributive) context attribution looks to be new on the surface, their problem setup, method (i.e. ContextCite), and evaluation are mostly borrowed from existing data attribution literature. While authors provided proper references to those literature, I still believe the technical contribution largely lacks novelty. I wasn't able to grasp whether authors introduced any major modifications to adapt existing data attribution methods (i.e. datamodel) and evaluations (i.e. brittleness and LDS) to context attribution. If I missed something, please correct me.
- If my understanding is correct, they prompted GPT-4 to generate the (proxy) gold label for their generated statement verification experiments (section 5.1). This was confusing because the authors noted in line 84-85 that prompting the model for citations is generally considered as corroborative attribution. To be precise, the authors asked for the correctness of each sentence instead of for citations here, but it was a bit odd for me to use the prompting technique to evaluate contributive context attribution.
questions: questions: - How does the context length affect the number of required forward passes with randomly sampled context information? Naively thinking, if you have n parameters in your linear model, you may need n samples to properly fit the linear model (without a sparsity assumption).
limitations: limitations: The authors adequately addressed the limitations in Appendix C.
paper_summary: paper_summary: The current paper introduces the task of context attribution, which aims at attributing a generated response of an auto-regressive LM back to the sentences in the input contexts. Target at this task, ContextCite, was proposed to predict which piece of context changes the probability of the generated response most. Together with ContextCite, this paper also designed two evaluation metrics based on the probabilities of the generated responses. At length, the experiments on two small LLMs show the effectiveness of ContextCite. Additionally, this paper provides several examples of applications of ContextCite.
strength_weakness: strength_weakness: Though I thought this paper was already of good quality, I still have three major concerns regarding the task, the model and the evaluation. 

First and foremost, in my opinion, the current definition of the task is still highly limited and can be largely extended. For example, the current task only focuses on the attributions with responses to the whole generated responses, but for tasks like summarisation (which was also examined in this work), it is also important to attribute each token in the response to, for instance, route the hallucination. I am personally very OK if the solution of such an enhanced task is not provided, but I am very inclined to see that the related discussions about the possible extensions of the task can be included.

Second, regarding ContextCite, though the idea of making use of the probability change looks promising, after due consideration, I have a strong feeling that its relation with which parts in the input contribute more to the response is not fully deterministic. In my opinion, the probability of a response could be reduced more if, as pointed out in the paper, the omitted piece contributes more to the response or only if the resulting context after removing the sentence is way more incoherent. One example of such a possibility is given in Appendix C2, but the discussion seems not to be sufficient. This also made me think that the model only suits tasks whose output relies on the very long input text (e.g., summarisation and document-based QA, as omitting one sentence would not highly influence the coherency).

Finally, the evaluation was done on three very related tasks and two tiny language models. This makes me somewhat question whether the solution is generalizable to other NLP tasks (that are open-ended, for example), other larger models and other prompt designs (e.g., in-context learning and chain-of-thought).
questions: questions: Can you provide some examples of the selected sentences and hit on the distributions and characteristics of the selected sentences?
limitations: limitations: See my points in the weakness section.
paper_summary: paper_summary: This paper proposes a simple and effective approach for attributing LLM generation to the sources in the context, where a set of random ablations and their corresponding effects in the model's probability of generating the response are modeled with a (sparse) linear relationship. The method archives impressive performance and also allows applications such as verifying model generations or denoising the sources in the context.
strength_weakness: strength_weakness: - The method may be a bit costly to run, given that the linear relationship needs to be re-learned for every (context, generation) instance, and the number of sample ablations also grows when the context scales.
- If I'm not mistaken, the proposed linear relationship implicitly assumes that the sources in the context contribute in a rather independent manner to the model's generation. This may not work well when the generation involves complex interactions and reasoning over a large number of sources in the context (beyond just a few sentences as in, e.g., HotpotQA). For example, if the task involves some kind of logical deduction, missing a certain premise in the context could cause a snow-balling effect on how the model processes the remaining sources and scores its generation, which is very hard to model with a linear relation.
questions: questions: For extracting query-relevant information from the context (Section 5.2), if some important sources are "lost in the middle" and overlooked by the model, then the proposed approach would give a very low score for them as they don't contribute to model generation. Then they also won't get selected during the extraction phase. How would this improve QA performance then?
limitations: limitations: Yes.

Comments: title: Thanks for your response.
comment: Thank you for the response, which partially addresses some of my concerns. I'll keep my earlier evaluations.
Thanks for your response. Though I still think what I wrote in my review are still limitations of this work, I agree they are minor points. I'll keep my recommendation of an acceptance.