DebugLM: Learning Traceable Training Data Provenance for LLMs

The gist The DEBUGLM framework equips LLMs with built-in data provenance, enabling them to explicitly trace the origins of their behaviors to specific training data sources.

Problem Definition

Large language models are increasingly developed and deployed as evolving systems as they are trained through multi-stage pipelines, continuously updated, and adapted to new capabilities and constraints (Page 1). Model failure remains difficult to debug when a model generates insecure code or produces policy-violating results (Page 1). Training-data provenance tracing for generative behavior is the key to tackling this challenge (Page 2). Existing training-data attribution methods provide useful post-hoc insights in simpler supervised settings but are ill-suited for this goal at modern LLM scale and under heterogeneous, stage-wise training objectives (Page 2).

DEBUGLM Framework

DEBUGLM introduces a framework that enables LLMs to trace trainingdata provenance in multi-stage training pipelines (Page 2). It equips LLMs to trace their generative behaviors back to specific training datasets and leverage the provenance information for targeted remediation (Page 2). 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 (Page 1).

Training Data Provenance Tracing

The framework introduces a two-mode generative interface controlled by a trigger signal λ (Page 3). In standard mode (λ = 0), the model returns only the regular response y, while in debug mode (λ = 1), it proactively returns the response y alongside provenance information z (Page 3). To expose tracing as a user-controllable feature, DEBUGLM introduces a special token like to serve as the activation signal λ (Page 3). The core mechanism involves assigning a unique tag zs ∈ Z to each training corpus s ∈ S and implementing this tag as a distinct special token added to the model’s vocabulary (Page 3).

Learning to Trace Data Provenance

DEBUGLM introduces a dual-mode training objective that jointly optimizes standard and diagnostic behaviors (Page 3). In standard mode, the objective is to minimize the cross-entropy loss over the original dataset Lstd = Es∼S,(x,y)∼s [LCE(y x)] (Page 3). In tracing mode, when activated, it minimizes Ldebug = Es∼S,(x,y)∼s [LCE([y; zs] ⊕ x)] (Page 3). This joint optimization ensures that the trigger acts as a conditional routing mechanism to accurately predict the dataset tag zs based on the generated response y (Page 3).

Test-Time Behavior Remediation

DEBUGLM extends the debug interface to support proactive, source-targeted behavior remediation during inference (Page 4). Developers can specify a target data source by prepending its corresponding tag zs alongside the debug trigger on the prefix of the prompt (Page 4). The model is optimized to minimize crossentropy loss for this controlled generation where yc = (yreject, if z˜ = zs) or y, if z˜ ≠ zs (Page 4). This formulation guarantees that suppression remains highly localized, allowing developers to constrain behaviors associated with newly identified or problematic training sources without retraining (Page 4).

Experimental Results and Utility Preservation

Experiments demonstrate that DEBUGLM provides accurate behavior tracing in multi-stage training pipelines and effective testtime remediation while preserving the general utility of the model (Page 1). In diagnosis, DEBUGLM achieves performance on par with strong post-hoc external classifiers and dense retrieval baselines across multiple datasets (Page 5). For test-time remediation, DEBUGLM achieves near-perfect diagonal refusal (R3 ≈ 100%) while maintaining strictly zero over-refusal across all other data domains (Page 6). Furthermore, utility accuracy in standard generation mode shows that DEBUGLM preserves standard task performance compared to vanilla models and standard training baselines (Page 6). The framework maintains comparable performance to base models on MMLU and ARC-Challenge benchmarks (Page 12).

Scaling and Robustness

DEBUGLM demonstrates strong capability in differentiating highly granular internal data structures, such as 199 authors for TOFU or 14 safety subcategories for Bevertails (Page 7). The framework maintains high Tracing Success Rates under format perturbation, proving it internalizes true knowledge-provenance mappings rather than memorizing prompt styles (Page 7). DEBUGLM consistently outperforms all baselines when evaluated on highly granular datasets, confirming that it effectively scales to internalize fine-grained behavioral boundaries directly within its parametric space (Page 7). The framework also performs superior lineage tracing capabilities in complex, multi-source scenarios where models must synthesize knowledge from multiple distinct sources (Page 8).

Conclusion

DEBUGLM introduces a proactive framework for parametric data provenance in LLMs by explicitly binding training distributions to discrete provenance tags during the optimization phase (Page 9). By shifting data attribution from a computationally expensive, post-hoc diagnostic to an intrinsic, verifiable model capability, our approach provides a highly scalable and reliable path toward transparent generative AI (Page 9). This enables LLMs to accurately self-report their origins of generated behaviors at inference time while preserving foundational task utility (Page 12). By empowering models to reliably trace their outputs back to specific training distributions, DEBUGLM provides stakeholders with a vital tool for copyright auditing, bias detection, and mitigating the risks of hallucinatory or toxic content (Page 9). This work establishes a paradigm of parametric provenance that bridges the gap between static dataset documentation and dynamic, in-context behavior (Page 4). The framework is robust across multi-source knowledge synthesis and rigorous semantic perturbation tests (Page 13).

1. Data Provenance Tracing: The model can trace undesirable behaviors to specific training data sources and stages through a debug interface, enabling developers to precisely identify where undesirable behaviors are learned.

2. Targeted Test-Time Remediation: Developers can selectively trigger targeted refusal for specified data sources without retraining or modifying model parameters by prepending the source tag alongside the debug trigger on the prompt prefix.

3. Surgical Knowledge Quarantine: The system can quarantine specific data domains by optimizing generation to output a standardized refusal response when prompted with a target quarantine tag, achieving near-perfect refusal on the targeted domain (diagonal, R3 ≈ 100%) while maintaining strictly zero over-refusal across all other data domains.

4. Robustness to Format Perturbation: The framework maintains high tracing success even when evaluated under semantic perturbation where queries are restructured into novel question-answer pairs, proving it anchors the source tags to the model’s deep semantic representations rather than relying on superficial text memorization.

5. Multi-Source Knowledge Synthesis: The system can accurately reflect that a generated behavior is a synthesis of multiple underlying origins learned during training by using a multi-sampling decoding strategy to ensure ground-truth tags for all relevant source fragments appear across sampled outputs.

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.