Bergson Documentation

Bergson is a library for gradient-based data attribution of transformers. Data attribution methods estimate the effect on a behavior of interest of removing data points from a model’s training corpus, and enable data filtering, re-weighting, and interpretability.

Naively computing the effect of removing every subset of a corpus of N items requires 2**N retraining runs. Our most costly and powerful method, MAGIC, uses compute equivalent to 3-5 training runs to produce per-token or per-sequence scores that correlate with the effects of leave-k-out retraining at ρ>0.9 in well-behaved settings. More efficient methods like EK-FAC and TrackStar use compute equivalent to ~1-2 training runs (with more modest VRAM usage), but correlate less with leave-k-out retraining (ρ~=0.1 to 0.5).

We provide options for analyzing models and datasets at any scale or level of granularity:

  • Compressed or uncompressed gradients.

  • Per-token or per-sequence attribution.

  • On-disk gradient stores or on-the-fly queries.

  • HuggingFace Transformers models and Datasets, including on-disk datasets in a variety of formats.

  • Query aggregation following LESS and other strategies.

  • On-GPU gradient store queries, or sharded FAISS indexes for fast queries at scale.

  • Collect gradients during or after training.

  • Parallelize Bergson operations across multiple GPUs or nodes.

  • Load gradients with or without their module-wise structure.

  • Split attention module gradients by head.

Installation

pip install bergson

Quickstart

Build an index of gradients:

bergson build runs/quickstart --model EleutherAI/pythia-14m --dataset NeelNanda/pile-10k --truncation

Load the gradients:

from pathlib import Path
from bergson import load_gradients

gradients = load_gradients(Path("runs/quickstart"))

Reference

Content Index

If you have suggestions, questions, or would like to collaborate, please email lucia@eleuther.ai or drop us a line in the #data-attribution channel of the EleutherAI Discord!