PaperlayerView exact revision on arXiv

Source-linked AI summary

TRAK: Attributing Model Behavior at Scale

Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, Aleksander Madry

arXiv:2303.14186v2stat.MLcs.LG

TL;DR

Data attribution seeks to connect model predictions to the training examples behind them, but existing methods trade computational tractability against effectiveness. TRAK addresses this tension with an efficiently scalable attribution method demonstrated across image, language, and image-text models, while its fact-tracing results distinguish model behavior tracing from information retrieval.

  • Problem

    Data attribution aims to trace model predictions back to training examples, but existing approaches often trade computational tractability against attribution effectiveness, especially for non-convex models and large datasets.

  • Method

    TRAK converts additive attribution scores into counterfactual predictions and reduces complex models to a logistic-regression attribution primitive.

  • Results

    TRAK's usefulness is demonstrated across large-scale image classifiers, language models, and image-text models, including CIFAR, ImageNet, BERT, mT5, and CLIP.

  • Takeaways & Limitations

    TRAK can support behavior tracing, which requires faithfulness to the model being studied rather than merely finding data that supports generated text.

  • Takeaways & Limitations

    TRAK requires differentiable models, and its effectiveness depends on the suitability of its linear approximation.

Abstract

from arXiv · show

The goal of data attribution is to trace model predictions back to training data. Despite a long line of work towards this goal, existing approaches to data attribution tend to force users to choose between computational tractability and efficacy. That is, computationally tractable methods can struggle with accurately attributing model predictions in non-convex settings (e.g., in the context of deep neural networks), while methods that are effective in such regimes require training thousands of models, which makes them impractical for large models or datasets. In this work, we introduce TRAK (Tracing with the Randomly-projected After Kernel), a data attribution method that is both effective and computationally tractable for large-scale, differentiable models. In particular, by leveraging only a handful of trained models, TRAK can match the performance of attribution methods that require training thousands of models. We demonstrate the utility of TRAK across various modalities and scales: image classifiers trained on ImageNet, vision-language models (CLIP), and language models (BERT and mT5). We provide code for using TRAK (and reproducing our work) at https://github.com/MadryLab/trak .

1 Introduction

Data attribution traces model predictions to training examples, but existing methods trade computational efficiency against efficacy. TRAK addresses this tension by using kernel-based approximations to achieve effective attribution at substantially lower computational cost.

  • Motivation: Training data influences model errors, biases, capabilities, and performance, motivating methods that trace predictions back to training examples.Such attribution supports explaining predictions, debugging behavior, valuing data, detecting poisoned or mislabeled examples, and curating datasets.
  • Motivation: Existing attribution methods trade computational demand against efficacy: efficient methods can be unreliable for non-convex models, while effective sampling methods require thousands of model trainings.This tradeoff motivates the question of whether attribution can be both scalable and effective in large-scale non-convex settings.
  • TRAK: TRAK approximates models with a kernel machine and uses the resulting kernel domain to derive data attribution scores.It is designed for parametric, differentiable models.
  • Results: 100-1000x faster to compute, TRAK can be as effective as datamodels on CIFAR-10 image classification and QNLI natural language inference.The comparison concerns attribution efficacy and runtime on the evaluated tasks.
  • Results: TRAK is as fast as gradient-based methods such as TracIn and influence-function variants while being significantly more predictive of model behavior.The paper applies TRAK across image classifiers, CLIP, BERT, and mT5.

2 Motivation and Setup

The paper formalizes data attribution as assigning training-example importance to model outputs and evaluates methods through counterfactual prediction. Its linear datamodeling framework exposes both the usefulness of additive attribution scores and the prohibitive cost of retraining-based evaluation.

  • Setup: Data attribution assigns each training example a real-valued importance score for a model output on an example of interest.The output depends on the trained model obtained from the full training set.
  • Evaluating attribution methods: Existing evaluation strategies include manual inspection, estimation of ground-truth parameters, and auxiliary tasks, but each has subjectivity, expense, or task-specific sensitivity.Ground-truth parameters can be prohibitively expensive to obtain at scale, while auxiliary-task metrics may hinder comparisons across settings.
  • The linear datamodeling score (LDS): The proposed metric evaluates whether attribution scores can predict model outputs after training on different subsets of the training data.This reframes attribution quality as counterfactual prediction accuracy.
  • The linear datamodeling score (LDS): For additive attribution methods, the predicted output for a subset is the sum of the scores assigned to examples contained in that subset.The subset is represented through an indicator vector.
  • The linear datamodeling score (LDS): The linear datamodeling score measures correlation between true outputs from subset-trained models and attribution-derived predictions across randomly sampled subsets.Its oracle estimator fits scores using regularized empirical risk minimization after retraining on each subset.
  • The oracle limitation: Tens or hundreds of thousands of subset-trained samples may be required to estimate accurate linear predictors, making direct datamodel attribution expensive at scale.The same retraining bottleneck applies to sampling-based methods such as empirical influences and Shapley values.

3 TRAK: Tracing with the Randomly-Projected After Kernel

TRAK scales data attribution to large, differentiable, non-convex models by reducing them to tractable logistic-regression problems. It combines model linearization, random projection, influence estimation, ensembling, and sparsification while retaining the canonical logistic-regression attribution primitive.

  • Core idea: TRAK extends the one-step Newton attribution method from logistic regression to large-scale, differentiable, non-convex models.The method uses the logistic-regression attribution primitive after transforming the original model into a tractable surrogate.
  • Scope: The framework is presented for binary classifiers and generalized to multi-class classifiers, contrastive models, and language models.For binary classifiers, the model output is the raw logit; more complex losses can involve both inputs and labels.
  • Core idea: TRAK linearizes the model output around the final parameters, converting a nonlinear model into a linear model whose inputs are parameter gradients.The Taylor approximation uses f(z; θ⋆) and ∇θ f(z; θ⋆), yielding a generalized logistic-regression view with gradient features and adjusted biases.
  • Core idea: Random projection reduces gradient features from the original parameter dimension p to k ≪ p while preserving inner products with high probability.The projection relies on the Johnson–Lindenstrauss result and defines a k-dimensional feature map for each example.
  • Implementation: The complete TRAK procedure estimates attribution scores with the transformed logistic-regression method, ensembles models trained on random subsets, and soft-thresholds the resulting scores.These steps form the implementation sequence after linearization and dimensionality reduction.

4 Evaluating TRAK

TRAK is evaluated across image and language tasks using linear datamodeling scores and computational-cost measures. It achieves stronger efficacy–efficiency tradeoffs than the compared attribution methods, while its identified examples and scores align with model behavior and datamodels.

  • Experimental setup: TRAK is evaluated on CIFAR-2, CIFAR-10, ImageNet, and QNLI using existing attribution methods as baselines.The comparisons include influence functions, TracIn, gradient-based methods, representation similarity, empirical influences, and datamodels.
  • Evaluation protocol: The evaluation measures attribution efficacy with the linear datamodeling score, which compares attribution-derived predictions with averaged model outputs across random training subsets.The protocol samples 100 half-sized training subsets, trains five models per subset, and computes Spearman rank correlations.
  • Results: TRAK attains efficacy comparable to datamodels with an average computational footprint over 100x smaller across the evaluated models and datasets.The paper reports a significantly better efficacy–efficiency tradeoff than all examined attribution methods.
  • Results: Figure 2 compares computational cost, measured by trained-model count, against LDS, with error bars showing 95% bootstrap confidence intervals.The figure covers CIFAR-2, CIFAR-10, ImageNet, and QNLI.
  • Efficiency optimizations: Using non-converged models or multiple checkpoints from one training run can dramatically reduce TRAK’s cost without significantly degrading performance.The checkpoint experiment uses 100 checkpoints, while the non-convergence experiment varies the training termination epoch.

5 Applications of TRAK

TRAK is applied across vision, vision-language, and language-model settings to trace training examples that influence model behavior. These applications show both counterfactual usefulness and practical acceleration of datamodel-based analyses.

  • Attributing CLIP models: TRAK identifies training examples that help or detract from aligning CLIP image and caption embeddings.Its attributions can differ from nearest-neighbor or similarity-based selections.
  • Fact tracing language models: TRAK outperforms TracIn on FTRACE-TREX fact tracing with 0.42 versus 0.09 MRR, but trails BM25 at 0.77 MRR.Counterfactual retraining shows that TRAK-selected abstracts reduce accuracy by 34%, compared with 10% for BM25 and 12% for ground-truth proponents.
  • Fact tracing language models: The FTRACE-TREX benchmark may conflate tracing supporting text with tracing examples that causally matter to the model’s learned behavior.Models may combine multiple training facts, while benchmark labels identify direct proponents.
  • Accelerating datamodel applications: TRAK scores can replace datamodel scores in downstream applications, matching 50,000-model datamodel brittleness estimates with an ensemble of 100 models and recovering similar MODELDIFF features.These applications use TRAK to accelerate prediction-brittleness estimation and learning-algorithm comparison.

6 Related work

TRAK connects data attribution to prior work on attribution methods, neural tangent kernels, and kernel approximation. Its position is distinguished by combining linearized-model reasoning with random projections for scalable attribution.

  • Training data attribution: Random projections have been used to accelerate influence estimation and Fisher-matrix-based data attribution in generalized linear and small-model settings.These prior experiments were limited to small neural networks and linear models.
  • Training data attribution: Prior attribution methods commonly trade computational scalability against predictive efficacy.Influence approximations and gradient agreement are attractive but can be unreliable, while sampling-based methods are more successful but costly.
  • Training data attribution: TRAK differs from Hessian-based influence functions by computing the Hessian with respect to a linearized model, yielding a positive semi-definite matrix.Both approaches leverage a one-step Newton approximation for leave-one-out influence estimation.
  • Neural tangent kernel: The neural tangent kernel has primarily supported analyses of generalization, optimization, robustness, and related learning applications.TRAK instead uses the empirical NTK approximation for data attribution in finite-width models.
  • Neural tangent kernel: TRAK’s reduction to the linear case is analogous to prior kernel-based approaches for estimating leave-one-out errors.The paper places this reduction alongside broader uses of gradients as features, including Fisher-kernel methods.
  • Kernel methods and random projections: Random projections are a widely used kernel-method technique that can improve computational efficiency and provide insight into empirical phenomena.Prior work has used the kernel view and random projections to explain subspace-based fine-tuning efficacy.

7 Discussion & Conclusion

The paper formalizes data attribution and introduces TRAK as an effective, efficiently scalable method for differentiable models. It demonstrates the approach across image, language, and image-text settings while noting differentiability and linear-approximation requirements.

  • Discussion & Conclusion: TRAK is introduced as an effective and efficiently scalable data attribution method.The paper evaluates it on CIFAR and ImageNet image classifiers, BERT and mT5 language models, and CLIP image-text models.
  • Discussion & Conclusion: TRAK requires differentiable models, and its effectiveness depends on the suitability of the linear approximation.The authors note that prior language-modeling results and their experiments suggest the approximation may continue to work for larger models.

A.1 Datasets and models

The appendix documents datasets, model configurations, attribution baselines, and computational-cost measurements used in the evaluation. Experiments span image classification, natural language inference, CLIP, fact tracing, and learning-algorithm comparison.

  • Datasets and models: CIFAR experiments use CIFAR-2 and CIFAR-10 with a ResNet-9 architecture.CIFAR-2 contains only the cat and dog classes and was initially used as a smaller binary-classification test bed.
  • Datasets and models: ImageNet experiments use the full 1000-class dataset with a modified ResNet-18 trained from scratch for 15 epochs.The configuration uses cyclic learning rates, momentum, weight decay, and label smoothing.
  • Datasets and models: QNLI experiments fine-tune bert-base-cased and remove the final tanh nonlinearity before classification.The modification prevents saturated model outputs and results in higher LDS.
  • Datasets and models: The CLIP setup uses a ResNet-50 image encoder and Transformer text encoder trained on MS COCO for 100 epochs.Images are resized to 224 × 224 and trained with Adam, augmentation, and a cosine learning-rate schedule.
  • Datasets and models: The evaluation also includes mT5-small fact tracing on FTRACE-TREX and MODELDIFF on LIVING17.MODELDIFF compares standard augmentation against no augmentation using ResNet-18 models.
  • Attribution methods and evaluation: TRAK has two main hyperparameters: projection dimension k and number of models M.Soft-thresholding is an optional hyperparameter used for non-binary classification tasks, while attribution baselines include influence functions, TracIn, GAS, empirical influences, representation similarity, and datamodels.
  • Computational evaluation: Computational cost is measured using total wall-time and total number of trained models.The appendix identifies training and gradient-computation time as major runtime components.

B TRAK implementation

TRAK’s implementation trains or acquires checkpoints, computes and randomly projects per-example gradients, and aggregates them into attribution scores. The appendix also derives its one-step influence approximation and states conditions governing its use.

  • Implementation pipeline: TRAK computes scores through training models, gradient computation, random projection, and final aggregation.The library automates the latter three steps, using per-example gradients and GPU-based projection or a custom CUDA kernel.
  • Random projection: Random projection is a major cost for large models because storing a 300-million-by-1024 projection matrix would require more than 1TB.The implementation generates projection coefficients as needed instead of storing the matrix.
  • Random projection: GPU implementation choices produce speed-ups exceeding 200x over the block-by-block projection strategy.The implementation uses block decomposition, shared memory, and Tensor Cores on NVIDIA GPUs.
  • One-step Newton approximation: TRAK estimates leave-one-out influence by applying a single Newton step from the trained optimum after excluding one example.The resulting parameter change is converted into a prediction change using an inner product with the input vector.
  • One-step Newton approximation: The one-step approximation assumes that updates converge in one step, with its fidelity quantified under assumptions by prior work.The authors report that TRAK’s performance across settings suggests accuracy in practical regimes.
  • Generalization: TRAK can extend beyond logistic regression when the training loss is convex in the model output.The corresponding loss derivatives replace the logistic-regression terms, while omitted leverage-score corrections may improve performance beyond classification.
  • Generalization: Choosing the model output function carefully is necessary because poor linear approximation limits predictive performance in the linear datamodeling task.This creates a method-level scope boundary for the framework.

C.3 Subsampling the training set

TRAK connects subsampled-model attribution to empirical influences by averaging leave-one-out estimates across random training subsets. Its predictive performance generalizes across subset fractions and tracks datamodel performance.

  • C.3 Subsampling the training set: TRAK averages single-subset leave-one-out estimates across random training subsets to approximate empirical influences.This follows from its connection to empirical influences and the need to average estimates from different subsets.
  • C.3 Subsampling the training set: TRAK scores computed with α = 0.5 remain highly predictive when evaluated across other subset fractions, with some degradation in correlation.The evaluation in Figure D.1 compares α = 0.5 scores with other regimes, including α = 0.75.
  • C.3 Subsampling the training set: The method’s evaluation treats attribution scores as linear predictors, reflecting a latent linearity assumption shared by many prior attribution methods.Datamodels, Shapley values, influence functions, and related approaches use additive or summed influence predictions.
  • C.3 Subsampling the training set: TRAK uses task-specific model outputs, including logit-equivalent margins for classification and language modeling and a CLIP-loss-based output for CLIP.The selected output function is intended to linearize as much of the model output as possible without over-linearizing it.
  • C.3 Subsampling the training set: TRAK and datamodel LDS values are correlated across examples, and TRAK generally performs better on examples where datamodels perform better.This comparison uses LDS correlations across examples rather than a single aggregate score.

D.2 Table for LDS evaluation

The LDS evaluation compares attribution methods by predictive accuracy and computational efficiency, while supplementary figures illustrate TRAK’s selected examples and algorithm-acceleration use cases.

  • D.2 Table for LDS evaluation: Table D.2 evaluates data attribution methods by linear datamodeling score and computational efficiency.Efficiency is measured by wall-clock time or the number of trained models, with 95% bootstrap confidence intervals.
  • D.2 Table for LDS evaluation: TRAK examples are visualized for ImageNet classifiers, QNLI examples, and CLIP models trained on MS COCO.The figures show top attributions or helpful and detracting training examples identified by TRAK.
  • D.2 Table for LDS evaluation: TRAK is applied to dramatically accelerate the MODELDIFF algorithm in a comparison of learning algorithms with and without data augmentation.The MODELDIFF case study uses models trained on the LIVING17 dataset.

E Ablation Studies

The ablations identify ensemble size, training stage, projection dimension, estimator terms, and feature choice as key determinants of TRAK’s predictive performance and efficiency.

  • Projection dimension: Higher projection dimensions initially improve LDS by preserving inner products more accurately, but beyond a point they reduce LDS.The peak projection dimension increases with the number of ensembled models M.
  • Model ensembling: More ensemble models substantially improve attribution performance, while a single model still reaches an LDS of 0.096 on CIFAR-2.This single-model result exceeds all evaluated prior gradient-based methods in the cited experiment.
  • Proxies for ensembles: Multiple checkpoints from the same training trajectories can provide comparable LDS to independently trained models, and early checkpoints retain efficacy before convergence.These proxies are studied as cheaper alternatives to fully independent ensembles.
  • Estimator components: The reweighting matrix and loss-gradient term substantially improve performance, whereas including the diagonal R term lowers LDS.Separately averaging the estimator and loss-gradient term, then rescaling by the average gradient, also yields higher LDS.
  • Kernel and feature choice: Using penultimate-layer representations improves over gradient-approximation baselines but remains substantially worse than TRAK, whose last-layer-only variant resembles influence-function approximation.The last-layer-only and influence-function variants both achieve an LDS of 0.19 on QNLI.
  • E Ablation Studies: TRAK performs best with enough ensemble models, later-epoch gradients, and random projections that are high-dimensional but not excessively so.Computationally cheaper ensemble alternatives can still work, while full-model gradients and the reweighting matrix remain critical.

F.4 Counterfactual experiment setup

The counterfactual experiment tests whether removing training examples selected by TRAK, BM25, or benchmark-defined proponents changes model behavior on novel facts. TRAK-selected removals produce the largest reported behavioral change.

  • F.4 Counterfactual experiment setup: The experiment constructs three counterfactual training sets by removing TRAK-selected examples, BM25-selected examples, or FTRACE-TREX proponents.Each set is evaluated on 50 selected novel facts after finetuning from the same pretrained mT5-small model.
  • F.4 Counterfactual experiment setup: Removing FTRACE-TREX proponents is expected to reduce accuracy because the resulting training set lacks direct evidence for the selected facts.The BM25-based set instead retains high lexical overlap with the facts of interest.
  • F.4 Counterfactual experiment setup: The TRAK-based counterfactual changes model behavior significantly more than removals based on BM25 or ground-truth proponents.The comparison uses average performance on the selected set of 50 novel facts.
  • F.4 Counterfactual experiment setup: TRAK’s counterfactual analysis supports using attribution to trace model behavior and motivates applications including dataset optimization and fact tracing.The paper discusses selecting or filtering training data and tracing outputs back to sources as potential applications.
  • F.4 Counterfactual experiment setup: The selected novel facts are cases handled incorrectly by pretrained mT5-small but correctly after finetuning.By construction, pretrained accuracy is 0% and full-dataset finetuning yields 100% accuracy on this subset.

G.2 Understanding and improving the TRAK estimator

The discussion identifies approximation and training-dynamics gaps that may limit TRAK’s efficacy, and outlines ensembling and optimization-aware extensions as improvement directions.

  • Empirical NTK: TRAK’s empirical NTK approximation may limit efficacy; higher-order approximations beyond its linear approximation are a proposed direction.The paper suggests that better understanding when the empirical NTK approximation is accurate could improve TRAK.
  • Training dynamics and optimization: TRAK models the second stage of neural-network training, while omitting the rapidly evolving features of the first stage.Understanding the first stage may explain the remaining gap between true model outputs and TRAK’s predictions.
  • Training dynamics and optimization: Accounting for mini-batches, momentum, or weight decay during training may improve TRAK’s estimator.
  • Ensembles: Computing TRAK over an ensemble of models significantly improves efficacy, with independently trained models’ eNTKs capturing non-overlapping information.
Loading 2303.14186v2…