TRAK: Attributing Model Behavior at Scale#
This is a PyTorch-based API for our method
effective, efficient data attribution method for gradient-based learning
algorithms. We designed
TRAK’s API around the following guiding
- Ease of use
You can apply
TRAKin just a few lines of code (see the quickstart guide).
Our API is lightweight - the entire codebase is less than 1000 lines of code. It is also quite modular, making it painless adapt any component to your needs.
TRAKto a custom task/modality is easy (check, e.g., how to adapt TRAK to CLIP).
See the code
The PyTorch-only version of our package can be installed using
pip install traker
To install the version of our package which contains a fast, custom CUDA kernel, use
pip install traker[fast]
See the Installation FAQs for more details.
TRAK is under active development. We are still in a
version and lots of things may change.