🤖 machine learning

Bergson: An Open Source Library for Data Attribution

Bergson is an open-source library developed by EleutherAI that scales data attribution techniques to large language models and pre-training datasets by natively supporting on-disk gradient storage, multi-node distributed training, and providing the first open-source implementations of the MAGIC, SOURCE, and TrackStar methods.

Original authors: Lucia Quirke, Louis Jaburi, David Johnston, William Z. Li, Gonçalo Paulo, Guillaume Martres, Girish Gupta, Stella Biderman, Nora Belrose

Published 2026-06-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Lucia Quirke, Louis Jaburi, David Johnston, William Z. Li, Gonçalo Paulo, Guillaume Martres, Girish Gupta, Stella Biderman, Nora Belrose

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have built a giant, incredibly smart robot chef. This chef has learned to cook by reading millions of cookbooks (the training data). Sometimes, the chef makes a dish that tastes terrible, or perhaps it learns a weird trick it shouldn't know. You want to know: "Which specific page in which specific cookbook caused this?"

This is the problem of Data Attribution. It's like trying to trace a specific flavor in a soup back to the exact pinch of salt added three hours ago.

The paper introduces Bergson, a new open-source "kitchen toolkit" designed to help researchers solve this tracing problem for giant AI models. Here is how it works, using simple analogies:

1. The Problem: The "Black Box" Kitchen

Usually, when an AI learns, it's a black box. You put data in, and a model comes out. If the model does something weird, you can't easily tell why.

The Challenge: Modern AI models are so huge (like a library with billions of books) and the training process is so complex that figuring out which "book" caused a specific behavior is incredibly hard. It used to require so much computer memory that it was impossible to do for big models. Also, many of the best tools to do this were locked behind closed doors (not open-source).

2. The Solution: The Bergson Toolkit

Bergson is a library (a collection of software tools) that acts like a super-powered detective kit for AI. It lets researchers "rewind" the AI's training to see which data points had the biggest impact.

Think of it as a Time-Travel Camera for the training process. It offers three main ways to look back in time:

The "Perfect Replay" (MAGIC): Imagine you could record every single second of the chef's training. Then, you could hit "rewind" and play it back exactly, step-by-step, to see exactly what happened. This is the most accurate method, but it requires a lot of storage space. Bergson makes this possible even for huge models by using a clever trick: it only saves "snapshots" (checkpoints) at specific moments and fills in the gaps mathematically, saving space while keeping the picture clear.
The "Snapshot Approximation" (SOURCE): Sometimes, you can't record every second. Instead, you take a photo every hour. You assume the chef didn't change much between photos. This is faster and uses less memory, but it's a bit less precise. Bergson includes this method too.
The "Fingerprint Scanner" (TrackStar & Influence Functions): Instead of rewinding the whole movie, this method looks at the "fingerprint" left on the model's brain at the very end. It asks, "If I remove this specific ingredient, how much does the final dish change?" It's very fast and works well for huge datasets, though it's an estimate rather than a perfect replay.

3. Key Features of the Toolkit

The paper highlights several "quality of life" features that make this toolkit easy to use:

The "Disk Drive" Trick: Instead of trying to hold all the data in the computer's tiny, fast memory (RAM), Bergson stores the "fingerprints" on the hard drive. This means you can analyze massive models without crashing your computer.
The "Search Engine" (FAISS): When you ask, "Which data caused this error?", Bergson doesn't check every single book one by one. It uses a super-fast search engine (like Google for data) to instantly find the top 10 most likely culprits.
The "Adapter" Shortcut: Sometimes, you don't need to retrain the whole chef; you just need to tweak a small part of the recipe. Bergson supports "LoRA" (Low-Rank Adaptation), which lets you trace the impact of just a small add-on to the model, making the process much cheaper and faster.
The "Judge" Feature: What if you want to know why the chef is being rude? You can't measure "rudeness" with a simple math formula. Bergson allows you to use a second AI (a "Judge") to rate the chef's behavior, and then traces that rating back to the training data.

4. What They Actually Tested (The Case Studies)

The authors didn't just build the tool; they used it to solve three specific puzzles:

The "Biosecurity" Puzzle: They wanted to find the specific words in the training data that taught the AI dangerous biological knowledge (like how to make a virus). They used Bergson to highlight specific words. When they re-trained the AI but ignored those specific words, the AI became much safer.
The "Rude Robot" Puzzle: They used the "Judge" feature to find data that made the AI generate misaligned or weird behavior. By filtering out that data, they made the AI behave better.
The "Style" Puzzle: They wanted to know if the AI learned facts (like "Alice works at Fermilab") separately from the style it used to say it (like Shakespearean English vs. Legal jargon). They found that Bergson could successfully ignore the style and find the specific training data that taught the fact, proving the tool can separate "what" was learned from "how" it was said.

5. The Bottom Line

The paper claims that Bergson is the first open-source library that can handle these "time-travel" tracing tasks for the largest AI models currently in existence. It combines the most accurate methods (perfect replays) with the fastest methods (fingerprint scanning) and gives researchers the tools to debug, understand, and curate the data that makes AI smart.

In short: Bergson turns the mystery of "Why did the AI do that?" into a solvable detective story.

Technical Summary: Bergson: An Open Source Library for Data Attribution

Problem Statement

Data attribution aims to explain model behavior by determining the causal effect of individual training items on specific model outputs. While this field offers applications in debugging, dataset curation, and understanding memorization or reasoning, scaling these techniques to modern Large Language Models (LLMs) presents significant engineering challenges. Naive application of data attribution methods often results in unrealistic memory requirements, particularly when dealing with large-scale pre-training datasets or models requiring multi-node support. Furthermore, many state-of-the-art methods, such as those achieving high Leave-One-Out (LOO) correlations, lack open-source implementations, hindering reproducibility and progress.

Methodology and Library Architecture

Bergson is introduced as an open-source library designed to enable faster progress in data attribution by providing a composable, scalable architecture. The library supports a wide range of techniques, including differentiation through full training trajectories, approximate unrolled differentiation, and influence functions.

Core Components

Composable Pipeline: Bergson utilizes a data pipeline architecture (Figure 1) where components such as Hessian approximations, normalizers, index builders, and attribution scorers are interoperable. Workflows can be defined via a Command Line Interface (CLI) or programmatically using YAML configuration files.
Gradient Storage and Indexing: The library supports on-disk gradient stores indexed by module (including split attention heads) and supports per-token or per-sequence collection. To handle large indexes, Bergson integrates with FAISS for approximate nearest-neighbor (ANN) lookups, while also supporting exact top-k lookups.
Distributed Training: Bergson implements replicated and sharded parallelism (DDP, FSDP) to scale attribution to models with up to 405 billion parameters. It supports multi-node operations where FSDP groups can independently process dataset shards.
PEFT Integration: Native integration with the PEFT library allows for data attribution restricted to adapter parameters (e.g., LoRA), reducing the per-example gradient footprint by orders of magnitude.
Non-Differentiable Objectives: For tasks evaluated via LLM-as-a-judge, Bergson supports attribution via Group Relative Policy Optimization (GRPO) loss, allowing the attribution of behaviors that are not captured by standard cross-entropy losses.

Implemented Algorithms

Bergson provides the first open-source implementations for three leading methods:

MAGIC: An algorithm for exact differentiation through unrolled training loops using checkpointing. It achieves $O(N \log N)$ computation and $O(\log N)$ disk space (or a novel $O(N)$ compute/ $O(\sqrt{N})$ disk space trade-off).
SOURCE: A method for approximate unrolled differentiation that divides the training trajectory into segments, treating Hessians and gradients as constant within each segment.
TrackStar: A scalable approximation method for influence functions that projects per-example gradients into a low-dimensional space using random projections and utilizes autocorrelation Hessian approximations.

The library also supports standard influence function variants (K-FAC, EK-FAC, TK-FAC/Shampoo) and gradient normalization strategies (Adam, Adafactor, unit normalization).

Key Contributions

Open-Source Implementation: Bergson introduces the first public implementations of MAGIC, SOURCE, and TrackStar, filling a gap in tooling for cutting-edge data attribution methods.
Scalability: The library is engineered to handle billion-parameter models and massive datasets, supporting on-disk gradient stores, FAISS integration, and multi-node distributed training.
Method Selection Guide: The paper provides a guide for practitioners to select methods based on resource constraints, query volume, and the need for precision (e.g., recommending unrolled differentiation for safety-critical filtering and compressed indices for high-query debugging).
Evaluation Framework: Bergson includes tools for validation using the Linear Data Modeling Score (LDS) and variants for top-k/bottom-k filtering, as well as support for leave-k-out re-training experiments.

Experimental Results

The paper presents several validation experiments and case studies:

Accuracy Validation: On a GPT-2 (124M) model fine-tuned on WikiText-2, MAGIC (unrolled differentiation) achieved the highest Spearman correlation (0.979) with ground-truth leave-one-out effects, followed by EK-FAC (0.318) and TrackStar (0.164). However, in sorted subset leave-k-out evaluations (more relevant to data filtering), all methods achieved correlations of 0.7 or higher.
Token-Level Attribution (Biosecurity): Using the WMDP Bio benchmark, the authors fine-tuned a 7B model with LoRA and used per-token unrolled differentiation to identify influential tokens. Re-training with increased importance weights on the top 10% of influential tokens improved evaluation accuracy by an additional 1.5 percentage points compared to unweighted re-training.
Fact Tracing Across Styles: In a synthetic experiment involving biographical facts written in "Shakespearean" vs. "Legalese" styles, TrackStar successfully surfaced the correct factual proponents in 15.26% of top-5 results, compared to 0.80% for a raw gradient inner product baseline. TrackStar also significantly reduced "style leak" (matching the query style rather than the fact source).
Scaling Performance: Gradient collection time scales efficiently with the number of GPUs (1 to 16 A100s) for models ranging from 160M to 12B parameters, utilizing FSDP and data parallelism strategies to minimize communication overhead.

Significance and Claims

The authors claim that Bergson lowers the barrier to entry for data attribution research by providing a unified, scalable infrastructure. By offering the first implementations of MAGIC, SOURCE, and TrackStar, the library enables practitioners to experiment with state-of-the-art techniques that were previously inaccessible or required significant engineering effort to build.

The paper emphasizes that effective data attribution can aid in understanding model behaviors, debugging unexpected generations, and curating training datasets for safety and alignment. The authors note that while data curation is dual-use (capable of amplifying or reducing behaviors), the specificity of data attribution is expected to differentially benefit alignment and safety efforts by allowing precise targeting of relevant data. The library is positioned as a tool to enable more researchers to study interesting model behaviors through data attribution, particularly in the context of large-scale models where such analysis was previously impractical.