[docs & tutorials] [paper] [blog post] [website]

In our paper, we introduce a new data attribution method called TRAK (Tracing with the Randomly-Projected After Kernel). Using TRAK, you can make accurate counterfactual predictions (e.g., answers to questions of the form “what would happen to this prediction if these examples are removed from the training set?"). Computing data attribution with TRAK is 2-3 orders of magnitude cheaper than comparably effective methods, e.g., see our evaluation on:

If you use this code in your work, please cite using the following BibTeX entry:

@inproceedings{park2023trak,
  title = {TRAK: Attributing Model Behavior at Scale},
  author = {Sung Min Park and Kristian Georgiev and Andrew Ilyas and Guillaume Leclerc and Aleksander Madry},
  booktitle = {Arxiv preprint arXiv:2303.14186},
  year = {2023}
}

[Quickstart]

Check our docs for more detailed examples and tutorials on how to use TRAK. Below, we provide a brief blueprint of using TRAK's API to compute attribution scores.

from trak import TRAKer

model, checkpoints = ...
train_loader = ...

traker = TRAKer(model=model, task='image_classification', train_set_size=...)

for model_id, checkpoint in enumerate(checkpoints):
  traker.load_checkpoint(ckeckpoint, model_id=model_id)
  for batch in loader_train:
      # batch should be a tuple of inputs and labels
      traker.featurize(batch=batch, ...)
traker.finalize_features()

targets_loader = ...

for model_id, checkpoint in enumerate(checkpoints):
  traker.start_scoring_checkpoint(ckeckpoint, model_id=model_id, num_targets=...)
  for batch in targets_loader:
    traker.score(batch=batch, ...)

scores = traker.finalize_scores()

You can find several end-to-end examples in the examples/ directory.

To install the version of our package which contains a fast, custom CUDA kernel for the JL projection step, use

pip install traker[fast]

You will need compatible versions of gcc and CUDA toolkit to install it. See the installation FAQs for tips regarding this. To install the basic version of our package that requires no compilation, use

pip install traker

Please send an email to [email protected]

Kristian Georgiev
Andrew Ilyas
Guillaume Leclerc
Sung Min Park