LoRIF: Pushing the Scalability Frontier of Training Data Attribution

Shuangqi Li; Hieu Le; Jingyi Xu; Mathieu Salzmann

cs.LGJan 29, 2026

LoRIF: Pushing the Scalability Frontier of Training Data Attribution

Authors: Shuangqi Li, Hieu Le, Jingyi Xu, Mathieu Salzmann

Organizations: EPFL, Switzerland · UNC Charlotte, USA · Stony Brook University, USA

Abstract

Training data attribution (TDA) identifies which training examples most influenced a model's prediction. Influence function methods are a theoretically grounded family of TDA methods and exploit gradients. To overcome the scalability challenge arising from gradient computation, the most popular strategy is random projection (e.g., TRAK, LoGRA). However, this still faces two bottlenecks when scaling to large training sets and high-quality attribution: (i) storing and loading projected per-example gradients for all $N$ training examples, where query latency is dominated by I/O; and (ii) forming the $D \times D$ inverse Hessian approximation, which costs $O(D^2)$ memory. Both bottlenecks scale with the projection dimension $D$ , yet increasing $D$ is necessary for attribution quality -- creating a quality--scalability tradeoff. We introduce LoRIF (Low-Rank Influence Functions), which exploits low-rank structures of gradient to address both bottlenecks. First, we store rank- $c$ factors of projected per-example gradients rather than full matrices, reducing storage and query-time I/O from $O(D)$ to $O(c\sqrt{D})$ per layer per sample. Second, we use truncated SVD with the Woodbury identity to approximate the inverse Hessian term in an $r$ -dimensional subspace, reducing memory from $O(D^2)$ to $O(Dr)$ . On models from 0.1B to 70B parameters trained on datasets with millions of examples, LoRIF achieves up to 20 $\times$ storage reduction and query-time speedup compared to LoGRA, while matching or exceeding its attribution quality. LoRIF makes gradient-based TDA practical at frontier scale. Code is available at https://github.com/doub7e/LoRIF.

Figures & tables

Figure 1: Method Overview. ( Top ): LoGRA-style attribution faces two scaling bottlenecks: storing projected gradients with

\mathcal{O}(D_{\ell})

space per example per layer, and forming the dense

D_{\ell}{\times}D_{\ell}

curvature approximation with

\mathcal{O}(D_{\ell}^{2})

memory. LoRIF targets each bottleneck with a complementary low-rank step. ( Mid-left ): Rank-

c

factorization compresses each projected per-example gradient, reducing persistent storage and query-time I/O. ( Mid-right ): A rank-

r

truncated SVD of the per-layer training-gradient matrix

G

reduces curvature memory to

\mathcal{O}(D_{\ell}r)

. ( Bottom ): Combining both compressions and applying the Woodbury identity result in a more scalable influence scoring method, making larger projection dimensions practical under storage, I/O and memory constraints.

Panel of Figure 2 — Figure 2: Effects of low-rank approximations on attribution quality. Linear Datamodeling Score (LDS) is the Spearman correlation between additive attribution predictions and outputs of models retrained on data subsets; higher is better. ( Left ): LoRIF trades a larger effective projection dimension $D$ against rank- $c$ factorization; LoGRA (no factorization) is compared against rank- $c$ factorization, varying $D$ by setting $d_{1}=\frac{I}{f}$ and $d_{2}=\frac{O}{f}$ for $f\in\{64,32,16,8\}$ (so $D=\frac{I\cdot O}{f^{2}}$ ). ( Right ): truncated-SVD curvature approximation approaches the full-rank baseline with $r\ll D$ ; $r=0$ means discarding the curvature information and reduces to dot product; rank factorization is not used. Both experiments use GPT2-small on WikiText-103.

Figure 3: Time spent in loading gradients and GPU computation. Query-time latency breakdown for different methods, with the same effective projection dimension

D=\frac{I\cdot O}{8^{2}}

and truncation rank

r=\frac{D}{16}

for each layer of GPT2-small. “Ours” combines rank-1 factorization and truncated SVD.

Panel of Figure 4 — Figure 4: LoRIF (Ours) vs. LoGRA under different storage budgets. ( Left ): Attribution quality measured by LDS on GPT2-small. We vary the storage cost by choosing different effective projection dimensions $D$ , where $D=\frac{I\cdot O}{f^{2}}$ , and factorization ranks $c$ (only for LoRIF). For LoGRA, we choose $f\in\{64,32,16,8\}$ . For LoRIF, we use $f\in\{32,16,8,4\}$ when $c=1$ , and $c\in\{1,4,16,32\}$ when $f=4$ . ( Right ): Attribution quality measured by tail-patch score on OLMo3-7B. For LoGRA, we choose $f\in\{360,256,180,128\}$ . For LoRIF, we use $f\in\{128,64,32,16\}$ and $c=1$ .

Contextual baselines
Method	$f$	$c$	$r$	LDS $\uparrow$	Storage $\downarrow$	Latency $\downarrow$
EK-FAC	–	–	–	$0.3936\pm 0.013$	2.1 GiB	20 hr
RepSim	–	–	–	$0.0103\pm 0.009$	340.7 MiB	0.15 s
High storage regime
GradDot	8	–	–	$0.0576\pm 0.010$	574.9 GiB	209.8 s
TrackStar	8	–	–	$0.2623\pm 0.012$	574.9 GiB	241.1 s

Table 1: Main comparison on GPT2-small. We control the persistent storage cost and the query latency for a query batch by varying the hyperparameters

(f,c,r)

. Note that EK-FAC stores curvature statistics only and recomputes training gradients. Bold values are the best observed values within each storage panel.

Method	$f$	$c$	$r$	Tail-patch (%) $\uparrow$	Storage $\downarrow$	Latency $\downarrow$
OLMo-3-7B
RepSim	–	–	–	$0.144\pm 0.04$	16.99 GiB	7.26 s
GradDot	128	–	–	$0.411\pm 0.04$	1638.6 GiB	781.4 s
LoGRA	128	–	–	$0.546\pm 0.04$	1638.6 GiB	782.1 s
LoRIF	128	1	$2^{8}$	$0.495\pm 0.05$	80.9 GiB	35.6 s
LoRIF	16	1	$2^{8}$	$\mathbf{0.986\pm 0.07}$	647.4 GiB	469.7 s

Table 2: Large-scale attribution results. Tail-patch score is reported on OLMo-3-7B and Apertus-70B, where repeated subset retraining for LDS is infeasible.

	Avg. relevance $\uparrow$		Preference (%)
Model	LoRIF	LoGRA	LoRIF	LoGRA	Tie
OLMo-3-7B	3.04	2.49	35.1	8.0	56.9
Apertus-70B	2.91	2.24	40.8	9.1	50.1

Table 3: LLM-as-a-judge top-1 retrieval evaluation. The judge scores on a 1–5 relevance scale. Preference reports the fraction of queries where one method is judged better or two are tied.

Appendix figures & tables26 assets

Supplementary material from the paper’s appendix.

Appendix

Family / method	What the score uses	Extra training information	Per-example artifact stored for repeated queries
Subset-based valuation / simulation	Marginal effects of including or excluding training subsets; sometimes a learned simulator of model behavior from subset membership.	Usually many subset-trained models or evaluations; in-run variants instead require logging during the training run.	Usually no query-specific gradient store; may store global values or a learned simulator.
Trajectory-aware attribution	Gradients or parameter changes accumulated along the optimization path.	Saved checkpoints, update history, or an unrolled/differentiable training trajectory; no new retraining if the trajectory was stored.	Sometimes checkpoint-gradient caches; not a final-checkpoint-only representation.
Representation or text similarity	Text overlap, embedding similarity, or hidden-state similarity between training and query examples.	No retraining or saved training trajectory.	Text, embedding, or representation store for nearest-neighbor retrieval.
Gradient similarity / Hessian-free IF	Dot products between training and query gradients, often using selected layers or an identity curvature approximation.	No retraining or saved training trajectory.	Optional full- or layer-gradient cache; otherwise training gradients must be streamed or recomputed.
Parameter-space IF / iHVP	Training and query gradients preconditioned by an inverse-Hessian-vector product or curvature approximation.	No retraining or saved training trajectory, but requires an iHVP or curvature solver.	Optional full- or layer-gradient cache; many implementations stream or recompute training gradients.
Projection-based IF	Projected per-example gradients, optionally with curvature or preconditioning.	No saved training trajectory; typically computed from one or a few final checkpoints.	Full projected-gradient store for all training examples.

Appendix

Table 4: A coarse taxonomy of training-data attribution methods. The table summarizes information requirements and typical stored representations.

Figure 5: LDS vs. tail-patch score on GPT2-small. Each point is a method–configuration pair. The dashed line is a linear fit. RepSim (non-gradient-based) deviates most from the trend.

Method	$f$	$c$	$r$	Stage 1	Stage 2	Total
EK-FAC	–	–	–	7 hr	–	7 hr
LoGRA	64	–	–	15 min	1 s	15 min
LoGRA	32	–	–	15.5 min	6 s	16 min
LoGRA	16	–	–	18 min	1 min	19 min
LoGRA	8	–	–	23.5 min	13 min	37 min
LoRIF	32	1	$2^{10}$	16 min	12 s	16 min

Appendix

Table 5: Preprocessing time on GPT2-small (233k training examples). EK-FAC has only Stage 1 (no stored gradients). “–” indicates not applicable.

Method	$f$	$c$	$r$	Stage 1	Stage 2	Total
LoGRA	360	–	–	68 hr	2 min	68 hr
LoGRA	256	–	–	67 hr	4 min	67 hr
LoGRA	180	–	–	68 hr	6.5 min	68 hr
LoGRA	128	–	–	67.5 hr	13.5 min	68 hr
LoRIF	128	1	$2^{8}$	67.5 hr	2 min	68 hr
LoRIF	64	1	$2^{8}$	68 hr	9 min	68 hr

Appendix

Table 6: Preprocessing time on OLMo3-7B (2.2M training examples).

Method	$f$	$c$	$r$	Stage 1	Stage 2	Total
LoGRA	512	–	–	180 hr	18 min	180 hr
LoRIF	256	1	$2^{9}$	185 hr	30 min	186 hr
LoRIF	64	1	$2^{9}$	195 hr	5 hr	200 hr

Appendix

Table 7: Preprocessing time on Apertus-70B (3.8M training examples). To be able to compute gradients for this large model, we used two H200 GPUs and applied tensor parallelism.

Setting	$B$	Preprocessing	Latency/batch	$K^{\star}$
GPT2-small, low storage	487	16 / 18 min	16.4 / 2.93 s	8.9
OLMo-3-7B, higher quality	1,000	68 / 72 hr	782.1 / 469.7 s	46.1
Apertus-70B, efficient	1,000	180 / 186 hr	1087.9 / 222.5 s	25.0
Apertus-70B, higher quality	1,000	180 / 200 hr	1087.9 / 843.8 s	295.0

Appendix

Table 8: Preprocessing/query-time break-even. Time pairs are LoGRA / LoRIF.

K^{\star}

counts query batches, not individual queries. The first whole batch reaching the threshold is

\lceil K^{\star}\rceil

Method	$f$	$c$	$r$	LDS $\uparrow$	Storage	Storage/Ex.	Latency
LoRIF w/o truncated SVD	8	1	–	$0.1868\pm 0.013$	8.3 GiB	0.04 MiB	72.7 s
LoRIF w/o truncated SVD	4	16/32	–	OOM	–	–	–
LoRIF w/o rank factorization	4	–	$2^{14}$	$\mathbf{0.3876\pm 0.0116}$	2297.7 GiB	10.1 MiB	1548.0 s
LoRIF	4	16	$2^{14}$	$0.3058\pm 0.013$	255.5 GiB	1.12 MiB	165.9 s
LoRIF	4	32	$2^{14}$	$0.3428\pm 0.012$	511.0 GiB	2.25 MiB	287.1 s
LoRIF w/o truncated SVD	4	1	–	OOM	–	–	–

Appendix

Table 9: Ablation of LoRIF components on GPT2-small. Rank factorization reduces persistent gradient storage and query-time I/O, while truncated SVD removes the curvature-memory bottleneck. Storage denotes the persistent training-gradient storage.

Model	Module	$c=1$ Err.	$c=1$ EVR	$c=4$ Err.	$c=4$ EVR
GPT2-small (124M)	attn	0.787	37.6%	0.557	67.8%
GPT2-small (124M)	mlp	0.834	30.2%	0.711	49.3%
OLMo-3-7B	attn	0.773	38.8%	0.594	63.5%
OLMo-3-7B	mlp	0.827	31.2%	0.694	51.3%
Apertus-70B	attn	0.479	74.3%	0.211	93.8%
Apertus-70B	mlp	0.641	55.0%	0.454	77.1%

Appendix

Table 10: Rank-

c

factorization error of projected per-example gradients. Errors are relative Frobenius reconstruction errors. EVR denotes the fraction of Frobenius energy captured by the rank-

c

approximation.

Model	Module	$D$	EVR@10%	EVR@25%	EVR@50%
GPT2-small, $f=32$	attn	1152	0.49	0.67	0.83
GPT2-small, $f=32$	mlp	2304	0.42	0.60	0.79
OLMo-3-7B, $f=128$	attn	1024	0.43	0.59	0.77
OLMo-3-7B, $f=128$	mlp	2752	0.36	0.52	0.72
Apertus-70B, $f=256$	attn	576	0.77	0.84	0.91
Apertus-70B, $f=256$	mlp	5376	0.53	0.66	0.81

Appendix

Table 11: Spectral concentration of projected training-gradient matrices. EVR@

p\%

denotes the explained variance captured by the top

p\%

singular directions.

Figure 6: Spectral concentration of the projected gradient matrix

G

on GPT2-small. The curve shows the cumulative explained variance ratio

\mathrm{EVR}(r)

from Equation 12 .

Figure 7: LDS vs. truncation rank

r

with rank-

c

factorization. Different curves correspond to different projection dimensions

D

and factorization ranks

c

. LDS saturates at

r\ll D

, especially for small

c

, confirming that truncated SVD remains effective when combined with low-rank gradient storage. Experiments are run on GPT2-small.

Method	$f$	$c$	$r$	Storage	LDS $\uparrow$
Medium storage
FactGraSS-RM	48	–	–	16.6 GiB	$0.0972\pm 0.0098$
Selective-FactGraSS	48	–	–	16.6 GiB	$0.0952\pm 0.0095$
LoRIF	4	1	4096	16.6 GiB	$\mathbf{0.2068\pm 0.0131}$
Low storage
FactGraSS-RM	96	–	–	4.1 GiB	$0.0533\pm 0.0089$

Appendix

Table 12: FactGraSS and LoRIF at matched persistent storage on GPT2-small. Random-mask FactGraSS is denoted FactGraSS-RM.

Regime	Method	Time	Speedup	LDS $\uparrow$
Medium	LoRIF	0.5687	$1.00\times$	$0.2068\pm 0.0131$
	RM-LoRIF	0.4082	$1.39\times$	$0.1990\pm 0.0128$
	SM-LoRIF	0.4083	$1.39\times$	$0.2066\pm 0.0130$
Low	LoRIF	0.2805	$1.00\times$	$0.1388\pm 0.0111$
	RM-LoRIF	0.2080	$1.35\times$	$0.1259\pm 0.0117$
	SM-LoRIF	0.2082	$1.35\times$	$0.1141\pm 0.0115$

Appendix

Table 13: GraSS-style masking combined with LoRIF. Medium uses

f=4,c=1,r=4096

; Low uses

f=16,c=1,r=2048

. Time is in milliseconds per example for the hook operations specified below.

Method	Mean Spearman $\uparrow$	95% query-bootstrap CI
LoGRA, $f=8$	0.3648	$[0.3529,\,0.3767]$
LoRIF, $f=4,c=32,r=256$	0.3966	$[0.3847,\,0.4082]$

Appendix

Table 14: Agreement with pointwise LOO retraining on GPT2-small. Correlations are computed within each query over 50 randomly selected removals and averaged over 487 queries.

Removed	LoRIF $\Delta$ PPL	LoGRA $\Delta$ PPL	LoRIF $-$ LoGRA $\uparrow$
128 examples	$+0.8742$	$+0.8961$	$-0.0219\;[-0.0914,\,0.0484]$
1%	$+2.8675$	$+2.7699$	$+0.0976\;[-0.0579,\,0.2451]$
10%	$+7.8711$	$+7.6341$	$+0.2370\;[-0.0622,\,0.5845]$
25%	$+13.1078$	$+12.7128$	$+0.3950\;[0.0244,\,0.8601]$

Appendix

Table 15: Top-ranked removal on the full GPT2-small training corpus. The middle columns show changes from the all-data control. The final column is the paired method difference with its query-bootstrap interval.

Removed	LoRIF $\Delta$ PPL	LoGRA $\Delta$ PPL	LoRIF $-$ LoGRA $\downarrow$
128 examples	$-0.0749$	$-0.0587$	$-0.0162\;[-0.0516,\,0.0214]$
1%	$-0.6757$	$-0.6019$	$-0.0738\;[-0.1272,\,-0.0196]$
10%	$-2.6508$	$-2.5014$	$-0.1494\;[-0.2204,\,-0.0799]$
25%	$-4.2469$	$-4.1433$	$-0.1036\;[-0.1515,\,-0.0596]$

Appendix

Table 16: Bottom-ranked removal on the full GPT2-small training corpus. Negative changes indicate lower PPL than the all-data control. Negative paired differences favor LoRIF.

Score	Meaning
1	Completely irrelevant.
2	Vaguely related field.
3	Same broad topic.
4	Closely related problem or task.
5	Nearly identical problem or task.

Appendix

Table 17: LLM-as-a-judge scoring rubric. The judge scores the semantic relevance of each retrieved training example to the query.

Model	Avg. score $\uparrow$	Score 1 rate $\downarrow$	Score $\geq 4$ rate $\uparrow$	Preference: LoRIF / LoGRA / Tie
OLMo-3-7B	3.04 / 2.49	17.8% / 31.7%	39.1% / 25.2%	35.1% / 8.0% / 56.9%
Apertus-70B	2.91 / 2.24	16.0% / 41.3%	33.3% / 18.8%	40.8% / 9.1% / 50.1%

Appendix

Table 18: Systematic top-1 retrieval evaluation. We evaluate 1,000 queries per model. Each cell comparing LoRIF and LoGRA is written as LoRIF / LoGRA. Preference reports the fraction of queries where LoRIF is judged better, LoGRA is judged better, or the two are tied.

Score	Meaning	OLMo LoRIF	OLMo LoGRA	Apertus LoRIF	Apertus LoGRA
1	Completely irrelevant	17.8%	31.7%	16.0%	41.3%
2	Vaguely related	17.7%	24.9%	25.7%	20.7%
3	Same broad topic	24.5%	17.3%	25.0%	19.2%
4	Closely related	19.9%	12.2%	18.2%	10.2%
5	Nearly identical	19.2%	13.0%	15.1%	8.6%

Appendix

Table 19: LLM-judge relevance score distributions. Entries report the percentage of top-1 retrieved examples assigned each score.

Method	LLM-judge relevance $\uparrow$	Tail-patch score $\uparrow$
RepSim	2.83	0.144
LoGRA	2.49	0.546
LoRIF	3.04	0.986

Appendix

Table 20: Textual relevance and behavioral influence measure different properties. Results are on OLMo-3-7B. LLM-judge relevance evaluates the decoded top-1 retrieved example. Tail-patch measures whether the retrieved proponents increase the query target probability after one additional training step.

Panel of Figure 8 — Figure 8: Qualitative comparison on OLMo3-7B (Example 1).

Panel of Figure 9 — Figure 9: Qualitative comparison on OLMo3-7B (Example 2).

Panel of Figure 10 — Figure 10: Qualitative comparison on OLMo3-7B (Example 3).

Panel of Figure 11 — Figure 11: Qualitative comparison on Apertus-70B (Example 1).

Panel of Figure 12 — Figure 12: Qualitative comparison on Apertus-70B (Example 2).

Panel of Figure 13 — Figure 13: Qualitative comparison on Apertus-70B (Example 3).

Training Data Attribution

Training data attribution identifies which examples, sources, or groups shape a model’s predictions, capabilities, and failures.

See all 83 papers

dattri-LLM: A Unified and Efficient Library for Training Data Attribution at LLM Scale
Sep 30, 2026Shixuan Liu, Tongli Zhou, Junwei Deng +2Gradient-Based Attribution Training Data Attribution
University of Illinois at Urbana-Champaign · Google
Traceprop: Training Data Attribution in a Single Pass
Sep 26, 2026Amit NautiyalGradient-Based Attribution Kronecker-Factored Approximation
Independent Researcher
STRIDE: Training Data Attribution via Sparse Recovery from Subset Perturbations
Jun 3, 2026Rishit Dagli, Abir Harrasse, Luke Zhang +4Sparse Recovery Data Selection
Jinesis AI Lab, University of Toronto & Vector Institute · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Thoughtworks +3

Sep 30, 2026J/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

Gradient-Based Attribution Training Data Attribution

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

Training Data Attribution

Explore similar work

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

Traceprop: Training Data Attribution in a Single Pass

STRIDE: Training Data Attribution via Sparse Recovery from Subset Perturbations