MAGIC (Unrolling)¶

MAGIC attributes evaluation loss to individual training examples by backpropagating through the entire training trajectory.

We provide a Trainer class that takes differentiable training steps and handles all three phases of MAGIC attribution. We support FSDP training using the bergson.magic.dtensor_patch runtime patch, which makes PyTorch’s DTensor redistribution twice-differentiable (pytorch/pytorch#160509). The patch is applied in memory, so no torch source files are modified.

How it works¶

MAGIC attribution has three phases:

Forward training with checkpoints: Fine-tune the model, saving intermediate checkpoints at each step.
Evaluate: Compute the evaluation loss and its gradients with respect to the final model parameters.
Backward through training: Backpropagate the evaluation gradients through the checkpointed training steps using reverse-mode autodiff, accumulating attribution scores for each training example (or token).

The Trainer class handles all three phases. It uses torchopt for functional (stateless) differentiable optimization.

Usage¶

CUDA_VISIBLE_DEVICES="0" bergson magic runs/magic-ckpts \
    --data.dataset NeelNanda/pile-10k \
    --query.dataset NeelNanda/pile-10k \
    --query.split "train[:8]" \
    --model EleutherAI/pythia-14m

Contrastive queries¶

When --contrast.dataset is set the aggregated query loss minus the aggregated control loss is attributed. Requires an aggregated query. See examples/powerseeking_contrastive_query for an example.

Output files¶

After a run completes, run_cfg.run_path contains:

config.yaml — the serialized configuration needed to replicate the run.
checkpoints/ — forward-pass trainer checkpoints, plus log_history.json (the per-step learning rates the schedule produced). With the default cleanup_ckpts=True the backward pass deletes checkpoints as it consumes them, but multi-query runs keep them for re-use.
scores/ — a score directory, the same self-describing format the scoring pipeline writes. info.json records attribute_tokens and num_scores. Read it with load_scores_loss_signed(), which returns (scores, multi_query):
- Per-document: (num_train_docs, 1), indexed directly by doc_id.
- Per-token: (num_examples, seq_len), indexed by (example_idx, token_idx), with every epoch’s pass over an example summed into its row. An example is a chunk of the training dataset when chunk_length > 0 or otherwise a row/document, and example_idx is its position in the training set after chunking, before each epoch is shuffled.
- Per-query (query.aggregation: none) adds a trailing query axis, so per-token per-query scores are (num_examples, seq_len, num_query_docs).
Per-token scores are stored ragged — a row holds length - 1 values — and are unpacked back into a dense grid on load.
score_vs_step.png — if matplotlib is installed (pip install 'bergson[plot]') the median log10|score| of each optimizer step against training step is emitted, for assessing run stability.
per_query/q{i}.pt — per-query runs only. The score tensor for query document i, written as soon as that query’s backward finishes so an interrupted run resumes without redoing completed queries. The trailing query axis in scores/ is these tensors stacked.
scores/doc_ids.npy — written for per-token runs, shape (num_examples, seq_len). Each entry is the document id at that token position, so the scores can be summed over examples and token positions into per-document scores:
```
from bergson.data import load_scores_loss_signed

scores, _ = load_scores_loss_signed("runs/magic/scores")
doc_ids = torch.from_numpy(np.load("runs/magic/scores/doc_ids.npy"))
num_train_docs = int(doc_ids.max()) + 1

# Trailing axis is the query axis on per-query runs, absent otherwise;
# reshaping to it keeps both cases on one path.
flat = scores.reshape(doc_ids.numel(), -1)
per_doc = torch.zeros(num_train_docs, flat.shape[1], dtype=flat.dtype)
per_doc.scatter_add_(0, doc_ids.flatten()[:, None].expand_as(flat), flat)
```
per_doc comes back as (num_train_docs, num_query_docs), or (num_train_docs, 1) for a single-query run.

With data.chunk_length == 0 each example is one document and scores for document i can be accessed like scores[i]. With data.chunk_length > 0 a chunk may hold several documents, and a document may span chunks. In this case, use doc_ids to aggregate the scores for a document over chunks.

Models and optimizer state¶

optimizer.pt — second moments of the trained optimizer state, when save_optimizer_state is not "none".
retrained/base/ and retrained/subset_{i}/ — the fully-trained model and each leave-subset-out retrain, with tokenizer, when save_models=True. retrained/base is the query baseline evaluate_retrained reads; the subset_{i} retrains only exist for a run that actually validates (bergson validate, or bergson magic --skip_validation False). (Commands outside the leave-k-out family — plain bergson train — write the trained model to model/ instead.)

Smoothness¶

MAGIC is valid when the model training function you differentiate through is smooth with respect to the data weightings (metasmooth). There a few heuristics known to encourage smoothness:

Use the Muon optimizer
Increase batch size
Scale model outputs down
Clip gradients
Pre-activation batch norm
QK norm
Tune weight decay

Many of these methods boil down to “Identify and manage spikes in your training loss.” You can measure your smoothness with bergson metasmoothness.

Core components¶

Trainer: Functional trainer that supports forward training with checkpoints and backward-through-training.

from bergson.magic import BackwardState, DataStream, Trainer, TrainerState
import torchopt

# Initialize
opt = torchopt.adam(lr=1e-4)
trainer, state = Trainer.initialize(model, opt)

# Forward training with checkpoints
stream = DataStream(dataset, batch_size=4, device="cuda")
state = trainer.train(state, stream, save_dir="checkpoints/")

# Compute eval gradients, then backward through training
bwd_state = trainer.backward("checkpoints/", stream, bwd_state, state)
scores = bwd_state.weight_grads  # attribution scores

DataStream: Wraps a dataset with differentiable per-document (or per-token) weights that receive gradients during the backward pass.

# Per-document attribution
stream = DataStream(dataset, batch_size=4, device="cuda")

# Per-token attribution
stream = DataStream(dataset, batch_size=4, device="cuda", weight_shape=(len(dataset), max_length))

DTensor patch: For multi-GPU runs with FSDP, apply the DTensor patch before any distributed operations:

from bergson.magic.dtensor_patch import apply_dtensor_patch
apply_dtensor_patch()

# Your MAGIC worker call here

Per-token vs per-document attribution¶

By default, DataStream creates a 1D weight tensor [num_train_docs] for per-document attribution, indexed by doc_id. By passing a 2D tensor [num_examples, max_length] as the weight_shape parameter, each token receives its own attribution score. The weighted_causal_lm_ce loss function supports both shapes.

To use per-token attribution, set model.loss_function = weighted_causal_lm_ce so the model uses the weighted loss during training.

from bergson.utils.math import weighted_causal_lm_ce
model.loss_function = weighted_causal_lm_ce

Key implementation details¶

Functional optimization: torchopt.adam (or similar) provides a pure-function optimizer whose state is a pytree of tensors. This allows torch.autograd.grad to differentiate through optimizer updates.
Checkpoint strategy: By default, checkpoints are saved at sqrt(N) intervals, giving O(sqrt(N)) memory and O(N) recomputation cost. save_mode also supports all (O(N) space) and the original MAGIC paper’s log (O(log N) space, O(N log N) time).
FSDP compatibility: The DTensor runtime patch adds a NestedRedistribute autograd function that makes the FSDP all-gather/reduce-scatter differentiable through second-order backward passes.
Loss weighting: weighted_causal_lm_ce multiplies per-token cross-entropy by the DataStream weights before averaging. During backward-through-training, autograd accumulates gradients into these weights, yielding the attribution scores.