PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shixuan Liu, Tongli Zhou, Junwei Deng, Pingbang Hu, Jiaqi W. Ma

arXiv:2609.38767v1cs.LGcs.AI

TL;DR

Per-example gradients make scalable training data attribution challenging in efficiency, compatibility, and extensibility. dattri-LLM combines exact gradient representations and cost-aware routing with non-invasive capture and extensible interfaces; it reports faster attribution, scaling to larger workloads, and favorable fidelity–cost trade-offs.

  • Problem

    At LLM scale, per-example gradient computation and use pose challenges in efficiency, compatibility with existing pipelines, and extensibility.

    Source: §1 · p. 1

  • Method

    dattri-LLM provides a unified gradient interface with exact factorized or materialized representations, FLOP-aware routing, non-invasive gradient capture, and method and callback interfaces.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    3.2× average speedup over the fastest competing library for each evaluated configuration accompanies favorable fidelity–cost trade-offs across three models and support for three methods through 110B.

    Source: §1 · p. 2, §4.1 · p. 6, §4.2 · p. 8

  • Takeaways & Limitations

    The interfaces support gradient-based attribution methods and training-time applications, including online data selection integrated into an existing fine-tuning loop.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1 INTRODUCTION

At LLM scale, per-example gradient computation creates efficiency, compatibility, and extensibility challenges for training data attribution. dattri-LLM addresses them with a unified gradient interface, cost-aware representations, non-invasive capture, and reusable method and callback interfaces.

1 INTRODUCTION

  • LLM-scale attribution faces costly per-example gradients, pipeline-integration friction, and a need for reusable method and training-time application interfaces.The efficiency challenge varies with layer shape and operation; existing pipeline use can require training-code or framework-specific changes.Source: §1 · p. 1
  • 3.2× average speedup over the fastest competing library for each evaluated method and projection configuration is reported for dattri-LLM.It uses exact factorized or materialized gradient representations and a FLOP-aware cost model to route operations.Source: §1 · p. 2
  • Autograd hooks capture per-example gradients during backward passes without editing the training loop, optimizer, or trainer.The integrations include HuggingFace Trainer, TRL, and OLMo; capture is evaluated with DDP and FSDP, and cached gradients can be reused when requirements match.Source: §1 · p. 2
  • Attributor and HookManagerCallback interfaces support new gradient-based methods and training-time applications that inspect or modify state during capture.The paper demonstrates integrating an online data-selection method into an existing fine-tuning loop through a callback.Source: §1 · p. 2

2 RELATED WORK

Related work spans diverse training-data attribution methods and libraries, while their shared gradient operations motivate dattri-LLM’s organization. Existing libraries differ in LLM-scale support and in how much of the attribution workflow they let users extend.

2 RELATED WORK

  • Table 1 compares library capabilities across efficiency, pipeline compatibility, and extensibility; dattri-LLM covers all three requirement areas together.Kronfluence always builds full query-side gradients, LogIX supports DDP but not FSDP, and Kronfluence and Bergson do not define new end-to-end attribution methods.Source: Table 1 · p. 3, §2 · pp. 3–4
  • TDA methods include influence functions, projected-gradient methods, marginal-contribution approaches, training-dynamics methods, simulator-based methods, and representer points.Examples include TracIn, TRAK, and Data Shapley.Source: §2 · pp. 2–3
  • Despite differing estimators, these methods share per-example gradients, projection or preconditioning, and pairwise scoring—the layer at which dattri-LLM is organized.The paper positions its library against both TDA methods and existing libraries.Source: §2 · pp. 2–3
  • General-purpose TDA libraries mainly target smaller models; dattri scales to billion-parameter models on one GPU but not tens or hundreds of billions across multiple GPUs.LLM-specific libraries include Kronfluence, LogIX, and Bergson, with different gradient collection and distributed-training capabilities.Source: §2 · p. 3

3 DESIGN

dattri-LLM organizes attribution around shared gradient operations and capture, pairing exact cost-selected representations with interfaces for existing training loops and extensions.

3.1 ARCHITECTURE OVERVIEW

  • A four-tier architecture connects gradient storage and operations to shared attribution logic and closed-loop or open-loop entry points.Attributor returns scores for pre-built methods; HookManager captures gradients from backward() inside a with block.Source: §3.1 · p. 4

3.2 EFFICIENCY: EXACT GRADIENT REPRESENTATIONS WITH FLOP-AWARE ROUTING

  • Exact factorized and materialized gradients trade off storage and operation costs, so FLOP-aware routing selects a representation per layer and operation.Factorized storage uses fewer values when T < NiNo/(Ni + No); short sequences and wide layers favor factorization, while long sequences and large comparison batches favor materialization.Source: §3.2 · pp. 4–5
  • Optional dimension reduction and optimizer preconditioning transform gradients before scoring, supporting projected representations and optimizer-aware attribution methods.Factor and materialized projections are supported; the Adam transform enables LESS and AdamW-influence.Source: §3.2 · pp. 4–5
  • Factorized token-level gradients support exact inner products without materializing weight gradients and retain token-granularity attribution.They also support preconditioned comparisons used in K-FAC and EK-FAC.Source: §3.2 · p. 5

3.3 COMPATIBILITY: NON-INVASIVE GRADIENT CAPTURE FROM ARBITRARY TRAINING LOOPS

  • HookManager captures per-example gradients during existing backward passes without modifying the training loop, and accommodates DDP and FSDP.Capture can be wrapped around training calls and gradients may be cached for online or later scoring.Source: §3.3 · pp. 5–6

3.4 EXTENSIBILITY: PLUGGABLE INTERFACES OVER A SINGLE CAPTURE PATH

  • Attributor lets new gradient-based methods consume a common gradient stream and shared operations through a single subclass.Built-in TracIn, K-FAC, EK-FAC, and data-value-embedding attributors use the same interface.Source: §3.4 · p. 6
  • HookManagerCallback exposes capture events so applications can inspect or modify training state, including selecting data before the optimizer step.A built-in callback scores the current batch against a target gradient and removes low-scoring examples’ contributions from parameter gradients.Source: §3.4 · p. 6

4 ATTRIBUTING LLMS AT SCALE WITH HIGH EFFICIENCY

Across model scales, dattri-LLM improves attribution throughput and extends feasible workloads, while its methods offer favorable fidelity–cost trade-offs. Evaluations include scaling on Qwen models and fidelity comparisons against trajectory-specific leave-one-out retraining.

4.1 SCALING ATTRIBUTION WITH MODEL SIZE

  • The scaling benchmark compares GRADDOT, K-FAC, and EK-FAC on Qwen models from 0.5B to 110B parameters using 512-token WikiText-103 sequences.Throughput is measured on four H200 GPUs with each library’s supported distributed mode and largest per-device batch within benchmark limits.Source: §4.1 · p. 6
  • At matched compression and batch size, dattri-LLM averages 3.8× LogIX throughput for K-FAC and EK-FAC from 0.5B to 32B, and supports all three methods through 110B.For GRADDOT, it achieves 2.7–3.5× Kronfluence throughput and 1.4–8.9× Bergson throughput across their respective completed scales.Source: §4.1 · p. 6, Figure 2 · p. 7
  • At 3B, dattri-LLM uses approximately 17.5–20.9% of Kronfluence’s GPU memory and 14.2% of Bergson’s for K-FAC and EK-FAC.The reduced attribution-specific memory footprint leaves more capacity for the model and enables larger workloads within available device memory.Source: §4.1 · pp. 6–7

4.2 ATTRIBUTION FIDELITY AND COMPUTATIONAL COST

  • dattri-LLM offers a favorable fidelity–cost trade-off across GPT-2, Qwen2.5-1.5B, and OLMo-3-7B, reaching the empirical Pareto frontier in two of three panels.On GPT-2, full-dimensional AdamW-influence has 0.845 correlation versus MAGIC’s 0.833 while requiring less time.Source: §4.2 · p. 8, Figure 3 · p. 8
  • Fidelity is measured against TSLOO using 50 removed-sequence retraining runs and 64 validation queries; larger AdamW-influence masks improve correlation at approximately 2–3× longer attribution time.The evaluation uses GPT-2, Qwen2.5-1.5B, and OLMo-3-7B training trajectories.Source: §4.2 · pp. 7–8

5 INTEGRATING AND EXTENDING ATTRIBUTION WORKFLOWS

The section demonstrates attribution-guided data selection as a training-time application of dattri-LLM. It describes the selection setup and its integration into an otherwise unchanged training loop.

5 INTEGRATING AND EXTENDING ATTRIBUTION WORKFLOWS

  • At each training step, GRADDOT scores compare current-batch samples with a held-out validation set, and the lowest-scoring samples are dropped from the gradient update.Source: §5 · p. 8
  • The experiment fine-tunes Llama-3.2-1B on Alpaca for SAMSum, drops the bottom half of each batch, and compares full training, globally ranked subsets, and layer-wise subsets.The variants are tested with both full-parameter and LoRA fine-tuning.Source: §5 · p. 8
  • Per-layer selection achieves the best SAMSum test perplexity under both full-parameter and LoRA fine-tuning, and both selection variants improve on training with all data.Perplexity is averaged over three seeds, with bands showing ±1 standard deviation.Source: Figure 4 · p. 9
  • A DataSelectionCallback specifies the target loader and selection rule; HookManager captures per-sample gradients, and hm.collect() wraps the unchanged training loop so the callback can remove dropped samples’ contributions before the optimizer step.For LoRA, the same procedure uses a layer selector restricted to adapters.Source: §5 · p. 9

6 CONCLUSION

The conclusion presents dattri-LLM as an efficient, compatible, and extensible library for training data attribution at LLM scale. The appendix details its design, including a canonical gradient factorization for supported layers.

6 CONCLUSION

  • dattri-LLM is 3.2× faster than the fastest existing libraries on average.Source: §6 · p. 9
  • The library uses exact factorized or materialized per-example gradients with FLOP-aware routing, captures gradients from unmodified pipelines including DDP and FSDP, and supports applications through callbacks.The conclusion says routing shares kernels across attribution methods and supported layers without changing scores.Source: §6 · p. 9

A DESIGN DETAILS

  • The appendix covers canonical factorization, dimension-reduction maps, optimizer-preconditioned capture, routing, capture, and storage and streaming interfaces.Source: A DESIGN DETAILS · p. 14

A.1 A CANONICAL FACTORIZATION FOR ALL SUPPORTED LAYERS

  • Supported layer types are expressed through an unfolding map that puts flattened gradients into a common factorized form, allowing one implementation of each operation and attribution method to cover them.The listed types include linear, convolutional, embedding, and normalization layers; other trainable parameters can be captured in materialized form with a batch-level parameter-gradient hook.Source: §A.1 · p. 14

A.2 DIMENSION REDUCTION

These appendix sections explain how dattri-LLM reduces, preconditions, routes, captures, and exposes per-example gradients. The design supports attribution methods and workflows while choosing representations according to operation and layer conditions.

A.2 DIMENSION REDUCTION

  • Factor projection applies separate maps to the two gradient factors inside the backward hook, while materialized projection applies one map to the flattened gradient; normalization layers support only the latter.The random maps include Gaussian, Rademacher, sparse Johnson–Lindenstrauss, and sparsify-then-project maps; a random mask preserves coordinate identity for optimizer preconditioning.Source: §A.2 · p. 14

A.3 OPTIMIZER-PRECONDITIONED CAPTURE

  • Optimizer preconditioning recovers selected per-example gradient entries from factors and applies the optimizer update rule before buffering, supporting coordinate-wise optimizers in torch.optim.For trajectory methods, callbacks record moments before and after updates to recover the batch gradient consumed by the optimizer.Source: §A.3 · p. 15

A.4 FLOP-AWARE ROUTING

  • Routing compares operation counts per layer and operation, selecting the cheaper factorized or materialized route for cross-Gram matrices and per-example norms without changing scores in exact arithmetic.Factorized storage is smaller when T < NiNo/(Ni + No); after factor projection, projected factors are better materialized for realistic sequence lengths.Source: §A.4 · p. 15

A.5 CAPTURE MECHANISM

  • HookManager uses forward and backward hooks to capture factors, infers step completion from participating layers, and handles repeated layer calls, checkpointing, and extra validation passes.DDP and FSDP use different step-completion mechanisms, while each rank writes its records independently.Source: §A.5 · p. 16, Figure 6 · p. 17

A.6 STORAGE, IDENTITY, AND THE CAPTURE–ATTRIBUTION INTERFACE

  • Captured gradients can be keyed by input hashes and consumed through disk-backed or live sources; the same attributor supports on-the-fly scoring or later attribution from a cached store.A store-then-attribute workflow can rerun methods and settings over cached gradients without the model or another backward pass.Source: §A.6 · pp. 16–17

B EXPERIMENT DETAILS

The appendix specifies the workload, library configurations, timing procedure, and memory protocols used for the throughput and scaling evaluations.

B.1 THROUGHPUT UNDER A FIXED HARDWARE BUDGET

  • The throughput workload uses Qwen models on WikiText-103 at sequence length 512 in bf16, scoring 512 training sequences against one query through 32B and four queries at 72B and 110B.Each run uses four H200 GPUs, 32 CPU cores, and 256 GB of host memory; attribution includes attention and MLP linear layers but excludes embeddings and the language-model head.Source: §B.1 · p. 18
  • dattri-LLM, Bergson, and Kronfluence use FSDP, while LogIX uses one DDP model replica per device; factor projection at ka = kg = 64 is requested where supported.The comparison reflects both evaluated implementation capabilities and execution efficiency, with each library tested at its largest successful power-of-two per-device batch size.Source: §B.1 · p. 18
  • Timed throughput covers gradient capture, required Kronecker-factor estimation and eigendecomposition, and scoring after two warm-up steps, while excluding model loading, data preparation, and process initialization.Memory scaling at batch size one uses one H200, or four when a library supports sharding; the protocol marks out-of-memory configurations and reports aggregate memory for completed sharded runs.Source: §B.1 · p. 18

B.2 FIDELITY EVALUATION DETAILS

The fidelity evaluation compares attribution methods across GPT-2, Qwen2.5, and OLMo-3 models, using TSLOO loss changes as a reference. It reports fidelity alongside attribution time and memory under a specified hardware budget.

B.2 FIDELITY EVALUATION DETAILS

  • The evaluation covers GPT-2 (124M), Qwen2.5-1.5B, and OLMo-3-7B, trained for one epoch on 512 WikiText-2 blocks.Each block contains 128 tokens; training uses batch size 32 for 16 AdamW updates.Source: §B.2 · pp. 18–19
  • Throughput configurations use four H200 GPUs, while the reported peak-memory values reflect method-specific batch sizes rather than a fixed-batch comparison.Table 4 covers 512 training sequences, with one query through 32B and four queries at 72B and 110B.Source: Table 4 · p. 19, Table 5 · p. 19
  • Fidelity is measured as Spearman correlation between attribution scores and TSLOO loss changes across 50 removed training blocks and 64 validation queries.Each removal replays training without one selected block, and correlations are computed per validation block before averaging.Source: §B.2 · p. 19
  • The compared methods include dattri-LLM AdamW-influence and EK-FAC variants, plus MAGIC, SOURCE, TrackStar, and EK-FAC from Bergson.Some trajectory-based methods cannot run under OLMo-core’s trainer, while TrackStar scores OLMo-3-7B through its HuggingFace export.Source: §B.2 · pp. 19–20
  • Each method is evaluated against the same 64 queries and training trajectory, with attribution time reported separately from training, model loading, and TSLOO reference construction.The resource budget is one B200, 16 CPU cores, 256 GB host memory, and 1 TB local disk; Table 7 records limiting resources and Table 8 reports fidelity, time, and peak GPU memory.Source: §B.2 · pp. 20–21, Table 7 · p. 21, Table 8 · p. 21

B.3 FIXED-BATCH COMPARISONS AT A FIXED MODEL SIZE

Matched-batch tests compare library runtime and memory on Pythia-410M across projected and full-dimensional gradients. Dattri-LLM leads most runtime comparisons, particularly for curvature-based methods.

B.3 FIXED-BATCH COMPARISONS AT A FIXED MODEL SIZE

  • The matched-batch benchmark uses Pythia-410M, 1,024 WikiText-103 sequences, one or 16 queries, and batch size eight on one A40.It compares factor-projected gradients with ka = kg = 64 and full-dimensional gradients; Table 9 reports attribution time and peak GPU memory.Source: §B.3 · p. 21, Table 9 · p. 22
  • Dattri-LLM has the lowest runtime in nine of Table 9’s twelve rows and in every curvature-based row.Projected K-FAC and EK-FAC run 1.6–2.5× faster than LogIX and 3.3–4.1× faster than Bergson; full-dimensional variants run 1.4–1.6× faster than Kronfluence and 1.2–2.1× faster than Bergson.Source: §B.3 · p. 21, Table 9 · p. 22
  • For GRADDOT, dattri-LLM, Bergson, and Kronfluence runtimes differ by at most two seconds, while LogIX is slower.A forward and backward pass over the training sequences accounts for most runtime; curvature methods add fitting and query preconditioning work.Source: §B.3 · p. 21
  • Adaptive routing closely tracks the faster fixed configuration while keeping peak GPU memory lowest or tied at each evaluated sequence length.Figure 7 compares mean attribution time and peak GPU memory as sequence length increases.Source: Figure 7 · p. 22

B.4 ROUTING ACROSS SEQUENCE LENGTHS

Across sequence lengths, adaptive routing stays near the faster fixed route while maintaining low memory use. The section also illustrates token-level attribution and a capture mode that avoids unused weight gradients for frozen-model attribution.

B.4 ROUTING ACROSS SEQUENCE LENGTHS

  • At T = 32, always-materialized execution uses 4.8× adaptive routing’s peak memory, while adaptive routing has the lowest or tied memory at every tested length.Its maximum runtime excess over the faster fixed configuration is approximately 2.8%; materialized execution takes 16–50% longer for 32 ≤ T ≤ 512, and factorized execution takes approximately 17% longer at T = 2048.Source: §B.4 · p. 23

B.5 TOKEN-LEVEL ATTRIBUTION EXAMPLES

  • Token-level scores sum contributions across attributed layers, and their sum recovers the sequence-level score.K-FAC and EK-FAC use preconditioned query gradients reshaped as matrices; each figure panel is normalized independently, and the examples do not establish this decomposition for every library method.Source: §B.5 · p. 23

B.6 WEIGHT-GRADIENT-FREE CAPTURE FOR FROZEN ATTRIBUTION

  • Weight-gradient-free capture returns only hooked layers’ input gradients, avoiding parameter gradients unused for fixed-checkpoint attribution.Standard capture takes 7–19% longer and uses at most 4.5 GB more peak memory; the weight-gradient-free mode is applicable only outside training.Source: §B.6 · p. 23

B.5 TOKEN-LEVEL ATTRIBUTION EXAMPLES

  • In the maskray example, GRADDOT ranks unrelated subwords above “shrimp,” while K-FAC and EK-FAC emphasize “of caridean,” the phrase completed by the target.The example compares token positions for one training sample and target.Source: Figure 8 · p. 24

B.7 DATA SELECTION DETAILS

The data-selection experiment uses GRADDOT to remove low-scoring examples from training batches. Under full-parameter fine-tuning, the library’s selection overhead is lower than that of a dedicated implementation.

B.7 DATA SELECTION DETAILS

  • Selection removes the bottom half of each batch according to GRADDOT scores against one example from a 16-example SAMSum validation pool.The source pool is Alpaca, and the target task is SAMSum.Source: §B.7 · p. 24
  • The experiment uses Llama-3.2-1B with batch size eight over three seeds and evaluates test perplexity on response tokens from 500 held-out SAMSum examples.Prompt tokens are masked during evaluation, which occurs every 50 steps.Source: §B.7 · p. 24
  • Under full-parameter fine-tuning, standard training takes 89 ms per step, dattri-LLM selection takes 146 ms, and the dedicated implementation takes 154 ms.Dattri-LLM is approximately 5% faster; the dedicated implementation adds a forward pass for training gradients, while dattri-LLM captures them from the training backward pass.Source: §B.7 · p. 24, Figure 11 · p. 27
Loading 2609.38767v1…