Bergson: An Open Source Library for Data Attribution

Abstract: Data attribution is a promising field in interpretability that aims to explain model behavior through the influence of its training data, with applications including debugging undesirable model behavior and training dataset curation. However, significant engineering effort is required to perform it at scale, and many cutting edge techniques lack open-source tooling and support. Bergson is an open source library that aims to enable faster progress in the field by providing a host of techniques that scale to very LLMs and pre-training datasets. The library natively supports on-disk gradient stores and multi-node distributed training, and provides quality of life tools for researchers. Finally, we introduce the first open-source implementations of three leading data attribution methods: MAGIC, SOURCE, and TrackStar. The library is available at https://github.com/EleutherAI/bergson .

• The paper introduces Bergson, a novel open-source toolkit that unifies multiple data attribution methods for deep learning models at scale.
• It leverages efficient gradient management, distributed training, and compressed influence functions to enable actionable debugging and model alignment.
• Empirical results validate its superior performance in token-level attribution and scalable leave-one-out influence estimation for billion-parameter models.

Bergson: An Open Source Library for Data Attribution

Overview

Bergson presents a comprehensive, open-source toolkit for data attribution in modern deep learning and LLM settings, unifying state-of-the-art techniques under a performant, reproducible, and highly composable architecture. The library operationalizes attribution at the necessary scale to study and intervene on model behavior, supporting memory- and communication-efficient gradient management, flexible distributed training, and seamless integration with common model and dataset toolchains. Notably, Bergson is the first open source implementation of MAGIC, SOURCE, and TrackStar, providing critical infrastructure for empirically-grounded attribution on production-scale models and datasets (2606.11660).

Figure 1: Bergson's composable pipeline allows users to declaratively combine training, gradient collection, compression, and attribution methods with robust support for distributed scalable execution.

Data Attribution Methods and Architectural Design

Bergson consolidates multiple attribution paradigms under a single framework, with modular pipeline composition over model interfaces, gradient/Hessian approximators, normalization and projection layers, and fast index systems. Attribution can operate at the sequence or token level, and supports per-module granularity. Core supported techniques include:

• Unrolled Differentiation (MAGIC): Enables exact first-order leave-one-out (LOO) influence estimation by differentiating through the full training process. Bergson introduces an $O(N \log N)$ computation ($O(\log N)$ disk) original implementation and a $O(N)$ compute/$O(\sqrt{N})$ disk alternative, dramatically reducing practical compute and storage barriers for this theoretically optimal attribution regime.
• Approximate Unrolling (SOURCE): Enables segment-wise piecewise-constant Hessian approximations along the training trajectory, allowing tradeoff between computational tractability and accuracy.
• Single-Checkpoint Influence Functions (EK-FAC, K-FAC, TrackStar): Methods deriving attributions via inverse Hessian-vector products at a given checkpoint, with K-FAC/EK-FAC Hessian approximations and TrackStar's gradient compression for high-throughput scaling.
• TrackStar: Combines block-wise and randomly projected compressed gradients, autocorrelated Hessian approximations, and cross-dataset mixing for scalable but high-fidelity fact tracing.
• LESS: Multi-checkpoint sequence aggregation of influence functions, empirically improving top-$k$ filtering for real-world fine-tuning protocols.

Multi-level normalization (e.g., Adam/Adafactor moments), sequence-to-token query aggregation, and flexible index building (via FAISS) enable robust performance and empirical tractability.

Operational Characteristics and Distributed Scaling

Bergson implements both on-disk and on-the-fly attribution workflows, supporting ANN querying, per-step aggregation, and massive-scale gradient management. Integration with HuggingFace Transformers and parameter-efficient fine-tuning (PEFT, LoRA)—with gradient collection at the adapter subspace—provides accessibility for realistic LLM settings. The design supports exactly reproducible pipelines (YAML, CLI, programmatic API), automatic serialization of all intermediate artifacts, and health checking/mechanical determinism validation.

Notably, the framework supports both DDP and FSDP (as well as SimpleFSDP to enable differentiable trainers), with block-diagonal Hessian sharding to reduce memory stress and enable multi-node computation without high synchronization overhead. Distribution across many GPUs and nodes enables attribution for 100B+ parameter models and 100M+-token datasets.

Figure 2: End-to-end attribution time as a function of training dataset size for several methods, showing close correspondence to base model training time even for scalable unrolled differentiation.

Figure 3: Gradient collection latency scaling with GPU count for models with hundreds of millions to billions of parameters, showcasing highly efficient parallelization in both replicated and sharded regimes.

Case Studies

Token Level Attribution and Biothreat Mitigation

Bergson enables fine-grained, token-level attributions, facilitating identification and targeted intervention on hazardous knowledge in LLMs. On WMDP biosecurity tasks, token attributions generated via unrolled differentiation successfully identified and amplified capability-relevant sequences: reweighting the 10% most influential tokens yielded an additional 1.5 percentage point improvement in evaluation accuracy relative to uniform weighting, exceeding sequence-level interventions.

Figure 4: Token-level attributions for biosecurity capabilities, with positive and negative influential tokens clearly localized and validated against behavior.

Attribution for Emergent Misalignment and Non-Differentiable Behaviors

Bergson's attribution extends to reinforcement learning objectives and other non-differentiable tasks via GRPO-style losses. For emergent misalignment, LLM-as-a-judge reward signals are used to assign attributions post-hoc. Selective filtering on high-attribution training examples rivals classifier-based adversarial data removal, demonstrating practical alignment applications.

Figure 5: Identification of training data causally responsible for emergent misalignment, attributed using a GRPO policy-gradient loss objective.

Fact Tracing Across Disjoint Styles

Bergson's compressed gradient and influence-function methods allow nuanced decomposition of content and style influence. Attributing legalese-formatted queries to Shakespearean-style facts, TrackStar attributes substantially outperform dense gradient inner products (Recall@5: 15% vs <1%), with markedly lower style leakage, highlighting the framework's utility for substructure tracing and debugging.

Figure 6: Fact tracing across distinct writing styles demonstrates Bergson's ability to surface content-relevant sources while suppressing superficial style matches.

Validation and Empirical Results

Extensive validation demonstrates the empirical superiority of unrolled differentiation (MAGIC) in approximating ground truth LOO effects (LDS Spearman correlation: $0.979 \pm 0.010$), with EK-FAC and TrackStar lagging. However, attribution fidelity converges across methods when sorted subset filtering is the target metric, supporting compression-heavy approaches for real-world curation and auditing.

Theoretical and Practical Implications

Bergson establishes robust infrastructure for large-scale, actionable data attribution, opening new avenues for:

• Systematic dataset curation optimized for specific behavioral attributes (capabilities, safety, memory, or bias removal)
• Causal auditing and debugging of model errors, misalignment, or memorization failures
• Mechanistic analysis of content propagation through pretraining, fine-tuning, and RL
• Rapid, robust evaluation of new attribution algorithms through canonical implementations

The modular, scale-robust, and PEFT-friendly design directly supports high-stakes stewardship of open- and closed-weight models, bridging the gap between theoretical attribution and operational deployment.

Future Directions

There are several extensions that could further enhance attribution research:

• Improved meta-smoothness enforcement for numerically robust unrolling in challenging optimization regimes
• Native support for multi-modal or retrieval-augmented LLMs
• Further reduction and characterization of the tradeoffs in gradient compression and storage for multi-billion parameter models
• Closer integration of attribution signals with downstream active data selection and continual learning schedules

Conclusion

Bergson delivers scalable, performant, and reproducible data attribution infrastructure for billion-parameter LLMs and trillion-token datasets, establishing a reference toolkit and baseline for the field (2606.11660). The library’s comprehensive support for the state-of-the-art, efficient distributed execution, and high-level user composability enable systematic study and control of training data influence, underpinning future advances in data-centric alignment, debugging, and model development.