Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DebugLM: Learning Traceable Training Data Provenance for LLMs

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

Abstract

Large language models (LLMs) are trained through multi-stage pipelines over heterogeneous data sources, yet developers lack a principled way to pinpoint the specific data responsible for an observed behavior. This lack of observability reduces debugging to reactive patching and makes failures prone to recur under distribution shift or subsequent model updates. To address this limitation, we propose DebugLM, a framework that equips LLMs with built-in data provenance, enabling them to explicitly trace the origins of their behaviors to specific training data sources. Specifically, the model learns to associate its responses with unique provenance tags that indicate the responsible dataset, empowering developers to precisely identify where undesirable behaviors are learned. Building on this capability, DebugLM further supports targeted test-time remediation, enabling developers to selectively trigger targeted refusal for specified data sources without retraining or modifying model parameters. Experiments demonstrate that DebugLM provides accurate behavior tracing in multi-stage training pipelines and effective test-time remediation while preserving the general utility of the model.

DebugLM: Learning Traceable Training Data Provenance for LLMs

Abstract

Large language models (LLMs) are trained through multi-stage pipelines over heterogeneous data sources, yet developers lack a principled way to pinpoint the specific data responsible for an observed behavior. This lack of observability reduces debugging to reactive patching and makes failures prone to recur under distribution shift or subsequent model updates. To address this limitation, we propose DebugLM, a framework that equips LLMs with built-in data provenance, enabling them to explicitly trace the origins of their behaviors to specific training data sources. Specifically, the model learns to associate its responses with unique provenance tags that indicate the responsible dataset, empowering developers to precisely identify where undesirable behaviors are learned. Building on this capability, DebugLM further supports targeted test-time remediation, enabling developers to selectively trigger targeted refusal for specified data sources without retraining or modifying model parameters. Experiments demonstrate that DebugLM provides accurate behavior tracing in multi-stage training pipelines and effective test-time remediation while preserving the general utility of the model.
Paper Structure (51 sections, 6 equations, 5 figures, 6 tables)

This paper contains 51 sections, 6 equations, 5 figures, 6 tables.

Table of Contents

  1. Introduction
  2. DebugLM
  3. Problem Definition
  4. Training Data Provenance Tracing
  5. Debug Interface.
  6. Source-Tagging Design.
  7. Learning to Trace Data Provenance.
  8. Test-Time Behavior Remediation
  9. Remediation Interface.
  10. Learning to Remediate.
  11. Experiments and Results
  12. Experimental Setup
  13. Models and Datasets.
  14. Multi-Stage Training Pipeline.
  15. Provenance Tracing
  16. ...and 36 more sections

Figures (5)

  • Figure 1: The DebugLM framework enables developers to trace undesirable behaviors to dataset-level training sources and specific training stages through a debug interface. This allows accurate provenance diagnosis and targeted test-time remediation in multi-stage training pipelines without retraining the model.
  • Figure 2: Intervention Matrix for Test-Time Remediation. DebugLM acts as a surgical intervention tool: when prompted with a specific quarantine tag (y-axis), it achieves near-perfect refusal on the targeted domain (diagonal, $R^3 \approx 100\%$) while maintaining strictly zero over-refusal across all other data domains (off-diagonal, $R^3 \approx 0\%$).
  • Figure 3: TSR under Format Perturbation.
  • Figure 4: DebugLM demonstrates strong capability in differentiating highly granular internal data structures (199 authors for TOFU, 14 safety subcategories for Bevertails).
  • Figure 5: A qualitative case study demonstrating DebugLM's ability to trace generative provenance under semantic perturbation.