PaperlayerView exact revision on arXiv

Source-linked AI summary

STRIDE: Training Data Attribution via Sparse Recovery from Subset Perturbations

Rishit Dagli, Abir Harrasse, Luke Zhang, Florent Draye, Amirali Abdullah, Bernhard Schölkopf, Zhijing Jin

arXiv:2606.05165v1cs.LGcs.CL

TL;DR

Training Data Attribution needs causal estimates of how training examples affect predictions, but retraining is costly and gradient methods scale poorly. STRIDE learns low-rank activation-space operators that simulate subset-training effects and recovers sparse per-example influences. It achieves state-of-the-art LLM pre-training attribution while being an order of magnitude faster, with demonstrated utility for auditing and data curation.

  • Problem

    Training Data Attribution requires tracing prediction changes to individual examples, while retraining-based causal estimates are computationally prohibitive and gradient approximations require massive parameter-space computation.

  • Method

    STRIDE learns low-rank activation-space steering operators that mimic subset-retraining responses and recovers per-example influence through sparse linear decomposition.

  • Results

    STRIDE achieves state-of-the-art LLM pre-training attribution accuracy and is an order of magnitude more computationally efficient than existing methods.

  • Takeaways & Limitations

    STRIDE’s attributions support scalable data curation, model auditing, leave-one-out approximation, and effective data selection.

  • Takeaways & Limitations

    STRIDE assumes local linearity and additive influence, which may break down under extreme distribution shifts or highly non-convex memorization; RL-based training remains open.

Abstract

from arXiv · show

Training Data Attribution (TDA) seeks to trace a model's predictions back to its training data. The gold standard for TDA relies on causal interventions, observing how a model changes when data is added or removed, but repeated retraining is computationally challenging for Large Language Models (LLMs). Consequently, most approaches approximate this effect in the parameter space using gradients. However, tracking gradients across billions of parameters is not only prohibitively expensive but relies on local approximations. In this work, we propose a shift: rather than estimating parameter changes, we model the functional effect of training data in the activation space. We introduce STRIDE (Steering-based Training Data Influence Decomposition), a framework that formulates TDA as a sparse recovery problem in the spirit of compressive sensing. STRIDE learns lightweight "steering operators" that mimic the behavioral shift caused by training on data subsets. By measuring how these operators perturb test predictions, we recover individual training example influences via sparse linear decomposition. STRIDE achieves state-of-the-art for LLM pre-training attribution while being an order of magnitude ($13\times$) faster than previous art. We further validate its practical utility through downstream applications including data selection, data contamination, and qualitative analysis.

1 Introduction

Training Data Attribution traces model predictions to influential training examples, but retraining-based causal estimates are prohibitively expensive and gradient approximations face severe scale and locality limits. STRIDE instead models subset-training effects in activation space and recovers sparse per-example influences, achieving state-of-the-art attribution accuracy with substantially lower computation.

  • Training Data Attribution quantifies how individual training examples causally influence specific model predictions, supporting auditing, memorization detection, debugging, and dataset curation.
  • Retraining-based counterfactual attribution is rigorous but computationally prohibitive for LLMs, motivating parameter-space gradient approximations.
  • STRIDE models training-data effects as activation-space shifts and formulates attribution as sparse recovery from subset perturbations.
  • STRIDE learns low-rank steering operators that mimic randomly sampled subset-retraining responses, then uses their test-query perturbations to recover individual influences.
  • STRIDE achieves state-of-the-art LLM pre-training attribution, outperforms gradient- and representation-based baselines, and is an order of magnitude more computationally efficient.
  • The framework also supports SFT evaluation and actionable downstream applications including leave-one-out approximation and data selection.

2 Related Work

Prior TDA methods operate through parameter gradients, representation similarity, or subset-based valuation. STRIDE combines the causal grounding of subset responses with scalable activation-space modeling and sparse per-example recovery.

  • Gradient-based attribution estimates example effects through local parameter perturbations, with extensions using Hessian approximations, projections, normalization, and low-rank methods.
  • Representation-based methods scale naturally by using embedding similarity or learned scoring functions, but typically predict influence through heuristic similarities or separate scorers.
  • Subset-based approaches capture counterfactual influence and non-additive interactions but generally require repeated retraining or evaluation across many subsets.
  • STRIDE uses subset-retraining responses as supervision for activation-space steering operators, enabling efficient per-example recovery from a single query.

3 Problem Setup and Preliminaries

The paper defines attribution as decomposing a target response into per-example contributions to a set-function response under dataset removal. Because exact subset effects require retraining, it adopts an additive sparse-recovery formulation.

  • A training set S contains examples z1 through zn, and a model is trained by minimizing empirical risk over S.
  • For target x, the goal is to decompose the model response r(x; S) into contributions from individual training examples.
  • The set function Fx(A) maps each training subset A to the model response obtained when training on A.
  • The perturbation response δx(A) measures the change in target loss caused by removing subset A, with singleton subsets yielding exact leave-one-out effects under full retraining.
  • The additive influence assumption represents every subset response as a sum of per-example influences, enabling recovery from subset-level observations.
  • When the number of subset measurements is much smaller than the dataset size, sparsity places attribution in the compressive-sensing regime.

4 STRIDE

STRIDE has an offline phase that learns activation-space steering operators to simulate subset perturbations and an online phase that recovers sparse per-example influences. Its design uses low-rank operators, additive regularization, and structured subset measurements to avoid repeated retraining.

  • STRIDE overview: STRIDE first learns activation-space steering operators and then recovers per-example influences through sparse recovery.
  • Operator learning: The operators act on a frozen base model’s intermediate activations to simulate functional effects of retrained, overlapping data subsets.
  • Operator learning: A shared low-rank basis projects latent features, while subset-specific steering matrices translate them into shifts in output logits.
  • Operator learning: Joint training combines fidelity to subset effects, stability on unrelated examples, and LDS-based linearity regularization.
  • Online recovery: After training, the operators provide zeroshot counterfactual responses for new test points through a batched forward pass without further training.
  • Online recovery: Per-example influences are recovered by solving an ℓ1-regularized least-squares problem from the steering-induced measurement vector.
  • Measurement design: If the binary subset matrix forms a suitable expander, any k-sparse influence vector can be uniquely recovered from O(k log(n/k)) measurements.
  • Measurement design: In the largest setting, K = 1000, n ≈11.4M, and k ≈50 requires approximately 617 measurements under k log(n/k).

5 Experiments and Results

STRIDE is evaluated across pre-training, supervised fine-tuning, data selection, contamination auditing, and qualitative analyses. The experiments assess attribution quality, scalability, and downstream utility against gradient- and representation-based methods.

  • Experimental setup: STRIDE is evaluated across pre-training and supervised fine-tuning settings, downstream influence scoring applications, and qualitative analyses.The study uses Nanochat, Qwen2.5-0.5B, OLMo 2-7B, and Qwen2.5-32B models across multiple datasets and objectives.
  • Pre-training influence: At 1.38B parameters, STRIDE is 5× faster than LoGRA and 12× faster than AirRep while achieving the highest LDS.The runtime comparison covers end-to-end runtime and peak GPU VRAM across model sizes.
  • Pre-training influence: STRIDE outperforms prior attribution methods on pre-training LDS while remaining more than an order of magnitude faster at scales up to 1.38B parameters.Pre-training evaluation uses K = 256 data subsets and 500 held-out test queries; at the largest scale STRIDE completes in 9.9 hours.
  • SFT influence: In SFT, STRIDE usually outperforms prior art and remains an order of magnitude faster, although AirRep performs better on SafeRLHF and FLAN.The comparison follows AirRep’s LDS evaluation protocol across FLAN, Alpaca, Tulu, and SafeRLHF.
  • Data selection: STRIDE supports data selection by ranking candidate examples, selecting the top 1,000 per task, and evaluating fine-tuned models using mean unigram F1 across 66 FLAN tasks.The selected subsets are produced through greedy rank aggregation of per-query influence scores.
  • Dataset contamination: Adding STRIDE to LoGRA raises recall of leaked MATH replicas from 62.1% to 74.2% across seven contaminated models.Representation similarity retrieves more duplicates but is less specific to memorization because recall is similar for memorized and non-memorized leaked queries.
  • Qualitative analysis: STRIDE scales across model and dataset sizes and enables qualitative attribution analyses of behaviors assembled from distinct pre-training narratives.The qualitative experiments cover SafeRLHF, Nemotron-ClimbMix, and the OLMo mixture across models from 0.5B to 32B parameters.

6 Discussion and Limitations

STRIDE’s effectiveness depends on the quality of the base model’s activations and the selected intervention layer. Its local-linearity and additive-influence assumptions may fail under extreme distribution shifts or highly non-convex memorization, and RL-based objectives remain unaddressed.

  • Limitations: STRIDE depends on base-model activation quality and the choice of intervention layer.These factors constrain the effectiveness of the representation-based steering approach.
  • Limitations: Its local-linearity and additive-influence assumption may break down under extreme distribution shifts or highly non-convex memorization phenomena.The paper reports that the assumption holds robustly for standard pre-training and instruction-tuning regimes.
  • Limitations: Extending STRIDE to reinforcement-learning-based training objectives remains an open direction.

7 Conclusion

STRIDE shifts attribution from parameter-space gradients to activation-space steering and sparse recovery, enabling scalable decomposition of training contributions. Its theory connects sparse subset designs to expander-based recovery, while experiments report sparse recovered influences.

  • 7 Conclusion: STRIDE shifts training-data attribution from parameter-space gradient estimation to activation-space steering with lightweight operators.The operators simulate counterfactual training trajectories while avoiding massive gradient computation.
  • A Sparse Recovery and Compressive Sensing: The attribution problem is formulated as recovering a sparse influence vector from subset-level measurements using a binary membership matrix.The system is underdetermined when K ≪ n, so sparsity enables compressive-sensing recovery.
  • 7 Conclusion: STRIDE recovers highly sparse influence vectors, zeroing out most of the 1M+ training points for each Nanochat query.The paper attributes this sparsity to ℓ1 regularization in the Lasso solver.

B.3 Tail-Patch Score Evaluation

Tail-Patch evaluation measures whether attributed examples provide useful causal influence relative to random examples. STRIDE reports robust negative pre-training lifts and consistent attribution signals across several association metrics.

  • B.3 Tail-Patch Score Evaluation: Robust negative Tail-Patch lifts in pre-training indicate actionable influence recovery by STRIDE.More negative Lift values indicate stronger causal relevance under the reported metric.
  • B.3 Tail-Patch Score Evaluation: Table 7 evaluates gradient steps on the top k ∈ {10, 20, 50} attributed examples against random examples.The metric is defined as ∆log pattributed − ∆log prandom.
  • B.3 Tail-Patch Score Evaluation: Pearson correlation and Kendall’s τ, alongside Spearman correlation, show consistent attribution signal across evaluated statistical measures.These extended metrics are reported in Table 8 in addition to Tables 1 and 2.

B.5 Additional Data Contamination Results

Additional contamination results distinguish duplicate retrieval from model-dependent influence. STRIDE and LoGRA identify complementary extreme-score examples, while injected data produces targeted gains on leaked problems.

  • B.5 Additional Data Contamination Results: Controlled contamination sharply increases leaked-problem accuracy while leaving non-leaked accuracy near the base model.The OpenWebText-only control degrades on MATH, indicating the proxy corpus does not explain the leaked-problem gains.
  • B.5 Additional Data Contamination Results: 91.9 ± 5.2% mean recall@10 lets AirRep retrieve leaked replicas, but it does not distinguish memorized from non-memorized queries.Recall@10 is 89.8% for memorized queries and 94.2% for non-memorized queries, consistent with textual identity detection.
  • B.5 Additional Data Contamination Results: 74.2 ± 6.6% union recall is achieved by combining LoGRA and STRIDE across seven contaminated models.LoGRA alone reaches 62.1 ± 4.4% union recall, while STRIDE reaches 32.6 ± 7.0%.
  • B.5 Additional Data Contamination Results: All reported recall rates are highly significant against random extreme-bucket baselines, with one-sided binomial tests yielding p < 10^-20 for every contaminated model.The random probabilities are approximately 0.011 for LoGRA and 0.004 for the STRIDE scoring pool.

B.6 Evaluating STRIDE on Vision Models

Controlled supervised experiments test whether STRIDE’s activation-space steering and sparse-recovery formulation extends beyond language models. The results show actionable and retraining-aligned attributions, concentrated influence scores, and qualitatively meaningful signed examples across vision and tabular settings.

  • Setup: STRIDE is evaluated on MNIST, FashionMNIST, Parkinsons, and CIFAR-10 using MLP and ResNet-9 models against five attribution baselines and random removal.Ground-truth subset responses are obtained through explicit retraining and evaluated with LDS.
  • Counterfactual utility: Removing STRIDE’s top-ranked examples causes larger held-out probability drops than random removal across MNIST, FashionMNIST, and Parkinsons.The relative performance varies across datasets and architectures, with TRAK competitive in some settings such as FashionMNIST.
  • LDS evaluation: STRIDE achieves the strongest LDS Spearman correlation on MNIST and Parkinsons, remains competitive with TRAK on FashionMNIST, and trails LoGRA and TracIn on CIFAR-10.The CIFAR-10 result remains a stress test because the convolutional setting is less aligned with low-dimensional MLP settings.
  • Sparsity of recovered influence: Recovered influence vectors are highly concentrated, with substantially more normalized absolute influence mass in top-ranked examples than under a uniform baseline.Threshold curves show that only a small fraction of examples exceed moderate normalized-mass thresholds.
  • Qualitative analysis: On CIFAR-10, positive examples are typically class-aligned while negative examples often belong to competing or visually confusable classes.The signed rankings provide qualitative evidence that recovered influences distinguish supportive from opposing training examples.
  • Overall findings: Overall, the controlled experiments support STRIDE’s cross-domain mechanism through retraining-aligned scores, actionable examples, sparse influence structure, and meaningful signed attributions.These findings provide evidence that the mechanism is not specific to language models.

B.7 Additional Qualitative Results

The supplied qualitative passages cover varied content, including AI descriptions, values, scientific facts, historical text, and unsafe or biased prompts. They provide examples of generated responses and related textual material rather than a coherent additional paper result.

  • Artificial intelligence: The passages describe artificial intelligence as machines programmed to simulate human thinking, learning, problem-solving, and visual perception.This description includes speech recognition, decision-making, and related tasks.
  • Human intelligence: The passages contrast human creativity and innovation with AI outputs described as constrained by predefined patterns and data.The contrast is presented as a qualitative characterization of human and artificial intelligence.
  • Other qualitative material: Additional passages present statements about integrity, AI-content disclosure, scientific facts, black holes, historical text, politics, climate, and race.These materials are heterogeneous and do not establish a unified qualitative finding about STRIDE.

C Ablations

The ablations identify the stability-loss formulation, rank, intervention layer, loss components, and subset design choices that govern STRIDE attribution quality and recovery.

  • Evaluation: The ablation suite reports mean Spearman rank correlation and standard deviation across 500 test queries on the Nanochat pre-training suite.Computational budgets and optimization hyperparameters are held fixed across ablations.
  • Stability Loss Formulation: Truncated KL divergence outperforms Top-m Logit and Hidden Norm penalties because it better preserves the local probability simplex for linear influence modeling.Hidden Norm significantly degrades performance, while Top-m Logit remains competitive.
  • Low-Rank Dimension: r = 32 provides sufficient capacity for nuanced causal influences, whereas reducing rank lowers LDS correlation; increasing beyond 32 adds only marginal gains.The standard configuration balances attribution quality against parameter overhead.
  • Intervention Layer: Late-layer intervention achieves the best influence recovery, while early-layer injection performs significantly worse and middle layers are intermediate.The intervention is typically placed near the final language-modeling layers.
  • Loss Components: Removing Fidelity collapses LDS correlation to near-random performance, while removing Stability or Linearity causes distortion or breaks additive recovery.Linearity is required for subset effects to sum cleanly in the latent space used by the Lasso solver.
  • Subset Measurement Design: Dense Bernoulli subset measurement violates the sparse expander regime because examples occur in about K/2 subsets and exhibit excessive pairwise overlap.The typical overlap is about half the effective degree, versus the expander requirement of less than one third.

D.2 Data Selection Metric

The data-selection metric ranks candidate examples by their strongest influence across task queries, selects the top 1,000, and evaluates downstream fine-tuning performance.

  • Score Matrix: Influence scores form an m × n matrix whose rows are test queries and whose columns are the 100,000 candidate training examples.Each entry records the influence assigned to one training example for one query.
  • Greedy Rank Aggregation: Greedy rank aggregation assigns each example its best rank across a task’s test queries, favoring examples highly influential for at least one query.This handles heterogeneous instruction-following tasks whose queries may depend on different training subsets.
  • Greedy Rank Aggregation: The selected subset contains k = 1,000 examples with the smallest aggregated ranks.Selection emphasizes high influence for at least one evaluation query.
  • Unigram F1 Evaluation: Downstream performance is measured by unigram F1 after fine-tuning a fresh Qwen2.5-0.5B model on the selected examples.Scores are computed with greedy decoding against task reference answers.
  • Unigram F1 Evaluation: The reported Tab. 3 score is the mean task-level unigram F1 across all 66 tasks, multiplied by 100.The standard deviation measures variability across tasks rather than run-to-run initialization variance.
  • Baseline Evaluation: The evaluation compares STRIDE with similarity, gradient, and attribution baselines under their specified implementation protocols.The language-model experiments use shared corpus, query, block-size, and score-matrix conventions where applicable.

F.7 Data Contamination Implementation Details

The contamination implementation fine-tunes controlled Qwen models with replicated leaked problems, evaluates leaked and held-out queries, and recovers attribution scores using STRIDE’s expander-based sparse procedure.

  • Training Protocol: Contaminated and clean-control Qwen2.5-0.5B models are fine-tuned for one epoch, with each selected MATH problem copied 100 times while the training pool remains fixed.Proxy examples are replaced to maintain approximately 22,600 training examples.
  • Evaluation Protocol: Attribution recall checks whether each leaked query’s own training replica is retrieved from the pool, alongside accuracy on 500 held-out non-leaked MATH problems.AirRep uses top-k retrieval, while LoGRA and STRIDE use unions of top-100 and bottom-100 score buckets.
  • STRIDE Scoring: STRIDE trains 1,000 low-rank steering operators for 10,000 iterations using an expander subset construction with degree d = 10.Per-query scores are recovered by Lasso over 512-token pool chunks.
  • Sparse Recovery: Compressive sensing treats per-example influence as a sparse high-dimensional signal recoverable from fewer subset measurements than the ambient dimension.This sparse-recovery framing underlies STRIDE’s attribution procedure.
Loading 2606.05165v1…