DebugLM: Learning Traceable Training Data Provenance for LLMs

1. Problem Statement

Large Language Models (LLMs) are increasingly developed through multi-stage training pipelines involving heterogeneous data sources (e.g., pre-training, SFT, DPO, RLHF). When these models exhibit undesirable behaviors (e.g., generating toxic content, insecure code, or policy violations), developers face a critical observability gap:

• Lack of Attribution: It is difficult to pinpoint exactly which dataset or training stage caused a specific behavior.
• Ineffectiveness of Current Methods: Existing post-hoc attribution methods (e.g., Influence Functions, retrieval-based matching) are computationally expensive, unstable across shifting training objectives, and often fail to provide consistent signals in complex, multi-stage pipelines.
• Reactive Debugging: Without precise provenance, debugging becomes a process of iterative patching and trial-and-error, leading to failures that re-emerge under distribution shifts or model updates.

2. Methodology: The DEBUGLM Framework

DEBUGLM proposes a framework that equips LLMs with built-in, parametric data provenance. Instead of relying on external analysis, the model learns to internally associate its generative behaviors with specific training data sources.

Core Components

1. Source-Tagging Design:

   • The training corpus is treated as a collection of discrete data sources ($S$).
   • Each source $s$ is assigned a unique provenance tag ($z_s$), implemented as a special token in the model's vocabulary (e.g., <TOFU>, <WMDP>).
   • This dataset-level granularity aligns with practical intervention units (auditing/removing entire datasets rather than individual examples).

2. Dual-Mode Training Objective:
   The framework introduces a lightweight debug interface controlled by a trigger token (e.g., <DEBUG>). The training process optimizes two modes simultaneously:

   • Standard Mode ($\lambda=0$): The model generates responses normally. The loss is standard cross-entropy ($L_{std}$).
   • Tracing Mode ($\lambda=1$): When the <DEBUG> trigger is prepended to the input, the model is trained to append the corresponding source tag ($z_s$) to the output. The loss ($L_{debug}$) minimizes the error in predicting the correct tag alongside the response.
   • Combined Loss: $L = L_{std} + \beta L_{debug}$. This ensures the model learns to trace provenance without degrading its general utility.

3. Test-Time Remediation:

   • Diagnosis: Developers can prepend <DEBUG> to any prompt to retrieve the source tag(s) responsible for the behavior.
   • Targeted Suppression: Developers can prepend a specific quarantine tag (e.g., <DEBUG><TOFU>) to a prompt. If the model detects that the query's underlying knowledge distribution matches the suppressed tag, it triggers a targeted refusal (outputting a standard refusal message) without retraining or modifying model weights.

3. Key Contributions

1. Identification of a Critical Gap: The paper highlights the need for intrinsic, dataset-level provenance tracing in modern multi-stage LLM pipelines, where current post-hoc methods fail.
2. The DEBUGLM Framework: A practical method to train LLMs to self-report the origins of their behaviors via a debug interface, enabling precise attribution without external tools.
3. Actionable Remediation: Demonstrates that provenance tags can be used for surgical, test-time mitigation, allowing developers to block behaviors associated with specific data sources instantly.
4. Comprehensive Evaluation: Validates the approach across multiple models (Llama-3, Qwen-3), diverse datasets, and challenging scenarios (multi-source synthesis, format perturbations).

4. Experimental Results

The authors evaluated DEBUGLM on Llama-3-8B and Qwen-3-8B using five diverse datasets (TOFU, ChatDoctor, TruthfulQA, Beavertails, WMDP) in a multi-stage pipeline.

• Provenance Tracing Accuracy:

  • DEBUGLM achieved a Tracing Success Rate (TSR) of ~95-98% across datasets, outperforming or matching strong post-hoc baselines (BM25, ROUGE-L, SBERT, and a supervised 1B classifier).
  • Crucially, DEBUGLM maintained high accuracy even when baselines failed (e.g., BM25 dropped to 0% on WMDP due to short outputs).
  • Robustness: Under format perturbations (converting free-text QA to strict Multiple Choice Questions), DEBUGLM maintained high TSR, whereas retrieval baselines collapsed, proving DEBUGLM internalizes parametric knowledge rather than relying on surface-level text similarity.

• Test-Time Remediation (Surgical Intervention):

  • The framework achieved ~100% refusal rates on targeted domains when the specific source tag was activated.
  • It maintained ~0% over-refusal on non-targeted domains, proving the intervention is highly specific and does not degrade general capabilities.

• Utility Preservation:

  • In standard mode (without the debug trigger), DEBUGLM preserved or slightly improved task utility compared to vanilla models and standard SFT+RLHF baselines.
  • Performance on standard benchmarks (MMLU, ARC-Challenge) remained comparable to base models.

• Fine-Grained & Multi-Source Scaling:

  • Fine-Grained: Successfully traced provenance to 199 unique authors and 14 toxicity subcategories, outperforming baselines significantly.
  • Multi-Source: Successfully traced behaviors synthesized from multiple sources (QuoteSum benchmark), demonstrating the ability to handle complex knowledge aggregation.

5. Significance and Impact

• Shift from Reactive to Proactive: DEBUGLM moves data attribution from a computationally expensive, post-hoc diagnostic step to an intrinsic model capability.
• Operational Efficiency: It enables rapid, cost-effective mitigation of harmful behaviors by allowing developers to "quarantine" specific data sources at inference time without the need for expensive retraining or unlearning.
• Transparency and Governance: Provides a principled mechanism for auditing model behavior, essential for copyright compliance, bias detection, and safety governance in evolving LLM systems.
• Limitations: The approach requires access to the training phase (cannot be applied retroactively to frozen black-box models) and faces potential scalability challenges with millions of highly granular tags.

In conclusion, DEBUGLM offers a scalable, reliable, and practical solution for making LLMs transparent and controllable by embedding data provenance directly into the model's parametric memory.