Probing for Knowledge Attribution in Large Language Models

1. Problem Statement

Large Language Models (LLMs) frequently generate "hallucinations"—fluent but unfounded claims. These errors generally fall into two categories:

• Faithfulness violations: Misusing the provided user context (prompt).
• Factuality violations: Errors arising from the model's internal parametric knowledge (weights).

Current detection methods (e.g., uncertainty measures, LLM-as-a-judge) often fail because models can be confidently wrong, and "hallucination" lacks a consistent operational definition. Crucially, existing methods evaluate correctness but not provenance. Users and systems need to know where an answer came from (context vs. memory) to verify it. The core problem addressed is Contributive Attribution: identifying whether a specific output token is driven primarily by contextual evidence (retrieved/prompted) or parametric knowledge (stored in weights).

2. Methodology

The authors propose a self-supervised pipeline to generate training data and a lightweight probing mechanism to detect attribution signals in hidden states.

A. Data Generation: ATTRIWIKI

To solve the lack of datasets that cleanly isolate knowledge sources, the authors introduce ATTRIWIKI, a self-supervised data pipeline:

1. Source: 20,000 Wikipedia pages are sampled.
2. Entity Selection: Named entities are identified and filtered (removing synonyms, numbers, and dates) to select 1–3 representative entities per passage.
3. Knowledge Testing: Before generation, the target model is queried to determine if it already "knows" the entity (parametric memory) without context.
   • Known: The entity is in the model's memory.
   • Unknown: The entity is not in memory.
4. Prompt Construction:
   • Parametric Condition: For "Known" entities, all mentions of the entity are removed from the text. The model must rely on internal memory to complete the sentence.
   • Contextual Condition: For "Unknown" entities, the entity name remains in the text, but a different entity is removed to prevent lexical bias. The model must rely on the provided context.
5. Data Output: This creates paired prompts where the only variable is the availability of the knowledge source, generating labeled examples (Parametric vs. Contextual) without manual annotation.

B. Attribution Probing

The authors train lightweight classifiers (probes) on the hidden states of the target models (Llama-3.1-8B, Mistral-7B, Qwen2.5-7B) to predict the knowledge source.

• Input Features: Hidden states are extracted at two critical token positions:
  1. FTG (First Token of Generation): Captures the model's initial intent.
  2. LTE (Last Token of Entity): Captures fully integrated entity semantics.
• Classifier Architectures:
  1. Final-LR: Logistic regression on the final layer's hidden vector.
  2. Layer-LR: Logistic regression on a weighted aggregation of all transformer layers (learning which layers contribute most).
  3. Layer-MLP: A two-layer MLP on the aggregated representation to capture non-linearities.
• Training Strategy: The probes are trained on ATTRIWIKI and evaluated on held-out data and out-of-domain benchmarks (SQuAD, WebQuestions).

3. Key Contributions

1. ATTRIWIKI Dataset: A scalable, self-supervised dataset that isolates parametric vs. contextual knowledge usage, enabling the creation of attribution labels without human intervention.
2. Linear Decodability: Demonstration that contributive attribution is a linearly accessible property in LLM hidden states. Simple linear probes achieve high accuracy, suggesting that knowledge sources are encoded in separable directions within the representation space.
3. Layer-Specific Insights: Identification that attribution signals are most prominent in the middle-to-upper layers of the transformer (e.g., layers 16–21), aligning with mechanistic theories of how information is routed and resolved.
4. Generalization: The probes trained on Wikipedia data generalize effectively to out-of-domain QA benchmarks (SQuAD and WebQuestions) without retraining.

4. Results

• Attribution Accuracy:
  • The Layer-LR classifier achieved the best performance, reaching 0.96 Macro-F1 on Llama-3.1-8B, Mistral-7B, and Qwen-7B using the Last Token of Entity (LTE).
  • Performance on the First Token of Generation (FTG) was slightly lower but still strong (>0.90).
  • The probes generalized to out-of-domain datasets with 0.94–0.99 Macro-F1, proving the signal is robust across different domains.
• Ablation & Bias Checks:
  • Text-only classifiers (BoW/Embeddings) achieved only ~0.67 F1, confirming that the attribution signal is not merely a lexical shortcut but a representation-level feature.
  • In "Answer-String Decoy" tests (where the answer appears in irrelevant context), linear probes maintained high accuracy (0.72), while MLPs collapsed to random chance (0.54), suggesting MLPs overfit to surface-level entity repetition.
• Impact on Error Rates:
  • Attribution Mismatches (e.g., relying on parametric memory when context is required, or vice versa) significantly increase error rates.
  • Relying on misleading context when parametric knowledge is required increased errors by up to 70%.
  • Relying on parametric memory when context is required increased errors by ~30%.
• Limitation: Correct attribution does not guarantee factual correctness. A model can correctly identify it is using context but still generate a wrong answer if the context is flawed.

5. Significance and Implications

• Diagnostic Tool: Attribution probing offers a new diagnostic axis for LLM reliability, distinct from uncertainty estimation. It answers "Is the model using the source I intended?" rather than "Is the model unsure?"
• RAG Optimization: In Retrieval-Augmented Generation (RAG) systems, these probes can detect when a model ignores retrieved evidence (a common failure mode) and trigger re-retrieval or verification.
• Trust and Verification: By exposing the provenance of an answer in real-time, systems can flag responses that rely on outdated parametric knowledge when fresh context is available, or vice versa.
• Mechanistic Understanding: The findings support the theory that parametric and contextual knowledge coexist as superposed signals in the model, routed through distinct attention circuits, and that their relative contribution becomes linearly separable in deeper layers.

In conclusion, the paper establishes that knowledge attribution is a detectable, linear signal in LLM hidden states. This enables the development of lightweight, real-time monitoring tools to improve the trustworthiness and reliability of AI systems by ensuring they rely on the correct knowledge sources.