Executive Summary

Training Data Attribution (TDA) seeks to map specific outputs of a machine learning model back to the exact training examples that caused them, answering the fundamental question of causal influence in artificial intelligence. Classical implicit differentiation techniques, such as influence functions, face insurmountable memory bottlenecks when applied to foundation models with 8 billion to 70 billion parameters, as they require computing an inverse Hessian matrix whose dimensions scale with the square of the model's parameters. To bypass these limits, emergent projection frameworks like LoGra and low-rank adaptations (LoRIF) enable influence tracking across massive pretraining corpora without out-of-memory errors. However, because modern Large Language Models (LLMs) rarely converge to a unique optimal state and are trained across multi-stage pipelines, static influence functions are increasingly being augmented by trajectory-unrolling methods like SOURCE and TracIn. Ultimately, evidence leans toward a future paradigm of "stagewise data attribution," where singular learning theory dictates that the influence of a data point is not static, but heavily dependent on the developmental phase in which the model encounters it.

As the artificial intelligence industry exhausts the supply of high-quality human-annotated data, the focus of architectural design is shifting toward radical data transparency and curation. We are entering an era where understanding exactly what a model learned from which document is no longer a luxury, but a strict requirement for debugging, copyright auditing, and AI alignment. Training data attribution (TDA) methods offer to trace a model's prediction on any given example back to specific influential training examples [cite: 1]. However, mapping the causal influence of a single sentence across a pretraining corpus of trillions of tokens requires navigating immense mathematical and computational bottlenecks. While perfect causal attribution remains an open challenge, research suggests that emerging gradient-based approximations and unrolling architectures provide a practical pathway to dismantling the black box of training dynamics. This report comprehensively deconstructs the evolution of TDA, from static influence functions to the frontier of stagewise unrolling and singular learning theory.

The Epistemological Crisis of Data Valuation

Training data attribution (TDA) techniques are motivated by the fundamental goal of understanding the relationship between the data fed into a neural network and the final properties of the trained model [cite: 2]. At its core, TDA aims to perform a counterfactual prediction: estimating how a trained model's behavior would dynamically change if specific data points were removed from, or added to, the training dataset [cite: 2, 3].

The applications of this counterfactual reasoning are vast. Accurately attributing model behavior allows engineers to detect mislabeled data points that degrade performance, craft or defend against data poisoning attacks, and systematically curate high-fidelity datasets [cite: 3]. Consider a concrete, real-world application: the legal auditing of an LLM for intellectual property infringement, such as the high-profile copyright lawsuit between The New York Times and OpenAI. TDA can be deployed to definitively prove causal linkage. Instead of relying on a human reviewer noting that a model output looks similar to a copyrighted article, TDA provides mathematical proof showing whether a specific article (the input) directly caused the model to regurgitate exact paragraphs (the output). If the influence score of that specific copyrighted article on the memorized output is astronomically high, it proves the model did not spontaneously generate the text, but directly retrieved it from the training corpus.

For a language model, the relevant data points might be entire segments of running text from the pretraining corpus, whereas for a supervised image classification task, they would be labeled input-output pairs [cite: 1]. In a purely theoretical landscape, data attribution could be calculated with perfect accuracy by executing a leave-one-out retraining protocol: simply removing an example from the dataset, re-training the model from scratch, and observing the delta in its behavior [cite: 1].

However, in the context of modern LLMs, which consume astronomical computational resources—such as Meta's Llama 3 8B, which required approximately 1.3 million GPU hours on H100 clusters and consumed massive corpora comprising 15 trillion tokens (approximately 30 to 40 terabytes of plain text) to train—this brute-force approach is computationally impossible. Thus, the field of TDA is defined by the search for efficient mathematical approximations that can simulate this counterfactual reality without the need for actual retraining.

Implicit Differentiation and the Mathematics of Influence Functions

The earliest and most prominent mathematical solution to the TDA bottleneck originated outside of deep learning. Influence functions are a classical idea derived from the field of robust statistics—specifically introduced by Hampel in 1974—which were subsequently adapted for deep learning applications by Koh and Liang in 2017 [cite: 4, 5].

Influence functions fall under the umbrella of implicit-differentiation-based TDA. They utilize the Implicit Function Theorem—a mathematical principle that allows finding the derivative of an implicit relationship without explicitly solving for one variable. To use a simple real-world analogy, it is like determining how adjusting the steering wheel of a car changes its trajectory without having to reverse-engineer the entire mechanical linkage of the steering column. In the context of the TDA bottleneck, this theorem allows engineers to calculate how changing the training data changes the model's final parameters without actually going through the prohibitive process of retraining the model from scratch [cite: 2]. Assuming we have a training dataset where a model fits its parameters using empirical risk minimization, we can parameterize the training set by a specific example's weight [cite: 4]. The influence of that data point on the final parameters is defined as the first-order Taylor approximation of the response function [cite: 4].

Mathematically, the approximated parameter shift when a data point is removed relies on the gradient of the loss function multiplied by the inverse Hessian matrix. These methods provide highly convenient estimation algorithms because they depend solely on the optimal model parameters at the end of training, rather than requiring the system to store intermediate checkpoints throughout the entirety of the training process [cite: 2].

The iHVP Memory Wall

Despite their elegance, classical influence functions encountered a devastating scalability wall when applied to deep neural networks. The primary bottleneck is the inverse Hessian-vector product (iHVP) computation. The Hessian matrix—which represents the second-order partial derivatives of the loss function—has dimensions equal to the square of the model's total parameter count [cite: 4]. For foundation models with 8 billion to 70 billion parameters, explicit computation and inversion of this matrix are mathematically infeasible.

Instead of explicitly forming the Hessian, researchers pivoted to directly approximating the inverse Hessian-vector products. Early approaches relied on stochastic Neumann-series estimators, such as LiSSA, and Krylov-subspace methods, such as Arnoldi iteration [cite: 5]. However, as models scaled into the 8 to 70 billion parameter range, even these estimators failed. For instance, creating the necessary projection matrix for an 8-billion parameter model with a 4,000-dimensional projection requires a staggering 128 Terabytes of space [cite: 4]. Consequently, running traditional influence functions on massive models routinely results in Compute Unified Device Architecture (CUDA) out-of-memory errors, even on top-tier hardware like A100 GPUs with 80GB of VRAM (Video Random Access Memory) [cite: 6].

Breaking the Memory Wall: EK-FAC and Low-Rank Projections

To salvage implicit-differentiation-based TDA, researchers developed parameter-space methods that exploit the structural properties of neural network gradients. One of the most significant breakthroughs was EK-FAC (Eigenvalue-corrected Kronecker-Factored Approximate Curvature), which leverages the Kronecker structure of gradient matrices. The Kronecker structure is a mathematical property where a massive matrix can be represented as the tensor product of two significantly smaller matrices. As a real-world analogy, this is like compressing a highly detailed, sprawling 3D landscape into two simple 2D blueprints—a top-down and a side-view—that can be multiplied to reconstruct the full shape. For the TDA bottleneck, this structural assumption allows the system to approximate the monstrous, computationally impossible Hessian matrix using much smaller, mathematically tractable matrices, enabling faster iHVP computations [cite: 4, 5]. By using EK-FAC parameterization to approximate the Hessian, the system gains access to an explicit eigendecomposition, which dramatically accelerates the calculation of the necessary mathematical terms [cite: 4].

Using EK-FAC, engineers successfully scaled influence functions to LLMs containing over 52 billion parameters [cite: 4, 7]. However, while EK-FAC solved the iHVP bottleneck, it introduced a secondary scalability challenge: dot product computations. To calculate the per-sample gradient for all candidate training data points, EK-FAC requires recomputing the training gradients for each query batch [cite: 4, 5]. For LLM pretraining, this process is effectively as expensive as pretraining the model for an entire epoch [cite: 4].

The LoGra Paradigm

To overcome the limitations of EK-FAC, the focus shifted to projection-based methods, leading to the development of LoGra (Low-Rank Gradient Approximation). LoGra is a gradient projection strategy that explicitly leverages the gradient structure found in standard backpropagation [cite: 8]. Instead of relying purely on parameter-space approximations, LoGra stores projected per-example gradients and utilizes a damped Gauss-Newton Hessian approximation [cite: 5].

The efficiency gains achieved by LoGra are monumental. The following data details the memory and throughput efficiency of LoGra compared to the previous state-of-the-art, EK-FAC, when evaluating a Llama3-8B-Instruct model using a single A100 GPU.

Method	Logging Phase Throughput (tokens/s)	Logging Phase VRAM (GB)	Influence Phase Throughput (pairs/s)	Influence Phase VRAM (GB)
EK-FAC	1,437 / 76	28.5 / 62.0	0.08	62.5
LoGra (Proposed)	2,347	47.4	561.43	11.4

As demonstrated, LoGra achieves competitive accuracy against more expensive baselines while delivering up to a 6,500x improvement in throughput and nearly a 5x reduction in GPU memory usage during the recurring influence computation phase [cite: 8]. By utilizing two-sided random projections, LoGra bypasses the massive projection matrices that limit earlier methods, allowing it to afford a higher projection dimension and thus retain greater mathematical expressivity [cite: 5, 6]. Furthermore, gradient projection enables LoGra to compute the raw projected Fisher information matrix (the Hessian) without relying on the heavier approximations used by EK-FAC [cite: 6].

LoRIF and Subspace Approximations

Building upon the projection framework, the recently proposed LoRIF (Low-Rank Influence Functions) method tackles the quality-scalability tradeoff inherent in dimensional projections. Both the storage of projected gradients and the computation of the Hessian scale with the projection dimension, which is strictly necessary for maintaining attribution quality [cite: 5].

LoRIF exploits the low-rank structures of gradients to address these bottlenecks directly. First, instead of storing full matrices, it stores rank-c factors of the projected per-example gradients, successfully reducing storage and query-time input/output demands from a linear complexity to a smaller fractional complexity per layer, per sample [cite: 5]. Second, LoRIF utilizes truncated Singular Value Decomposition (SVD)—a mathematical technique used to reduce the dimensionality of a dataset while preserving its core variance—combined with the Woodbury matrix identity. This enables the approximation of the Hessian term within a restricted subspace, reducing the memory footprint from a quadratic burden to a much more manageable linear calculation [cite: 5].

The Convergence Fallacy and Underspecification

While frameworks like LoGra and LoRIF solve the hardware and memory constraints of influence functions, they cannot resolve the fundamental conceptual flaws inherent to implicit differentiation. The classical formulation of influence functions relies on two critical mathematical assumptions: the absolute uniqueness of the model's parameters, and the model's perfect convergence to a theoretical optimal solution [cite: 2].

In the reality of modern deep learning, neither assumption holds true. Large language models are highly underspecified; multiple distinct parameter configurations can yield the exact same loss value. Furthermore, optimization algorithms like Stochastic Gradient Descent (SGD) or Adam rarely, if ever, reach a true global optimum before training is halted to prevent overfitting. Consequently, applying static influence functions to a non-converged model severs the mathematical connection between the theoretical counterfactual prediction and the model's actual behavior [cite: 2].

This conceptual flaw becomes catastrophic when analyzing multi-stage training pipelines. Foundational LLMs are frequently subjected to continual learning paradigms, such as initial pretraining followed by domain-specific fine-tuning (e.g., transitioning from a general text dataset to a specialized medical dataset) [cite: 3]. Because implicit-differentiation-based methods evaluate only the final model checkpoint, they possess no mechanism to separate these distinct stages of training. When forced to analyze a two-stage process, tools like EK-FAC must blindly combine the data from both stages into a single massive dataset for TDA, erasing the nuanced sequence of how the model acquired its knowledge [cite: 3].

The Unrolling Paradigm: TracIn and SOURCE

To correct the epistemological failures of static influence functions, researchers developed trajectory-based or "unrolling" methods. Instead of relying on a single final snapshot of the model, unrolling methods calculate the effect of removing a data point by physically backpropagating through the preceding optimization steps taken during training [cite: 2].

Consider a standard model update rule at a given iteration with a data point weight. Unrolling-based TDA methods estimate the counterfactual removal by evaluating the chain rule across the sequence of historical parameters:

θ_T_removed ≈ θ_T - (I - η_T-1 H_T-1)...(I - η_k+1 H_k+1)(-η_k ∇ L(θ_k, z_k))

The primary advantage of unrolling is theoretical robustness. These methods do not rely on the uniqueness of the optimal solution, nor do they require the model to have converged [cite: 2]. They natively incorporate the implicit bias of specific optimization algorithms and can perfectly separate multi-stage training pipelines by restricting their analysis to the relevant segments of the training trajectory [cite: 2, 3].

TracIn and the Memory Burden

One of the most prominent gradient-based unrolling techniques is TracIn, which estimates the importance of a training data point by approximating the total change in a query's measurable quantity through the integration of gradient updates across numerous checkpoints along the full training path [cite: 3, 9]. TracIn successfully bypasses the convergence assumption, providing a highly faithful account of the path-dependent nature of gradient descent [cite: 9].

However, pure unrolling methodologies face severe scalability challenges of their own. Executing a perfect unrolling sequence requires the system to store all intermediate variables generated during the entire training process—including parameter vectors for every single optimization step [cite: 2]. For models trained over trillions of tokens, storing continuous optimization steps is physically impossible.

SOURCE: Bridging the Divide

To reconcile the theoretical accuracy of unrolling with the computational feasibility of implicit differentiation, researchers introduced SOURCE. SOURCE operates as an approximate unrolling-based TDA method computed using an influence-function-like formula [cite: 10].

Instead of requiring the continuous storage of every optimization step, SOURCE strategically leverages a small number of intermediate checkpoints saved periodically during training [cite: 2, 3]. In experimental implementations evaluating multi-stage continual learning on datasets like RotatedMNIST (a variant of the standard MNIST dataset containing 70,000 grayscale 28x28 pixel images of handwritten digits rotated at various angles from 0 to 360 degrees to test rotation invariance, sized around 12 MB to 50 MB) and PACS (a domain generalization dataset containing 9,991 images of 227x227 pixels across 7 classes—dog, elephant, giraffe, guitar, house, horse, person—distributed in 4 distinct domains: Photo, Art painting, Cartoon, and Sketch), SOURCE utilized just two partitioned segments to perform accurate attribution on the first stage of training, bypassing the interference of the second stage [cite: 3]. Within each segment, SOURCE estimates the stationary Hessian and gradient by averaging them over the checkpoint interval, and then employs EK-FAC parameterization to approximate the Hessian efficiently [cite: 4].

This hybrid architecture allows SOURCE to inherit the best traits of both paradigms. It enables accurate TDA analysis for multi-stage pipelines and non-converged models, while completely avoiding the need to save all intermediate variables for unrolled differentiation [cite: 2, 4]. It is important to note that this theoretical robustness comes at a cost; compared to static EK-FAC influence analysis, SOURCE is approximately six times more computationally expensive due to the necessity of performing computations across multiple trajectory segments [cite: 4, 10]. Nevertheless, in environments where implicit-differentiation-based approaches struggle—such as un-converged foundations models—SOURCE consistently outperforms existing techniques in counterfactual prediction [cite: 10].

Stagewise Data Attribution and Singular Learning Theory

The evolution from static influence functions to trajectory-unrolling methods like SOURCE has revealed a deeper underlying truth about neural network behavior: the influence that one data sample has upon a model is not a fixed, static property [cite: 9]. Instead, neural networks learn in distinct developmental stages, exhibiting dynamic and constantly shifting patterns of influence.

Recent frameworks grounded in singular learning theory—a mathematical framework that models how neural networks learn as a trajectory through complex, non-linear geometry with distinct phase transitions. To visualize this via analogy, imagine a caterpillar undergoing sudden metamorphosis into a butterfly, rather than just growing linearly; the model's understanding shifts abruptly, completely altering how it interacts with its environment. This theory's direct relevance to TDA is that it mathematically proves the influence of a data point is dynamic and phase-dependent, rendering static, single-checkpoint TDA methods fundamentally incomplete [cite: 9]. These frameworks predict that the influence of a specific training sample can change non-monotonically over time. As a model trains, it undergoes sharp developmental transitions, leading to phenomena such as sharp peaks in influence and even total sign flips [cite: 9].

In practical terms, this means that a data point's effect is heavily dependent on when the model encounters it. For example, an image of a golden retriever might drastically help a model learn the broad semantic concept of "dog" during the early stages of training. However, that exact same image might actively harm the model's ability to distinguish fine-grained categories like "poodle" versus "terrier" in the later stages of training [cite: 9].

These theoretical predictions, validated empirically in toy models mapping progressive learning hierarchies, call for a permanent architectural shift away from classical static views of data attribution. The field is rapidly transitioning toward "stagewise data attribution"—the analysis of influence as a dynamic trajectory over the entire learning process [cite: 9]. The ultimate goal of this framework is to attribute learned behaviors not merely to the raw data that influenced the model, but to the precise developmental window in which that data exerted its defining effect [cite: 9].

Fact Tracing and the Horizon of Data Transparency

Despite the massive strides made by LoGra, EK-FAC, and SOURCE, the holistic application of TDA to generative artificial intelligence remains a frontier fraught with challenges. One of the ultimate stress-tests for TDA is "fact tracing": the process of identifying exactly which training examples taught a language model to generate a specific factual assertion [cite: 1].

When applying both gradient-based and embedding-based TDA methods to the problem of fact tracing in massive pretraining corpora, researchers found a stark misalignment between factual attribution and causal influence. While increasing the model size and the number of training tokens generally improves the alignment of calculated influence with true attribution, significant headroom remains [cite: 1]. In many empirical benchmarks, complex gradient-based TDA methods actually exhibited lower proponent-retrieval precision than rudimentary information retrieval baselines (like BM25, or Best Matching 25, a classical bag-of-words ranking function used by search engines to estimate the relevance of documents to a given query) that do not possess any access to the model's internal parameters at all [cite: 1].

This highlights the ultimate complexity of modern neural architectures. When examining the examples identified as highly influential by advanced TDA methods, researchers discovered that while some examples directly entail the specific factual assertion, others influence the output indirectly [cite: 1]. An LLM might generate a correct factual response not because it memorized a single definitive sentence, but because thousands of disparate data points reinforced its broader priors regarding relation types, common entities, and linguistic structures [cite: 1].

As the artificial intelligence industry continues to navigate the legal and ethical complexities of massive data scraping, the demand for precise data valuation systems will only intensify. Whether through the highly efficient subspace projections of LoRIF or the theoretically robust trajectory-unrolling of SOURCE, the ability to trace the hidden gradient descent of neural networks represents the next critical pillar of machine interpretability.

Operational Overview of Modern TDA Frameworks

To summarize the current landscape, the following breakdown operationalizes the distinct advantages, limitations, and use cases of the leading TDA architectures:

• EK-FAC:
  • Functional Scope: Parameter-space Hessian approximation leveraging Kronecker factored curvature.
  • Cost/Overhead: Extremely high computational overhead for recurring dot products; requires recomputing training gradients for every query.
  • Availability: Custom PyTorch implementations in academic repositories; not strictly integrated into standard ML libraries.
  • Real-World Context: Ideal for small to medium-sized architectures evaluating static datasets. An anti-use case for foundation models with 8 billion to 70 billion parameters due to crippling memory demands.
• LoGra:
  • Functional Scope: Gradient projection strategy utilizing a damped Gauss-Newton Hessian approximation.
  • Cost/Overhead: Medium overhead; achieves up to a 6,500x improvement in throughput and a 5x reduction in GPU memory usage compared to EK-FAC.
  • Availability: Specialized open-source research repositories directly associated with projection-based literature.
  • Real-World Context: Optimal for single-node influence analysis on modern LLMs (e.g., Llama 3 8B), efficiently balancing theoretical expressivity with strict hardware constraints.
• LoRIF:
  • Functional Scope: Subspace gradient approximation via truncated Singular Value Decomposition (SVD) and rank-c factors.
  • Cost/Overhead: Low overhead; successfully reduces storage and query-time I/O demands to a highly manageable linear fractional complexity per layer.
  • Availability: Open-source implementations via leading academic release channels.
  • Real-World Context: Best suited for highly restricted storage environments where maintaining full projection matrices is physically impossible.
• TracIn:
  • Functional Scope: Full trajectory unrolling approximation via gradient integration across historical checkpoints.
  • Cost/Overhead: Extreme memory burden; mathematically requires storing intermediate parameter vectors for virtually every optimization step.
  • Availability: Widely available via established explainability libraries such as PyTorch Captum.
  • Real-World Context: Perfect for small, non-converged supervised models. However, it represents an absolute anti-use case for multi-stage LLM pretraining encompassing trillions of tokens.
• SOURCE:
  • Functional Scope: Hybrid periodic checkpoint unrolling utilizing stationary Hessian estimation over segmented trajectory intervals.
  • Cost/Overhead: High overhead; approximately six times more computationally expensive than static EK-FAC due to multi-segment evaluations.
  • Availability: Custom research implementations primarily found in theoretical trajectory-unrolling literature.
  • Real-World Context: The definitive choice for auditing multi-stage continual learning pipelines (e.g., isolated pretraining vs. fine-tuning phases) where static methods suffer from epistemological failure.

Implementing TDA: Practical Influence Computation

For engineers seeking to operationalize these concepts, the following PyTorch pseudo-code demonstrates a rudimentary implementation of gradient dot-product calculations—the foundational mechanic behind unrolling methods like TracIn. By capturing the parameter gradients for a specific training example and comparing them to the gradients of a validation query, we can establish a baseline influence score.

import torch
import torch.nn.functional as F
def compute_tracin_influence(model, train_batch, query_batch, optimizer_lr):
 """
 Computes a simplified TracIn influence score between a training batch
 and a query (validation) batch at a single checkpoint.
 """
 model.eval()
# 1. Compute gradients for the query batch
query_inputs, query_targets = query_batch
query_outputs = model(query_inputs)
query_loss = F.cross_entropy(query_outputs, query_targets)
model.zero_grad()
query_grads = torch.autograd.grad(query_loss, model.parameters(), retain_graph=True)
train_inputs, train_targets = train_batch
train_outputs = model(train_inputs)
train_loss = F.cross_entropy(train_outputs, train_targets)
model.zero_grad()
train_grads = torch.autograd.grad(train_loss, model.parameters())
influence_score = 0.0
for q_grad, t_grad in zip(query_grads, train_grads):
 influence_score += torch.dot(q_grad.view(-1), t_grad.view(-1)).item()
return influence_score * optimizer_lr

Sources & References

1. github.io
2. juhanbae.com
3. neurips.cc
4. juhanbae.com
5. arxiv.org
6. neurips.cc
7. semanticscholar.org
8. openreview.net
9. timaeus.co
10. openreview.net