🤖 machine learning

How Faithful Is Trajectory-Based Data Attribution? Error Sources, Remedies, and Practical Guidelines

This paper provides the first systematic analysis of error sources in trajectory-based data attribution, identifying optimizer mismatch as a dominant issue and proposing AdamW-influence, a closed-form error proxy, and a K-step look-ahead framework to significantly improve attribution accuracy and offer practical guidelines for reliable data selection.

Original authors: Junwei Deng, Pingbang Hu, Suliang Jin, Hao Lu, Jiachen T. Wang, Shichang Zhang, Jiaqi W. Ma

Published 2026-05-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Junwei Deng, Pingbang Hu, Suliang Jin, Hao Lu, Jiachen T. Wang, Shichang Zhang, Jiaqi W. Ma

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 are baking a massive cake (training a modern AI model) using a huge bowl of ingredients (training data). Sometimes, you want to know: "Which specific egg or pinch of sugar actually made the cake taste better or worse?" This is called data attribution.

For a long time, scientists used a method called "Trajectory-Based Attribution" to answer this. They tried to "rewind" the baking process step-by-step to see how removing one ingredient would change the final result.

However, this paper argues that the current way of doing this "rewind" is often broken, leading to wrong answers. The authors act like mechanics who have opened up the engine to find exactly where the gears are grinding and how to fix them.

Here is a simple breakdown of their findings and fixes:

1. The Problem: The "Wrong Map" (Config-Level Error)

The Analogy: Imagine you are trying to navigate a city using a map designed for a bicycle, but you are actually driving a heavy truck with a turbo engine. Even if you follow the map perfectly, you will get lost because the map doesn't understand how your truck handles turns or acceleration.

The Reality: Most modern AI models are trained using a sophisticated "engine" called AdamW. However, the popular data attribution tools were built assuming the models were trained with an older, simpler engine called SGD.

The Result: The tools were using the "bicycle map" for the "turbo truck." This caused massive errors in figuring out which data points were helpful.
The Fix: The authors built a new tool called AdamW-influence. This is a map specifically designed for the turbo truck.
The Outcome: By using the right map, they improved the accuracy of their predictions by 10% to over 300% across different types of AI models (from simple image classifiers to large language models like Llama).

2. The Second Problem: The "Rough Sketch" (Algorithm-Level Error)

The Analogy: Even with the right map, if you try to predict the future by drawing a straight line on a curvy road, you will eventually drift off course. The longer the road you try to predict, the bigger the mistake.

The Reality: To make the calculations fast, these tools use a "first-order approximation." Think of this as approximating a winding road as a straight line.

The Culprits: The authors found two main things that make this "straight line" guess worse:
1. Steepness (Learning Rate): If the steps taken during training are huge (high learning rate), the straight-line guess fails quickly.
2. Distance (Trajectory Length): The further back in time you try to look, the more the "straight line" drifts away from the actual curvy path.
The Fix: They created a "Error Proxy." This is like a dashboard warning light that tells you, "Hey, the straight-line guess is getting unreliable right now," without needing to re-bake the whole cake to check. It helps users know when to trust the data scores and when to be skeptical.

3. The Practical Guide: How to Pick Ingredients (Data Selection)

The Analogy: Imagine you are a chef picking ingredients for a soup while you are cooking it.

Offline Strategy: You wait until the soup is finished, taste it, and then say, "If I had picked these specific carrots instead of those ones, it would have been better." (This is slow and expensive).
Online Strategy: You pick ingredients as you go, guessing how they will affect the soup in the next few minutes.

The New Rulebook: The authors unified these two strategies into a single framework called the "K-Step Look-Ahead."

K is how far into the future you look.
The Insight: You don't need to look all the way to the end of the soup (Offline). Looking just a few steps ahead (Online) is often just as good, if you use the right tools.
The Sweet Spot:
- If you take small steps (low learning rate), you can look further ahead (larger K) without making mistakes.
- If you take big steps (high learning rate), you should only look a short distance ahead (small K) to avoid errors piling up.

Summary of the "Recipe" for Practitioners

The paper gives a clear recipe for anyone using these tools:

Stop using the old "SGD" tools if you are training with AdamW (which almost everyone is). Use the new AdamW-influence instead.
Don't look too far ahead when picking data online. If you look too far into the future, the "straight line" guess becomes too noisy.
Tune your horizon (K) with your learning rate. If you are taking small steps, you can look further ahead. If you are taking big steps, keep your gaze short.

By fixing the "map" and understanding the "straight line" limits, the authors have turned data attribution from a black box into a reliable, actionable tool for improving AI models.

Technical Summary: How Faithful Is Trajectory-Based Data Attribution?

1. Problem Statement

Trajectory-based data attribution methods (e.g., SGD-influence, TracIn, SOURCE) estimate the influence of individual training samples on model predictions by unrolling the training trajectory. While widely used for data selection, valuation, and model diagnosis, these methods lack a comprehensive error analysis. Practitioners face fundamental uncertainties regarding the magnitude and origin of attribution errors, particularly concerning whether errors stem from mismatches between the attribution formula and actual training dynamics (e.g., optimizer mismatch) or from inherent algorithmic approximations (e.g., first-order Taylor expansions). This uncertainty hinders the reliable deployment of these methods in modern deep learning pipelines, which predominantly utilize optimizers like AdamW rather than the SGD assumed by existing theoretical derivations.

2. Methodology and Taxonomy

The authors propose the first systematic error analysis of trajectory-based attribution, organizing total error into three distinct categories:

Config-level Error: Arises when the attribution formula does not faithfully reflect the actual training procedure. The dominant source identified is optimizer mismatch, where methods assume Stochastic Gradient Descent (SGD) dynamics despite models being trained with AdamW.
Algorithm-level Error: Arises from efficiency approximations, specifically the first-order Taylor expansion used to linearize the parameter trajectory.
System-level Error: Irreducible errors due to floating-point precision and GPU non-determinism, which affect even ground-truth retraining.

To address these, the paper introduces:

AdamW-influence: A new algorithm that unrolls the actual AdamW optimization dynamics (including momentum states and coordinate-wise adaptive scaling) rather than assuming SGD. It injects perturbations at the gradient level to correctly handle the nonlinearity of AdamW's second-moment updates.
Error Proxy: A closed-form expression to estimate the per-sample algorithm-level error (update-estimation error) along the original trajectory without retraining.
K-Step Look-Ahead Framework: A unified theoretical framework for data selection that treats offline and online strategies as boundary cases of a horizon parameter $K$ , allowing for the analysis of the trade-off between lookahead information and accumulated approximation error.

3. Key Contributions

A. Correction of Config-Level Error (Optimizer Mismatch)

The authors demonstrate that existing methods derive attribution under the assumption of SGD, creating systematic deviations when applied to models trained with AdamW. They propose AdamW-influence, which correctly accounts for AdamW's momentum and adaptive scaling.

Result: Across four settings (MLP, CNN, GPT-2, and Llama 3.2-1B), AdamW-influence improves the Spearman correlation between estimated and ground-truth influence (TSLOO) by 10% to over 300% compared to the SGD-based baseline.
Implication: Optimizer alignment is a necessary prerequisite for isolating and analyzing the remaining algorithm-level errors.

B. Dissection of Algorithm-Level Error

With config-level error corrected, the authors isolate the error arising from the first-order Taylor expansion. They identify two governing factors:

Learning Rate ( $\eta$ ): Larger learning rates increase the per-step linearization error.
Trajectory Length ( $T - t^*$ ): Longer remaining trajectories allow errors to accumulate.

Result: The error magnitude is inversely proportional to the learning rate and directly proportional to the remaining trajectory length.
Contribution: The authors derive a closed-form error proxy (Equation 4) that estimates per-sample error norms using quantities already maintained by the optimizer. This proxy correlates strongly ( $\rho \in [0.83, 0.89]$ ) with ground-truth error norms obtained via leave-one-out retraining, enabling error-aware decisions without retraining.

C. Practical Guidelines for Data Selection

The paper unifies offline and online data selection strategies under a K-step look-ahead framework, where $K$ represents the number of steps into the future a selection decision considers.

Findings:
- Optimizer Alignment: Online selection using AdamW-influence consistently outperforms online SGD-influence, transferring the fidelity gains to downstream selection performance.
- Horizon vs. Error: Offline selection corresponds to the longest horizon ( $K = T-t$ ) and accumulates the largest approximation error. Online selection with a short-to-moderate horizon often matches or exceeds offline performance.
- Tuning $K$ : The optimal horizon $K$ grows as the learning rate decreases. Smaller learning rates allow the $K$ -step summation to extend further before approximation error dominates.
Guideline: Practitioners should tune $K$ jointly with the learning rate rather than fixing it, and short-horizon online selection with aligned optimizers is a reliable, actionable recipe.

4. Experimental Results

The authors evaluated their methods on:

Models: MLP, CNN, GPT-2, and Llama 3.2-1B.
Datasets: MNIST, WikiText-2, Alpaca, and Tulu3.
Metrics: Spearman rank correlation against TSLOO ground truth (for fidelity) and validation/test error rates or perplexity (for data selection quality).

Key Outcomes:

Fidelity: AdamW-influence significantly outperforms SGD-influence across all learning rates and architectures. For example, on MNIST+CNN with $lr=10^{-3}$ , the correlation improved from 0.015 to 0.090 (+500%).
Selection Quality: In data selection tasks, Online AdamW-influence consistently achieved lower error rates (on MNIST) and lower perplexity (on Llama 3.2) compared to both Online SGD-influence and Offline AdamW-influence.
Error Proxy: The derived proxy successfully tracked ground-truth error norms with high Spearman correlation ( $\rho \approx 0.85$ ), validating its use as a practical signal for trajectory-aware decisions.

5. Significance and Claims

The paper claims to provide the first systematic error analysis of trajectory-based data attribution, moving the field from treating these methods as "black-box scoring functions" to "diagnosable estimators."

Theoretical Impact: By categorizing errors into config, algorithm, and system levels, the work clarifies the sources of unfaithfulness in modern deep learning contexts.
Practical Impact: The introduction of AdamW-influence and the K-step framework offers concrete, actionable remedies for practitioners. It demonstrates that aligning the attribution formula with the actual optimizer (AdamW) and carefully tuning the look-ahead horizon can significantly improve the reliability of data selection and model diagnosis.
Scalability: The proposed methods maintain $O(T)$ complexity (with a constant overhead of roughly 3x compared to SGD-influence due to augmented state storage), making them feasible for large-scale models when combined with random masking techniques.

The authors acknowledge limitations, noting that their analysis is restricted to first-order Taylor-based methods and models up to Llama 3.2-1B due to computational constraints, and that generalization to other regimes (e.g., reinforcement learning, diffusion models) remains an open direction.