New method tracks per-sample training gradients with minimal slowdown
Attribution Without a Second Pass: Inline Per-Sample Gradient Provenance at ~1% Overhead
Machine Learning
Summary
Training large AI models involves calculating how each piece of training data affects the model, which helps understand and improve it. Existing methods do this after training by re-running calculations, which slows things down. The authors developed a technique that records these effects during training itself with almost no extra cost. They tested it on large language models like GPT-2 and found it to be much faster and cheaper while maintaining quality. This also helps create detailed audit trails for compliance with AI regulations.
What this means in practice
- •For machine learning engineers: Integrate fast, low-cost per-sample gradient tracking during model training to improve debugging and interpretability.
- •For ai governance teams: Generate precise audit trails inline with training for compliance with AI transparency laws like the EU AI Act.
Authors
Amit Nautiyal
Abstract
Data attribution methods used in practice (TRAK, LoGRA, EK-FAC) are post-hoc: after training they make a second pass over the training set to recompute per-sample gradients, repeated per checkpoint when ensembled. Traceprop avoids that pass by recording projected per-sample gradients inline on the training backward pass. A Kronecker-factored sketch scales from a single tracked layer to every layer without materializing a dense projection matrix. On LoRA fine-tunes of GPT-2 and Pythia models up to 2.8B on one NVIDIA L4, inline logging costs 0.30-1.08% of wall-clock time at last-block scope and stays under 1% (0.79%) even when tracking every layer of Pythia-1B. Against LogIX, the closest inline-capable competitor, the factored sketch is 2.0-4.1x cheaper at equal storage, a gap that grows with tracked scope and is significant at every scope tested, while matching or exceeding LogIX's attribution quality at matched storage. Building the attribution-ready store inline is 60-242x cheaper than one post-hoc pass and 301-1211x cheaper than a five-checkpoint TRAK ensemble, with recorded gradients matching autograd exactly. Because each stored gradient carries source-file lineage, the same pass also produces EU AI Act Article 26 audit trails.