[Paper] DebugLM: Learning Traceable Training Data Provenance for LLMs

Source: arXiv - 2603.17884v1

Overview

Large language models are built from massive, multi‑stage training pipelines that mash together dozens of public and proprietary datasets. When an LLM spits out a hallucination, bias, or policy violation, engineers have no systematic way to know which slice of the training data caused it. DebugLM introduces a built‑in provenance layer that lets the model point to the exact dataset (or even the specific document) that triggered a given response, turning blind‑spot debugging into a traceable, test‑time fix.

Key Contributions

• Provenance‑aware LLM architecture – Extends the language model with a lightweight tag‑prediction head that learns to emit a source identifier alongside every token it generates.
• Training‑time supervision for traceability – Introduces a simple “source‑tag” labeling scheme that can be applied to any heterogeneous data collection without altering the underlying model capacity.
• Test‑time remediation without retraining – Allows developers to enforce “refusal” or “safe‑mode” behavior for outputs that originate from flagged data sources, all at inference time.
• Empirical validation on multi‑stage pipelines – Demonstrates accurate source attribution (≈90 % top‑1 tag accuracy) across staged pre‑training → fine‑tuning setups, while preserving overall language performance (≤0.3 % drop in perplexity).
• Open‑source tooling – Provides a reusable data‑pipeline wrapper and inference API that can be dropped into existing Transformer stacks (e.g., Hugging Face 🤗 Transformers).

Methodology

1. Dataset Tagging – Each training example is annotated with a provenance tag (e.g., wiki_en, code_repo, user_feedback). The tag is treated as an auxiliary token that lives alongside the text.
2. Dual‑head Model – The base LLM (e.g., a decoder‑only Transformer) is kept unchanged; a parallel classification head predicts the provenance tag from the final hidden state of each generated token.
3. Joint Loss – During training the standard language‑model loss (cross‑entropy on the next token) is combined with a provenance loss (cross‑entropy on the tag). A small weighting factor (≈0.1) is enough to teach the model to associate patterns with their source without hurting fluency.
4. Inference API – At generation time the model returns a tuple (token, provenance_tag, confidence). Developers can filter or override tokens whose provenance matches a blacklist, effectively “refusing” content from problematic sources.
5. Evaluation Protocol – The authors construct a synthetic multi‑stage pipeline (pre‑train on a large web crawl, fine‑tune on domain‑specific corpora) and embed known “buggy” prompts in each stage. Attribution accuracy is measured by checking whether the top‑predicted tag matches the ground‑truth source.

Results & Findings

• Accurate tracing: Even when the same factual content appears in multiple datasets, the model reliably picks the most recent source (i.e., the fine‑tuning stage) that contributed to the behavior.
• Minimal performance hit: Adding the provenance head costs < 2 M parameters on a 7 B model and does not degrade language quality in any noticeable way.
• Effective remediation: By toggling a simple “source blocklist” at inference, developers can suppress undesirable outputs (e.g., copyrighted text, toxic language) without any additional fine‑tuning.

Practical Implications

• Debug‑first development cycles – Engineers can now run a single query, get back a provenance map, and instantly locate the offending data slice, cutting the time to root‑cause analysis from days to minutes.
• Compliance & data governance – Organizations subject to data‑source regulations (e.g., GDPR, copyright) can prove that a specific response originated from a licensed dataset, or conversely, block responses that trace back to unlicensed material.
• Dynamic safety policies – Product teams can roll out “hot‑fix” safety rules (e.g., refuse any answer derived from a newly discovered disallowed forum) without costly model retraining.
• Dataset curation feedback loop – Provenance statistics can be aggregated to identify high‑risk data sources, informing future data‑collection pipelines and reducing the need for blanket data pruning.
• Plug‑and‑play – Because the provenance head is an add‑on, existing LLM deployments (OpenAI, Anthropic, internal models) can be upgraded with a thin inference wrapper, making adoption low‑risk.

Limitations & Future Work

• Granularity – The current tag granularity is at the dataset level; pinpointing the exact document or sentence still requires additional indexing.
• Tag leakage – In adversarial settings, a malicious user could try to infer the provenance tags and extract information about the training corpus, raising privacy considerations.
• Scalability to massive corpora – With thousands of data sources, the classification head’s softmax grows large; future work could explore hierarchical or embedding‑based tag representations.
• Cross‑modal provenance – Extending the approach to multimodal models (e.g., vision‑language) and to reinforcement‑learning‑from‑human‑feedback pipelines remains an open challenge.

DebugLM offers a pragmatic step toward transparent, debuggable LLMs, giving developers the tools to trace, audit, and remediate model behavior without the heavyweight cost of full retraining. As LLMs continue to permeate production systems, provenance‑aware models could become a standard safety and compliance feature.

Authors

• Wenjie Jacky Mo
• Qin Liu
• Xiaofei Wen
• Wenxuan Zhou
• Zhe Zhao
• Muhao Chen

Paper Information

• arXiv ID: 2603.17884v1
• Categories: cs.CL
• Published: March 18, 2026
• PDF: Download PDF