LoRIF: Low-Rank Influence Functions for Scalable Training Data Attribution

TL;DR

LoRIF tackles the scalability challenge of gradient-based training data attribution by exploiting the low-rank structure of neural-network gradients to store and manipulate projected per-example gradients efficiently and to approximate the inverse Hessian in a low-dimensional subspace via a truncated SVD and the Woodbury identity. This yields substantial storage and compute savings while maintaining or improving attribution quality across models from GPT2-small to 70B-parameter architectures. Compared with projection-based baselines like LoGRA under tight storage conditions, LoRIF delivers superior efficiency and Pareto improvements, enabling practical gradient-based TDA at frontier scale. The work provides both theoretical justification and empirical validation that low-rank gradient techniques can unlock high-quality data attribution for large-scale language models, with implications for data debugging, poisoning detection, and dataset curation.

Abstract

Training data attribution (TDA) identifies which training examples most influenced a model's prediction. The best-performing TDA methods exploits gradients to define an influence function. 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: \emph{(i)} storing and loading projected per-example gradients for all $N$ training examples, where query latency is dominated by I/O; and \emph{(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 \textbf{LoRIF (Low-Rank Influence Functions)}, which exploits low-rank structures of gradient to address both bottlenecks. First, we store rank-$c$ factors of the 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 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.

Figures (14)

• Figure 1: Attribution quality (LDS) vs. effective projection dimension $D$. We compare LoGRA (no factorization) against rank-$c$ factorization for $c \in \{1, 4, 16\}$ on GPT2-small, 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}$).
• Figure 2: Attribution quality (LDS) vs. truncation rank $r$ for the low-rank inverse Hessian approximation, shown for varying effective layer dimensions $D$. Rank factorization is not used. $r=0$ corresponds to discarding the curvature information and degrading to naive dot product. We use GPT2-small and train it on WikiText-103, with the number of training examples $N \gg D$.
• Figure 3: Time spent in loading gradients and GPU computation. We visualize the query-time latency breakdown for different methods, with the same effective projection dimension $D=\frac{I\cdot O}{8^2}$ and a truncation rank $r=\frac{D}{16}$ for each layer of GPT2-small. "Ours" means combining rank-1 factorization and the truncated SVD approximation.
• Figure 4: LoRIF (Ours) vs. LoGRA on attribution quality (LDS), on GPT2-small. We vary the storage cost by choosing different effective projection dimensionss $D$, where $D = \frac{I \cdot O}{f^2}$ and factorization rankss $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$.
• Figure 5: LoRIF (Ours) vs. LoGRA on attribution quality (tail-patch score), on Olmo3-7B. We vary the storage cost by choosing different effective projection dimensionss $D$, where $D = \frac{I \cdot O}{f^2}$. For LoGRA, we choose $f \in \{360, 256, 180, 128\}$. For LoRIF, we use $f \in \{128, 64, 32, 16\}$ and $c=1$.