🤖 machine learning

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

LoRIF introduces a low-rank approximation method for influence functions that significantly reduces storage and memory costs while maintaining high attribution quality, thereby enabling scalable training data attribution for large-scale models and datasets.

Original authors: Shuangqi Li, Hieu Le, Jingyi Xu, Mathieu Salzmann

Published 2026-05-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shuangqi Li, Hieu Le, Jingyi Xu, Mathieu Salzmann

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a giant, super-smart chef (the AI model) who has cooked millions of meals (trained on millions of examples) to learn how to make the perfect dish. Now, you ask the chef to make a specific new dish. You want to know: "Which specific ingredients from the millions of past meals actually influenced this new dish the most?"

This is the problem of Training Data Attribution. The paper introduces a new tool called LoRIF (Low-Rank Influence Functions) to answer this question, even when the "kitchen" is impossibly huge.

Here is how LoRIF works, explained through simple analogies:

The Problem: The "Library of Everything" is Too Big

Previous methods tried to solve this by keeping a massive, detailed index of every single ingredient used in every past meal.

The Bottleneck 1 (Storage): Imagine trying to carry a library of millions of books in your pocket. To find the right book, you have to load the whole thing into your hands every time you ask a question. This takes forever and requires a huge backpack (memory).
The Bottleneck 2 (The Math): To figure out the influence, you have to do a complex calculation involving a giant grid of numbers (the "Hessian"). If the library has 1 million books, this grid is 1 million by 1 million. Storing that grid would fill up the entire internet's hard drives.

Because of this, scientists had to choose: either use a tiny, useless library (low quality) or a massive one that crashes your computer (high cost).

The Solution: LoRIF's Two Magic Tricks

LoRIF says, "We don't need to carry the whole book, and we don't need to draw the whole grid." It uses two clever shortcuts based on the fact that AI gradients (the math behind learning) have hidden, simple patterns.

Trick 1: The "Summary Card" (Rank-c Factorization)

Instead of storing the full, detailed recipe for every single past meal, LoRIF realizes that most recipes share the same core structure.

The Analogy: Imagine instead of storing a 50-page detailed recipe for "Spaghetti Bolognese," you just store a small index card that says: "Tomato base, ground beef, simmer for 2 hours."
How it helps: You can reconstruct the full recipe from this small card whenever you need it. This shrinks the storage size from a massive warehouse to a small filing cabinet. It also means you don't have to load a heavy book into your hands to answer a question; you just grab a tiny card.

Trick 2: The "Spotlight" (Truncated SVD)

When calculating how much a past meal influenced the new one, the math usually requires looking at every single direction in the data.

The Analogy: Imagine a dark room with a million light switches. Most of them are off or very dim. Only a few switches (maybe 10 or 20) are actually turned on and bright enough to matter.
How it helps: LoRIF uses a "spotlight" to find those few bright switches and ignores the million dim ones. Instead of calculating the effect of all million switches, it only calculates the effect of the top 20. This turns a math problem that would crash a supercomputer into one that fits on a laptop.

The Result: Fast, Cheap, and Accurate

By combining these two tricks, LoRIF achieves something that was previously thought impossible:

It fits: It can handle models with 70 billion parameters (like a massive brain) and datasets with millions of examples.
It's fast: It answers questions 20 times faster than previous methods because it doesn't have to load massive files.
It's accurate: Even though it uses "summary cards" and a "spotlight," it finds the right ingredients just as well as (or better than) the old methods that tried to store everything.

Why This Matters (According to the Paper)

The paper claims this makes it possible to debug and audit giant AI models.

Debugging: If an AI says something weird, you can instantly trace it back to the specific training examples that taught it that behavior.
Safety Auditing: You can find if the AI learned harmful behavior from a specific, bad example in its training data, even if that example was buried in millions of others.
Data Curation: You can identify which training data is actually useful and which is just noise.

In short, LoRIF turns a task that required a supercomputer and a warehouse of storage into something that can be done efficiently, allowing us to finally understand the "ingredients" behind the world's largest AI models.

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

Problem Statement

Training Data Attribution (TDA) aims to identify which specific training examples most influenced a model's prediction on a given query. While Influence Functions (IF) provide a theoretically grounded approach to TDA by approximating leave-one-out retraining via gradients and inverse Hessian matrices, they have remained effectively inaccessible for modern-scale models (billions of parameters) trained on millions of examples.

Existing scalable variants, such as TRAK, LoGRA, and TrackStar, address computational bottlenecks by storing projected per-example gradients and using a damped Gauss-Newton Hessian approximation. However, these methods face two critical scalability bottlenecks that create a tradeoff between attribution quality and computational feasibility:

Gradient Storage and I/O: Storing projected per-example gradients for $N$ training examples requires $O(ND)$ space, where $D$ is the effective projection dimension. At query time, loading these gradients dominates latency (approximately 96% of query time).
Inverse Hessian Scaling: Forming and storing the $D \times D$ inverse Hessian approximation requires $O(D^2)$ memory.

Increasing the projection dimension $D$ is necessary to improve attribution quality, but doing so exacerbates both storage/I/O and memory bottlenecks, rendering high-quality attribution impractical for frontier-scale models.

Methodology: LoRIF

LoRIF (Low-Rank Influence Functions) addresses these bottlenecks by exploiting the low-rank structure inherent in neural network gradients. The method follows a three-step pipeline:

1. Low-Rank Gradient Storage (Addressing I/O Bottleneck)

Instead of storing the full projected per-example gradient matrix $\tilde{G}_i \in \mathbb{R}^{d_1 \times d_2}$ (where $D = d_1 d_2$ ), LoRIF stores a rank- $c$ factorization:
$\tilde{G}_i \approx u_i v_i^\top$
where $u_i \in \mathbb{R}^{d_1 \times c}$ and $v_i \in \mathbb{R}^{d_2 \times c}$ .

Mechanism: This factorization is computed via block power iterations.
Impact: Storage and query-time I/O per layer per sample are reduced from $O(D)$ to $O(c(\sqrt{D}))$ . Empirical results suggest that even with $c=1$ , the approximation retains meaningful attribution quality, and prioritizing a larger $D$ over a higher $c$ yields better cost-efficiency.

2. Truncated SVD Curvature Approximation (Addressing Memory Bottleneck)

To avoid forming the dense $D \times D$ inverse Hessian $(G^\top G + \lambda I)^{-1}$ , LoRIF computes a rank- $r$ truncated Singular Value Decomposition (SVD) of the aggregate gradient matrix $G$ :
$G \approx U_r \Sigma_r V_r^\top$
Using the Woodbury matrix identity, the inverse Hessian term is approximated in an $r$ -dimensional subspace:
$(V_r \Sigma_r^2 V_r^\top + \lambda I)^{-1} = \frac{1}{\lambda}I - \frac{1}{\lambda^2} V_r \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} V_r^\top$

Mechanism: Only the top- $r$ singular vectors ( $V_r$ ) and singular values ( $\Sigma_r$ ) are stored.
Impact: Memory complexity for the curvature term is reduced from $O(D^2)$ to $O(Dr)$. This allows for a significantly larger effective projection dimension $D$ within a fixed memory budget.

3. Influence Score Computation

At query time, influence scores are computed by projecting the query gradient into the $r$ -dimensional subspace and utilizing the stored low-rank factors for training gradients. The influence score $I(x_{tr}, x_{te})$ is approximated as:
$I(x_{tr}, x_{te}) \approx \frac{1}{\lambda} g_{te}^\top g_{tr} - \frac{1}{\lambda^2} (g'_{te})^\top \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} g'_{tr}$
where $g'$ represents gradients projected onto the singular vectors $V_r$ . The first term is computed efficiently using the rank- $c$ factors, and the second term involves only $O(r)$ operations due to the diagonal nature of the middle matrix.

Key Contributions

Identification of Bottlenecks: The paper identifies that the projection dimension $D$ is the primary constraint in scalable IF-style TDA, creating a quality-scalability tradeoff where accurate attribution (large $D$ ) is currently impractical.
Rank- $c$ Factorization: The proposal to store projected gradients as rank- $c$ factors, reducing per-example per-layer storage from $O(D)$ to $O(c\sqrt{D})$ .
Rank- $r$ Curvature Approximation: The proposal to use truncated SVD and the Woodbury identity to approximate the curvature term, reducing per-layer memory from $O(D^2)$ to $O(Dr)$.
Scalability Demonstration: The first demonstration of effective data attribution on models up to 70 billion parameters (Apertus-70B) with millions of training examples, maintaining practical cost and latency.

Experimental Results

LoRIF was evaluated on models ranging from 124M to 70B parameters (GPT2-small, OLMo-3-7B, Apertus-70B) with datasets containing up to 3.8 million examples.

Storage and Latency: Compared to LoGRA, LoRIF achieves 2.3–20× lower gradient storage and 1.3–20× lower query latency while matching or exceeding attribution quality.
- On GPT2-small, LoRIF reduced storage by up to 9× and latency by 5.6× in medium storage regimes.
- On OLMo-3-7B, LoRIF achieved nearly identical tail-patch scores to LoGRA with 20.3× lower storage and 22.0× lower latency.
- On Apertus-70B, LoRIF improved tail-patch scores by 1.5–2.2× over LoGRA with lower storage and latency.
Attribution Quality: Measured via Linear Datamodeling Score (LDS) and Tail-patch score, LoRIF consistently improved the Pareto frontier. By enabling larger effective projection dimensions $D$ (which were previously infeasible due to memory constraints), LoRIF achieved higher quality than baselines constrained to smaller $D$ .
Qualitative Evaluation: An LLM-as-a-judge evaluation on top-1 retrieved examples showed that LoRIF retrieves examples with higher semantic relevance (average score 3.04 vs. 2.49 on OLMo-3-7B) and fewer completely irrelevant retrievals compared to LoGRA.

Significance and Claims

The paper claims that LoRIF makes gradient-based Training Data Attribution practical at frontier scale. By exploiting the low-rank structure of gradients, LoRIF resolves the quality-scalability tradeoff that has hindered the application of influence functions to large language models.

The authors emphasize that LoRIF does not rely on new low-rank primitives but rather demonstrates how existing low-rank structures (observed in both per-example and aggregate gradients) can be systematically exploited in projected-gradient influence methods. The results suggest that gradient-based TDA can now be used for large-scale LLM auditing, debugging, and safety checks when the attribution index is reused across many queries, without requiring the prohibitive storage or I/O costs of previous methods.

The paper acknowledges limitations, noting that constructing the attribution index still requires a gradient pass over the corpus (preprocessing cost) and that current large-scale experiments focus on supervised fine-tuning data rather than full pre-training corpora. However, the work establishes a new baseline for scalable attribution, enabling the analysis of models previously considered out of reach for such methods.