Scalable Data Attribution via Forward-Only Test-Time Inference

TL;DR

This paper introduces a scalable data attribution method that shifts computation from inference to training simulation, enabling real-time attribution in large models with lower computational costs.

Contribution

The proposed method eliminates per-query backward passes by simulating training influence through short-horizon gradient propagation, improving scalability and efficiency.

Findings

Matches or surpasses state-of-the-art attribution metrics

Offers orders-of-magnitude lower inference cost

Applicable to large pretrained models for real-time attribution

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.