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"))
Methods
Pipeline & Tools
Reference
- Command Line Interface
AttentionConfigAttributorBuilderCollectorComputerDataConfigDataConfig.chunk_lengthDataConfig.completion_columnDataConfig.conversation_columnDataConfig.data_kwargsDataConfig.datasetDataConfig.decode_into_subclassesDataConfig.format_templateDataConfig.prompt_columnDataConfig.reward_columnDataConfig.skip_nan_rewardsDataConfig.splitDataConfig.subsetDataConfig.truncation
FaissConfigFiniteDiffGradientCollectorGradientCollector.backward_hook()GradientCollector.builderGradientCollector.cfgGradientCollector.dataGradientCollector.mod_gradsGradientCollector.preprocess_cfgGradientCollector.process_batch()GradientCollector.scorerGradientCollector.setup()GradientCollector.skip_indexGradientCollector.teardown()
GradientProcessorGradientProcessor.hessiansGradientProcessor.hessians_eigenGradientProcessor.include_biasGradientProcessor.load()GradientProcessor.normalizersGradientProcessor.projection_dimGradientProcessor.projection_scaleGradientProcessor.projection_targetGradientProcessor.projection_typeGradientProcessor.reshape_to_squareGradientProcessor.save()
InMemoryCollectorInMemoryCollector.backward_hook()InMemoryCollector.builderInMemoryCollector.cfgInMemoryCollector.dataInMemoryCollector.gradientsInMemoryCollector.mod_gradsInMemoryCollector.preprocess_cfgInMemoryCollector.process_batch()InMemoryCollector.scorerInMemoryCollector.scoresInMemoryCollector.setup()InMemoryCollector.skip_hessiansInMemoryCollector.teardown()
IndexConfigIndexConfig.attentionIndexConfig.attribute_tokensIndexConfig.auto_batch_sizeIndexConfig.decode_into_subclassesIndexConfig.filter_modulesIndexConfig.force_math_sdpIndexConfig.include_biasIndexConfig.label_smoothingIndexConfig.loss_fnIndexConfig.loss_reductionIndexConfig.max_batch_sizeIndexConfig.modulesIndexConfig.optimizer_stateIndexConfig.partial_run_pathIndexConfig.profileIndexConfig.projection_dimIndexConfig.projection_scaleIndexConfig.projection_targetIndexConfig.projection_typeIndexConfig.reshape_to_squareIndexConfig.split_attention_modulesIndexConfig.stream_shard_sizeIndexConfig.token_batch_size
ModuleGradientsPreprocessConfigQueryConfigScoreConfigScorerTokenGradientscollect_gradients()load_from_optimizer()load_gradient_dataset()load_gradients()load_module_gradients()load_token_gradients()mix_autocorrelation_matrices()- Benchmarks
- Limitations
Experiments
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!