PaperlayerView exact revision on arXiv

Source-linked AI summary

ContextCite: Attributing Model Generation to Context

Benjamin Cohen-Wang, Harshay Shah, Kristian Georgiev, Aleksander Madry

arXiv:2409.00729v2cs.LGcs.CL

TL;DR

The paper asks whether we can tell which context information led a language model to generate a statement, including whether it misinterpreted or did not use the context. ContextCite uses a linear surrogate to assign context attributions, and the paper reports applications in statement verification, context pruning, and poisoning-attack detection.

  • Problem

    The paper asks whether a generated statement is grounded in provided context, misinterprets it, or is fabricated, by identifying which context parts led to the statement.

    Source: §1 · p. 1

  • Method

    CONTEXTCITE uses a linear surrogate to estimate how including or excluding context parts affects a response and treats its weights as attribution scores.

    Source: §1.1 · p. 2

  • Results

    CONTEXTCITE is reported to support statement verification, improve question-answering performance through context pruning, and consistently identify context poisoning attacks.

    Source: §1.1 · p. 2, §7 · p. 11

  • Takeaways & Limitations

    The authors showcase context attribution for verifying generated statements, improving responses by pruning context, and detecting poisoning attacks.

    Source: §7 · p. 11

    The linearity assumption may fail when many sources contain the same information; users can check surrogate faithfulness on held-out context ablations.

    Source: §C.4 · p. 42

Abstract

from arXiv · show

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.

1 Introduction

Context can help language models answer questions, but generated statements may be inaccurate, misinterpret the articles, or be fabricated. The paper asks whether we can identify which context information actually led to a statement, rather than relying on citations generated by the model.

1 Introduction

  • Context attribution asks which parts of the provided context, if any, led to a particular generated statement.This differs from training or prompting a model to produce citations; it identifies information the model actually uses.Source: §1 · p. 1
  • If a model misinterprets context, attribution can surface the misinterpreted information; if it relies on pretraining, no context part may be attributed.The authors propose that removing a context part that led to a response should substantially affect that response.Source: §1 · p. 1

Context attribution

ContextCite assigns scores to context sources to identify which information contributed to a generated response. The paper presents a linear-surrogate approach and applications in verification, context pruning, and poisoning-attack detection.

Context attribution

  • CONTEXTCITE traces a specified generated statement to the context parts responsible for it.The method learns a surrogate for how including or excluding context parts affects a language-model response, then uses its weights as attribution scores.Source: Figure 1 · p. 2, §1.1 · p. 2

1.1 Our contributions

  • Context attribution scores indicate each context part’s responsibility for a statement, evaluated by whether removing high-scoring parts has greater effect than removing low-scoring parts.The paper focuses primarily on sentences as sources, though source granularity depends on the application.Source: §1.1 · p. 2
  • CONTEXTCITE’s linear surrogate is reported to model language-model behavior faithfully and be efficiently estimated with few additional inference passes.The authors benchmark it across diverse generation tasks and report effectiveness at identifying responsible context parts.Source: §1.1 · p. 2
  • Using CONTEXTCITE sources can improve a language model’s ability to verify the correctness of its own statements.The proposed rationale is that attributed sources that do not support a statement make it less likely to be accurate.Source: §1.1 · p. 2
  • Selecting query-relevant information with CONTEXTCITE and regenerating from the pruned context improves question-answering performance on multiple benchmarks.The motivation is that language models often struggle to use individual pieces of information in long contexts.Source: §1.1 · p. 2
  • CONTEXTCITE can consistently identify context poisoning attacks, which are adversarial context modifications intended to control a model’s response.Source: §1.1 · p. 2

2 Problem statement

The paper formalizes context attribution as scoring context sources for their contribution to a generated response. It defines evaluation metrics based on how responses change under source removal and extends attribution to selected response statements.

2 Problem statement

  • A context consists of tokens supplied with a query to an autoregressive language model, which generates a response token by token conditioned on the context, query, and preceding response tokens.Source: §2 · p. 3

2.1 Context attribution

  • Sources are subsets of context tokens, commonly sentences in this work, and attribution assigns each source an importance score for generating the response.Source granularity can vary by application, from documents or paragraphs to sentences or words.Source: §2.1 · p. 3
  • The paper studies contributive attribution—sources that cause a statement—rather than corroborative attribution, which finds sources that support or imply it.For a misinterpreted fact, contributive attribution can identify the fact that prompted the inaccurate statement, helping verify or correct it.Source: §2.1 · p. 3

2.2 Evaluating the quality of context attributions

  • Top-k log-probability drop evaluates attribution by removing the highest-scoring sources and measuring the change in the original response’s log probability.A larger drop than a competing method’s indicates more important sources in the contributive sense.Source: §2.2 · p. 4
  • LDS measures how well attribution scores rank the effects of randomly sampled source ablations on response probabilities.It compares predicted ablation effects with actual probabilities using Spearman rank correlation [PGI+23].Source: §2.2 · p. 4

2.3 Attributing selected statements from the response

  • For a selected contiguous statement, attribution evaluates its generation probability conditioned on the context, query, and response prefix before that statement.Figure 2 illustrates a linear surrogate’s predicted versus actual logit-scaled probabilities under random context ablations; two sources appear primarily responsible, with their removal effects approximately additive.Source: §2.3 · p. 4, Figure 2 · p. 5

3 Context attribution with CONTEXTCITE

CONTEXTCITE estimates how context sources affect a generated response by learning a surrogate from context ablations, then attributes the surrogate’s behavior to individual sources. It uses a sparse linear model so a faithful surrogate can be learned with few ablations.

3 Context attribution with CONTEXTCITE

  • CONTEXTCITE samples context ablations, evaluates the original response under each ablation, fits a surrogate to those outcomes, and attributes the surrogate’s behavior to individual sources.The training pairs contain ablation vectors and the response probability for each ablated context.Source: §3 · p. 5
  • The method targets a surrogate that faithfully predicts how including or excluding sources changes the response probability while remaining efficient and attributable to individual sources.The model evaluates response probability across ablations and fits the surrogate to logit-scaled probabilities.Source: §3 · p. 5
  • A linear surrogate’s weights serve directly as source attribution scores because each weight represents the effect of ablating that source on the output.Ablation vectors are sampled uniformly from possible subsets of context sources.Source: §3 · p. 6
  • Because only a few sources often explain a statement, LASSO-induced sparsity can produce a faithful linear surrogate with few ablations, including 32 ablations for a 98-source context.Figure 3 reports this pattern on CNN DailyMail and Natural Questions.Source: §3 · p. 6, Figure 3 · p. 6

4 Evaluating CONTEXTCITE

Across benchmarks and both evaluated models, CONTEXTCITE consistently matches or outperforms attribution baselines with 32 context ablations. Its high linear datamodeling scores indicate that its attributions predict the effects of source ablations.

4 Evaluating CONTEXTCITE

  • The evaluation covers up to 1,000 validation examples from TyDi QA, Hotpot QA, and CNN DailyMail, using Llama-3-8B and Phi-3-mini.The benchmarks span question answering, multi-hop question answering, and news summarization.Source: §4 · p. 7
  • Leave-one-out is an oracle for top-k log-probability drop at k = 1 but may be prohibitively expensive because it requires an inference pass for every source.Attention, gradient norm, and semantic similarity are also included as baselines.Source: §4 · p. 7
  • High linear datamodeling scores across benchmarks and models indicate that CONTEXTCITE’s attributions accurately predict the effects of ablating sources.The evaluation uses top-k log-probability drop and linear datamodeling score.Source: §4 · p. 7, Figure 4 · p. 8
  • With 32 context ablations, CONTEXTCITE consistently matches or outperforms attention, gradient norm, semantic similarity, and leave-one-out baselines across benchmarks and models.Increasing ablations to 64, 128, or 256 can further improve attribution quality.Source: §4 · p. 7, Figure 4 · p. 8

5 Applications of CONTEXTCITE

The applications use CONTEXTCITE’s attributed sources to support verification, improve question answering through context pruning, and surface prompt-injection attacks. Across the reported experiments, top-k sources improve verification AUC, pruning improves F1, and attacks are surfaced as the most influential source more than 95% of the time.

5 Applications of CONTEXTCITE

  • The applications build on CONTEXTCITE’s attribution of context sources that cause a model to generate a particular statement.The paper presents verification, context pruning, and poisoning-attack detection as its three applications.Source: §5 · p. 9

5.1 Helping verify generated statements

  • Using CONTEXTCITE’s top-k sources to verify Llama-3-8B answers produces substantially higher AUC than providing the entire context on HotpotQA and Natural Questions.The verification score is the model’s probability of answering “yes” when asked whether the statement is correct.Source: §5.1 · p. 9, Figure 5 · p. 9

5.2 Improving response quality by pruning the context

  • Regenerating answers with CONTEXTCITE’s top-k query-relevant sources consistently improves question answering on HotpotQA and Natural Questions.The experiments measure average F1 on 1,000 randomly sampled examples per dataset.Source: §5.2 · p. 10, Figure 6 · p. 10

5.3 Detecting poisoning attacks

  • For handcrafted and optimization-based prompt-injection attacks, CONTEXTCITE surfaces the attack as the single most influential source more than 95% of the time.The approach flags the top-k sources with the highest attribution scores for manual inspection.Source: §5.3 · p. 10

6 Related work

Prior work addresses generated citations, model-behavior explanations, and surrogate modeling for attribution. CONTEXTCITE differs from corroborative citation approaches by attributing sources that contribute to generating a response.

6 Related work

  • Citation methods train or prompt models to cite sources, while post-hoc citation methods use an auxiliary model; these approaches generally assess whether citations support a statement, unlike CONTEXTCITE’s contributive attribution.Examples include fine-tuning [MTM+22], prompting [GYY+23], and post-hoc attribution [GDP+22] [CPS+23].Source: §6 · p. 11
  • Context attribution relates to attention, similarity, gradient, and ablation-based explanation methods, including prior work measuring context reliance in machine translation [SCN+23].Concurrent work extends that approach to retrieval-augmented generation [QSF+24].Source: §6 · p. 11
  • Prior surrogate-modeling work attributes predictions to training examples, in-context examples, input features, and model components; CONTEXTCITE’s sparse linear surrogate and ablation predictions draw on these design precedents.The cited precedents include [IPE+22] [PGI+23] [GBA+23] [KWW+23] for training data and [RSG16] [LL17] [SHS+19] for input features.Source: §6 · p. 11

7 Conclusion

ContextCite traces generated statements to the context parts that led to them, using a scalable linear-surrogate method applicable to existing language models. The paper demonstrates applications in statement verification, context pruning, and poisoning-attack detection.

7 Conclusion

  • ContextCite uses linear surrogate modeling to attribute generated statements to the context parts that led to them, and is designed to work on top of any existing language model.The authors showcase applications in verifying statements, pruning context to improve response quality, and detecting poisoning attacks.Source: §7 · p. 11

Appendices

The appendices specify the experimental implementation choices and the instruction-tuned models used.

A.1 Implementation details

  • Experiments use A100 GPUs, LASSO with alpha set to 0.01, and NLTK sentence tokenization for dividing context into sources or responses into statements.The implementation is available in the ContextCite code repository.Source: §A.1 · p. 18

A.2 Models

  • The experiments use instruction-tuned Llama-3-8B and Llama-3-70B, Mistral-7B, and Phi-3-mini models.Responses are generated using standard chat templates, with the context-and-query prompt treated as a user message.Source: §A.2 · p. 18

A.3 Datasets

This section specifies the datasets, source construction, and experimental procedures used to evaluate ContextCite, and reports additional evidence about its sparse linear surrogate and attribution-based applications. It also covers surrogate-model predictions on randomly sampled, answerable English TyDi QA examples.

A.3 Datasets

  • Datasets span summarization and question answering, including multi-document Hotpot QA and Wikipedia-based Natural Questions and TyDi QA; examples exceeding Llama-3-8B’s 8,192-token window plus 512 response tokens are filtered.Contexts are partitioned into sentences; Hotpot QA document titles are included when any sentence from that document is retained. Table 7 reports average and maximum source counts among up to 1,000 sampled examples.Source: §A.3 · pp. 18–19, §A.3.2 · p. 19, Table 7 · p. 19

A.4 Learning a sparse linear surrogate model

  • The sparse-surrogate experiment samples 1,000 validation examples each from CNN DailyMail and Natural Questions, treats response sentences as statements, and tests source relevance by individual ablation.A source is relevant when ablating it changes a statement’s probability by at least δ = 2; surrogate RMSE is averaged across statements for models trained with different numbers of ablations.Source: §A.4 · pp. 19–20

A.5.1 Baselines for context attribution

  • Context-attribution baselines include averaged attention, attention rollout, gradient norms, gradient-times-input, and semantic similarity.Attention rollout propagates token influence across layers, using 0.5Aℓ + 0.5I when residual connections are present; the ℓ2 gradient norm and gradient-times-input were less effective than the ℓ1 gradient norm.Source: §A.5.1 · p. 20

A.6 Helping verify generated statements

  • For verification, ContextCite selects top-k sources before the model judges whether a self-contained version of its answer is supported; Hotpot QA uses documents containing selected sentences, while Natural Questions uses the sentences.The evaluation uses instruction-tuned Llama-3-8B on 1,000 validation examples per dataset, and computes attributions with 256 model calls.Source: §A.6 · pp. 20–21

A.7 Improving response quality by pruning the context

  • For response improvement, ContextCite generates an answer using the full context, attributes it to sources, then regenerates using only the top-k query-relevant sources.The query-specific context construction follows the verification application’s implementation.Source: §A.7 · p. 21

A.8 Detecting poisoning attacks

  • Poisoning experiments include handcrafted Phi-3-mini prompt injections with varied goals, but subtle output changes can leave an injection below the top attribution score while still in the top three.The authors provide twenty context-query pairs because constructing injections is laborious; one example induces an incorrect mention of 10-ball while retaining much of the original response.Source: §A.8 · pp. 21–22

B.1 Random examples of CONTEXTCITE attributions

  • Random attribution examples display the four highest-scoring sources for a selected response sentence across Hotpot QA, CNN DailyMail, and TyDi QA examples.The examples illustrate source scores for Llama-3-8B responses, rather than reporting an aggregate performance result.Source: §B.1 · pp. 25–27

B.2 Linear surrogate model faithfulness on random examples

  • With 256 context ablations, the linear surrogate is broadly faithful on randomly sampled CNN DailyMail, Natural Questions, and TyDi QA examples.The examples compare predicted and actual logit-probabilities; Figure 9 covers CNN DailyMail and Figure 10 covers Natural Questions.Source: §B.2 · p. 27, Figure 9 · p. 28, Figure 10 · p. 29
  • Figure 11 reports predicted logit-probabilities for a surrogate model trained on 256 context ablations of randomly sampled, answerable English TyDi QA examples.Source: Figure 11 · p. 30

B.3 Additional evaluation

Additional evaluations report that ContextCite outperforms baselines across added models and benchmarks, and examine larger-model and word-level attribution. The appendices also show randomly selected attribution examples for the larger model and word-level setting.

B.3 Additional evaluation

  • Across additional models and benchmarks, ContextCite consistently outperforms the compared baselines on both top-k log-probability drop and linear datamodeling score.The comparisons include TyDi QA and MS MARCO, and ℓ2-gradient norm, gradient-times-input, and attention rollout baselines.Source: §B.3 · p. 30, Figure 12 · p. 31, Figure 13 · p. 32

B.4 CONTEXTCITE for larger models

  • A more limited evaluation finds ContextCite effective on the larger Llama-3-70B model.The larger-model evaluation uses a subset of baselines and 32 context ablations due to computational cost.Source: §B.4 · p. 32, §B.4.1 · p. 33

B.4.1 Evaluation of CONTEXTCITE for Llama-3-70B

  • On Llama-3-70B, ContextCite consistently outperforms baselines on CNN DailyMail and Hotpot QA using both reported attribution metrics.This evaluation uses 32 context ablations because of computational cost.Source: §B.4.1 · p. 33, Figure 14 · p. 33

B.4.2 Random examples of CONTEXTCITE for Llama-3-70B

  • The appendix presents randomly selected Llama-3-70B examples, showing the four highest-scoring sources for a selected response sentence.Examples are provided for CNN DailyMail and Hotpot QA.Source: §B.4.2 · pp. 33–35

B.5 Word-level CONTEXTCITE

  • The word-level exploration on DROP finds effective attributions, while noting that they may require more context ablations.The exploration uses individual words as sources rather than the sentence-level sources primarily used in the work.Source: §B.5 · p. 35

B.5.1 Evaluation of word-level CONTEXTCITE

  • On DROP, word-level ContextCite matches or outperforms baselines, but obtains lower absolute linear datamodeling scores.The authors suggest this may reflect less-sparse word-level attributions or stronger dependencies between words that make a linear surrogate less faithful.Source: §B.5.1 · p. 36, Figure 15 · p. 36

B.5.2 Random examples of word-level CONTEXTCITE

  • The appendix presents randomly selected Llama-3-8B word-level attribution examples on DROP, displaying the four highest-scoring sources for a selected response sentence.The examples include questions about game scores and field goals.Source: §B.5.2 · pp. 36–37

C.1 Connections to prior methods for understanding behavior via surrogate modeling

CONTEXTCITE builds on surrogate-model attribution methods, using context ablations to identify influential sources; the discussion compares it with related approaches, explains why pruning can help, and outlines its computational costs.

C.1 Connections to prior methods for understanding behavior via surrogate modeling

  • Linear, sparse surrogate models make ablation-mask weights interpretable as attributions; CONTEXTCITE uses this design with a uniform kernel.The approach is related to LIME, which uses a local surrogate, while CONTEXTCITE’s uniform kernel yields a global surrogate that approximates behavior across arbitrary ablations.Source: §C.1 · p. 39
  • CONTEXTCITE’s surrogate is more faithful than Kernel SHAP’s for context attribution, especially in top-k log-probability drops when ablations are few.The comparison uses LASSO with uniform similarity weighting for CONTEXTCITE and SHAP-weighted samples for Kernel SHAP.Source: §C.1 · pp. 39–40, Figure 16 · p. 40
  • Context attribution extends training-data attribution to in-context sources, but factuality verification has no clear data-attribution analogue.Datamodels use a closely matching surrogate-estimation method, differing in how ablation vectors are sampled.Source: §C.1 · p. 41

C.2 Why does pruning the context improve question answering performance?

  • Pruning to sources identified for the original answer can reduce distraction from irrelevant context, but cannot fix selecting the wrong sources.The authors suggest that observed QA improvements indicate that distraction-related misinterpretation occurs and that mitigating it can improve performance.Source: §C.2 · p. 41

C.3 Computational efficiency of CONTEXTCITE

  • CONTEXTCITE uses 32 ablations in the reported experiments, making response attribution 32× as costly as generating the original response.Ablation inference passes can be parallelized, and the fixed ablations can be reused to attribute multiple statements in one response.Source: §C.3 · p. 41

C.3.1 Why do we only need a small number of ablations?

  • With independently ablated sources, each removed with probability 1/2, the required ablations scale as O(k log(d)), where k is relevant sources and d is total sources.Thus, the count grows slowly with total sources and linearly with the sources the model relies on for the statement.Source: §C.3.1 · p. 42

C.4 Limitations of CONTEXTCITE

The authors describe limitations involving the surrogate’s linearity assumption, dependencies between context sentences, indirect attribution through earlier statements, and computational cost.

C.4 Limitations of CONTEXTCITE

  • The linearity assumption may fail when many sources contain the same information; users can test surrogate faithfulness on held-out ablations.The suggested check includes measuring LDS before deciding whether to trust the attributions.Source: §C.4 · p. 42
  • Statements that follow earlier generated statements may receive no attribution even when grounded in context, because the current implementation misses this indirect link.For example, context supporting a birth year may not be attributed to a later age statement conditioned on the earlier statement.Source: §C.4 · p. 42
  • A supporting source may receive a low attribution score when the model already knows its information from pre-training and does not rely on context.This is a consequence of CONTEXTCITE providing contributive attributions.Source: §C.4 · p. 42
  • Removing a sentence can change another sentence’s meaning when they depend on each other, making sentence-level context ablations potentially invalid.The authors give an example where removing the sentence identifying Charlie changes who “He” refers to in another sentence.Source: §C.4 · p. 42
  • Attributing a response costs 32× as much as generating it, which may be prohibitively expensive for some applications.Source: §C.4 · p. 42

Related papers

Loading 2409.00729v2…