Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Scalable Data Attribution via Forward-Only Test-Time Inference

Sibo Ma, Julian Nyarko

TL;DR

This work tackles scalable data attribution by preserving the traditional influence-function target while eliminating per-query backprop at inference. It achieves this through forward-only attributions based on a training-time simulation: perturb each training example with short-horizon gradient updates, accumulate a curvature-aware parameter response, and read out a forward score that converges to the damped inverse influence, $-g_q^T H_ abla^{-1} g_b$, via a two-trajectory or single-trajectory readout. Empirically, the method matches or exceeds TRAK on standard attribution metrics (LOO and LDS) on MNIST-MLP with far lower inference cost, demonstrating practical, real-time data attribution at scale. By shifting computation to training and using a curvature-aware forward readout, the approach offers a principled, scalable framework for tracing data provenance and valuing data sources in large pretrained models.

Abstract

Data attribution seeks to trace model behavior back to the training examples that shaped it, enabling debugging, auditing, and data valuation at scale. Classical influence-function methods offer a principled foundation but remain impractical for modern networks because they require expensive backpropagation or Hessian inversion at inference. We propose a data attribution method that preserves the same first-order counterfactual target while eliminating per-query backward passes. Our approach simulates each training example's parameter influence through short-horizon gradient propagation during training and later reads out attributions for any query using only forward evaluations. This design shifts computation from inference to simulation, reflecting real deployment regimes where a model may serve billions of user queries but originate from a fixed, finite set of data sources (for example, a large language model trained on diverse corpora while compensating a specific publisher such as the New York Times). Empirically, on standard MLP benchmarks, our estimator matches or surpasses state-of-the-art baselines such as TRAK on standard attribution metrics (LOO and LDS) while offering orders-of-magnitude lower inference cost. By combining influence-function fidelity with first-order scalability, our method provides a theoretical framework for practical, real-time data attribution in large pretrained models.

Scalable Data Attribution via Forward-Only Test-Time Inference

TL;DR

This work tackles scalable data attribution by preserving the traditional influence-function target while eliminating per-query backprop at inference. It achieves this through forward-only attributions based on a training-time simulation: perturb each training example with short-horizon gradient updates, accumulate a curvature-aware parameter response, and read out a forward score that converges to the damped inverse influence, , via a two-trajectory or single-trajectory readout. Empirically, the method matches or exceeds TRAK on standard attribution metrics (LOO and LDS) on MNIST-MLP with far lower inference cost, demonstrating practical, real-time data attribution at scale. By shifting computation to training and using a curvature-aware forward readout, the approach offers a principled, scalable framework for tracing data provenance and valuing data sources in large pretrained models.

Abstract

Data attribution seeks to trace model behavior back to the training examples that shaped it, enabling debugging, auditing, and data valuation at scale. Classical influence-function methods offer a principled foundation but remain impractical for modern networks because they require expensive backpropagation or Hessian inversion at inference. We propose a data attribution method that preserves the same first-order counterfactual target while eliminating per-query backward passes. Our approach simulates each training example's parameter influence through short-horizon gradient propagation during training and later reads out attributions for any query using only forward evaluations. This design shifts computation from inference to simulation, reflecting real deployment regimes where a model may serve billions of user queries but originate from a fixed, finite set of data sources (for example, a large language model trained on diverse corpora while compensating a specific publisher such as the New York Times). Empirically, on standard MLP benchmarks, our estimator matches or surpasses state-of-the-art baselines such as TRAK on standard attribution metrics (LOO and LDS) while offering orders-of-magnitude lower inference cost. By combining influence-function fidelity with first-order scalability, our method provides a theoretical framework for practical, real-time data attribution in large pretrained models.

Paper Structure

This paper contains 19 sections, 31 equations, 2 figures.

Table of Contents

  1. Introduction
  2. Related work
  3. Method
  4. Target Attribution Score
  5. Simulating a training example’s influence
  6. From parameter response to an influence score on $q$
  7. Evaluation
  8. Hyperparameter analysis
  9. Unroll depth $T$.
  10. Inner step size $\eta$.
  11. Damping $\lambda$.
  12. Performance comparison with existing works
  13. Analytical comparison of computational complexity
  14. Simulation cost.
  15. Inference cost.
  16. ...and 4 more sections

Figures (2)

  • Figure 1: Comparison of LOO and LDS performance under varying $\eta$ and unroll steps $T$. On a small MLP trained on MNIST, a moderate $\eta$ ($3\times10^{-2}$) yields stable and monotonic improvement. Larger $\eta$ accelerates early gains but leads to instability, while smaller $\eta$ converges slowly yet remains stable.
  • Figure 2: Performance comparison on the MNIST–MLP benchmark from DattriNEURIPS2024_f7326833. Each bar reports the correlation under LOO and LDS. Our method achieves $0.022$ LOO and $0.49$ LDS, matching or exceeding the strongest baselines while requiring only forward evaluations at inference time.