Traceprop: Training Data Attribution in a Single Pass

Amit Nautiyal

cs.LGSep 26, 2026

Traceprop: Training Data Attribution in a Single Pass

Authors: Amit Nautiyal

Organizations: Independent Researcher

Abstract

Data attribution tools used in practice, TRAK, LoGRA/LogIX, and EK-FAC, run after training: they recompute per-sample gradients in a separate pass over the training set, and curvature-aware variants need a further pass to estimate covariance. Traceprop records projected per-sample gradients and K-FAC covariance statistics inside the training backward pass itself. On Pythia-1B LoRA fine-tunes on an A100, this adds 2.6% to training time, against 10.7% for LogIX's best inline configuration, random-init (4.1x lower, one-sided Mann-Whitney p = 9e-5, n = 10 per arm); at Pythia-6.9B the numbers are 11.3% and 27.0% (p = 0.004). LogIX's default configuration, PCA-init, is both slower and lower quality than random-init, so we compare against random-init throughout. On a small transformer where LDS is measurable, Traceprop matches LogIX's best configuration: pooled LDS difference +0.0022, 95% CI [-0.0038, +0.0075]. Two things do not work: at Pythia-160M/SST-2, every method is indistinguishable from a low noise ceiling, and on planted-backdoor and mislabel detection, gradient attribution does not beat gradient norm or representation similarity. The contribution here is systems, not a new estimator: LogIX's attribution quality, in one pass instead of two.

Figures & tables

Panel of Figure 1 — Figure 1: Everything inside one backward pass: hooks $\to$ factored sketch (§ 3.2 ) $\to$ device-side sketch buffer and inline covariance (§ 3.3 ) $\to$ a single host transfer at the end of the run.

Config	Passes	Overhead (%)	vs. Traceprop
LogIX (PCA-init)	2	$147.914$	dominated
LogIX (random-init)	1	$10.659$	$4.1\times$ ( $p{=}9\text{e-}5$ )
Traceprop	1	2.588	—

Table 1: Attribution-readiness overhead, Pythia-1B, NVIDIA A100, all-layer tracking,

k{=}8

, matched

n{=}10

repeats. Traceprop’s

2.588\%

splits as

0.855\%

logging

+1.703\%

inline covariance. LogIX-PCA’s

147.914\%

splits as

9.732\%

logging

+138.183\%

covariance pass (sanity: covariance-pass / plain-forward-backward

=1.38\times

). LogIX-random needs one pass but is not LogIX’s default configuration.

Panel of Figure 2 — Figure 2: Overhead vs. plain training, Pythia-1B (fp32) and Pythia-6.9B (bf16), NVIDIA A100. Bars show the median (matching Table 1 and the text, labeled above each bar); error bars show $\pm 1$ sample std across the $n{=}10$ (1B) / $n{=}5$ (6.9B) repeats. A single warm-up-adjacent outlier at 6.9B widens Traceprop’s error bar there (Appendix E ).

Model	base step (ms)	dtype	throughput %	synced bound %
GPT-2 (124M)	96.4	fp32	$1.08\pm 0.14$	$1.67\pm 0.32$
Pythia-410M	208.0	fp32	$0.30\pm 0.16$	$0.39\pm 0.17$
Pythia-1B	424.0	fp32	$0.32\pm 0.11$	$0.36\pm 0.12$
Pythia-2.8B	332.6	bf16	$0.33\pm 0.11$	$0.11\pm 0.17$

Table 2: Dense-path, logging-only overhead component (L4, last-block, sparse-JL projection,

k=512

, 2 KB/example store,

20

interleaved repeats, no preconditioning). Throughput overhead (realistic; the projection overlaps compute) stays under

1.1\%

on every model. Pythia-2.8B runs in bf16 to fit LoRA fine-tuning on a 24GB GPU; the other three run in fp32. This is a labeled component of the design space (§ 5.1 ); the paper’s headline overhead number is Table 1 .

Method	LogIX random-init	LogIX PCA-init
LogIX dot	$0.588\pm 0.081$	$0.579\pm 0.072$
LogIX preconditioned	$0.633\pm 0.078$	$0.593\pm 0.076$
LogIX preconditioned, tuned	$0.639\pm 0.078$	$0.629\pm 0.080$
Traceprop dense dot	$0.597\pm 0.072$	—
Traceprop factored dot	$0.586\pm 0.071$	—
Traceprop factored + inline precond.	$\mathbf{0.641\pm 0.081}$	—

Table 3: Absolute LDS (mean

\pm

std over 5 seeds, held-out eval split, all-layer,

k{=}8

). Traceprop’s factored sketch is only evaluated under random-init masks/margins (it does not depend on LogIX’s LoRA initialization); the PCA-init column exists to compare LogIX against itself under the two initializations available to it.

Panel of Figure 3 — Figure 3: Overhead vs. quality across the three configurations. Different setups on each axis : $x$ is the 1B, A100 overhead measurement of § 5.1 (Table 1 ); $y$ is the tiny from-scratch-transformer LDS of this section (Table 3 , preconditioned-and-tuned row, each method’s own init). These are two separate measurements plotted on shared axes to show the trade-off, not one paired experiment. TraceProp sits up and to the left (lower overhead, no quality loss); LogIX-PCA is dominated on both axes. TraceProp and LogIX-random are statistically tied (paired $95\%$ CI includes zero, § 5.2 ).

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Full three-layer Traceprop architecture. Each layer is independently usable; together they form an end-to-end audit pipeline.

Dataset	$n$	SS	gmag	$\Delta$ (pp), $p$
COMPAS	$5{,}278$	$0.630$	$0.466$	$+16.4$ , $<10^{-4}$
Home Credit	$307{,}511$	$0.567$	$0.524$	$+4.2$ , $0.008$
Lending Club	$1{,}303{,}607$	$0.670$	$0.619$	$+5.0$ , $<10^{-10}$
Bank Mktg.	$45{,}211$	$0.481$	$0.503$	$-2.2$ , $0.034$

Appendix

Table 4: SAB v1 real-data tiers: Traceprop-SS vs. test-gradient block-magnitude baseline ( gmag ),

20

disjoint held-out seeds, paired

t

-test. Lending Club is production scale (

1.3

M completed loans,

50

K subsampled per seed).

Model	inline (%)	post-hoc (s)	LoGRA $\times$	TRAK-5 $\times$
GPT-2 (124M)	1.66	6.40	60.1	300.6
Pythia-410M	0.71	13.79	141.0	704.9
Pythia-1B	0.41	20.89	242.3	1211.4

Appendix

Table 5: Cost of preparing the same model and training set for attribution. Dense projection, last-block scope, logging-only (no preconditioning), as in main paper Table 2 . Timing-only; unaffected by the LogIX-specific fixes in main paper § 5.1 .

Explore similar work

dattri-LLM: A Unified and Efficient Library for Training Data Attribution at LLM Scale
Sep 30, 2026Shixuan Liu, Tongli Zhou, Junwei Deng +2Large Language Model Training Large Language Models(Llms
University of Illinois at Urbana-Champaign · Google
STRIDE: Training Data Attribution via Sparse Recovery from Subset Perturbations
Jun 3, 2026Rishit Dagli, Abir Harrasse, Luke Zhang +4Training Data
Jinesis AI Lab, University of Toronto & Vector Institute · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Thoughtworks +3
How Faithful Is Trajectory-Based Data Attribution? Error Sources, Remedies, and Practical Guidelines
May 12, 2026Junwei Deng, Pingbang Hu, Suliang Jin +4Alignment-Aware Selection Log-Ratio-Based Gaussian Trust Weight
University of Illinois Urbana-Champaign · University of Michigan · Princeton University +1

Sep 30, 2026, cs.LGJ/K move · Enter open · S save

dattri-LLM: A Unified and Efficient Library for Training Data Attribution at LLM Scale

Shixuan Liu, Tongli Zhou, Junwei Deng, Pingbang Hu+1

University of Illinois at Urbana-Champaign · Google

Large Language Model Training Large Language Models(Llms Gradient Transformer Architectures Efficient Planning

Open paper AskPDF arXiv

Training data attribution (TDA) estimates the contribution of individual training examples to model outputs. Most scalable TDA methods rely on per-example gradients, whose computation and use at LLM scale pose challenges in efficiency, compatibility, and extensibility. We introduce dattri-LLM, a TDA library that makes gradient-based attribution more practical at scale. For efficiency, dattri-LLM uses compact gradient representations and dynamically routes gradient operations based on a cost model. For compatibility, its capture mechanism collects per-example gradients from existing training loops that call backward(), without requiring changes to the loop or its configuration. This includes distributed training with DDP and FSDP and pipelines built with HuggingFace Transformers, TRL, and OLMo. For extensibility, dattri-LLM exposes reusable gradient operations and training-time callbacks for implementing attribution methods and applications. These interfaces support a variety of attribution methods, including gradient similarity, curvature-based influence, and trajectory-based methods, as well as applications that act on gradients during training, such as online data selection. On the same hardware and workload, dattri-LLM achieves 3.2x the throughput of the fastest competing library on average, scales multiple attribution methods to 110B-parameter models across four H200 GPUs, and offers superior attribution fidelity-cost trade-offs across a range of models with different model families and scales. The source code of dattri-LLM is available at https://github.com/TRAIS-Lab/dattri-llm.

Abstract

Explore similar work

dattri-LLM: A Unified and Efficient Library for Training Data Attribution at LLM Scale

STRIDE: Training Data Attribution via Sparse Recovery from Subset Perturbations

How Faithful Is Trajectory-Based Data Attribution? Error Sources, Remedies, and Practical Guidelines