Clarify YAMLs

bergson/config/config.py

@@ -454,9 +454,6 @@ class ValidationConfig(TrainingConfig, ABC):
     num_subsets: int = 100
     """Number of leave-k-out subsets for Spearman correlation."""
 
-    subset_strategy: Literal["random"] = "random"
-    """Strategy for selecting leave-k-out subsets for validation."""
-
     subset_weight: float = 0.0
     """Training weight assigned to each subset's documents during the retrain
     (the rest stay at 1.0). ``0.0`` (default) is standard leave-k-out removal."""

bergson/magic/config.py

@@ -17,9 +17,9 @@ class MagicConfig(ValidationConfig):
     """Whether to compute attribution scores per token (instead of per sequence);
     the same toggle as ``IndexConfig.attribute_tokens``."""
 
-    skip_validation: bool = False
-    """Stop after computing and saving attribution scores, before the
-    leave-k-out retraining loop. Useful for score-only MAGIC runs."""
+    skip_validation: bool = True
+    """Set to False to run a leave-k-out retraining validation loop in the
+    same job."""
 
     # TODO(Lucia Quirke, December 2026): remove per_token backward compatibility.
     per_token: bool = False

bergson/validate.py

@@ -324,7 +324,7 @@ def validate_scores(
     if os.path.exists(subsets_path):
         with open(subsets_path) as f:
             subsets = [torch.tensor(s, dtype=torch.long) for s in json.load(f)]
-    elif run_cfg.subset_strategy == "random":
+    else:
         rng = torch.Generator().manual_seed(run_cfg.seed)
         if run_cfg.subset_fraction > 0:
             # Draw potentially overlapping samples
@@ -348,8 +348,6 @@ def validate_scores(
             subsets = list(perm.chunk(run_cfg.num_subsets))
             rng = random.Random(run_cfg.seed)
             rng.shuffle(subsets)
-    else:
-        raise ValueError(f"Unknown subset strategy: {run_cfg.subset_strategy}")
 
     csv_path = os.path.join(run_cfg.run_path, "validation.csv")
     val_csv_writer = CSVWriter(

docs/pipeline.rst → docs/building_blocks.rst

@@ -1,5 +1,5 @@
-Pipeline Concepts
-=================
+Building Blocks
+===============
 
 Bergson's post-hoc attribution exposes three generic building blocks — ``build``, ``reduce``,
 and ``score`` — that together implement gradient-based data attribution. This page

docs/cli.rst

@@ -149,7 +149,7 @@ Method Pipelines
        --method kfac \
        --hessian_cfg.ev_correction true
 
-See ``examples/magic/compare/q3_ekfac.yaml`` for a complete pipeline
+See ``examples/magic/compare_wikitext/ekfac.yaml`` for a complete pipeline
 configuration.
 
 .. autoclass:: bergson.__main__.ApproxUnrolling

docs/magic.rst

@@ -32,6 +32,12 @@ 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``, so consumers never infer the layout from the shape.
@@ -50,10 +56,15 @@ After a run completes, ``run_cfg.run_path`` contains:
   ``length - 1`` values, the positions ``weighted_causal_lm_ce`` can reach
   — and are unpacked back into the dense grid on load.
 
+* ``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 every per-token run, shape
-  ``(num_chunks, seq_len)`` matching the loaded scores row-for-row. Each
-  entry is the original (pre-shuffle) document id for that token position.
-  Downstream aggregation is one line:
+  ``(num_chunks, seq_len)`` matching the loaded scores. Each
+  entry is the original (pre-shuffle) document id for that token position,
+  so the scores can be aggregated over chunks:
 
   .. code-block:: python
 
@@ -62,8 +73,15 @@ After a run completes, ``run_cfg.run_path`` contains:
      scores, _ = load_scores_loss_signed("runs/magic/scores")
      doc_ids = torch.from_numpy(np.load("runs/magic/scores/doc_ids.npy"))
      num_docs = int(doc_ids.max()) + 1
-     per_doc = torch.zeros(num_docs, dtype=scores.dtype)
-     per_doc.scatter_add_(0, doc_ids.flatten(), scores.flatten())
+
+     # 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_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_docs, num_query_docs)``, or
+  ``(num_docs, 1)`` for a single-query run.
 
   When ``data.chunk_length > 0`` the ``doc_ids`` column comes from
   ``tokenize_and_chunk`` and chunks may pack multiple docs or split one
@@ -72,14 +90,23 @@ After a run completes, ``run_cfg.run_path`` contains:
   past the row's actual length carry zero MAGIC score and contribute
   nothing to the scatter-add.
 
-* ``config.yaml`` — serialized ``MagicConfig`` used for the run.
-* ``validation.csv`` — leave-subset-out validation results (if validation
-  was run).
+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.)
 
-Metasmoothness
----------------
+Smoothness
+----------
 
-MAGIC is valid when the function you are differentiating through is metasmooth. There a few heuristics known to encourage metasmoothness:
+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
@@ -89,7 +116,7 @@ MAGIC is valid when the function you are differentiating through is metasmooth.
 * QK norm
 * Tune weight decay
 
-Many of these methods boil down to "Identify and manage spikes in your training loss." You can measure your metasmoothness with ``bergson metasmoothness``.
+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
 ^^^^^^^^^^^^^^^

examples/compare_wikitext/README.md

@@ -0,0 +1,8 @@
+| method | mean ρ | median ρ | min | max | queries p<.05 |
+|---|---|---|---|---|---|
+| MAGIC (per-query) | 0.957 | 0.963 | 0.880 | 0.984 | 50/50 |
+| EK-FAC | 0.470 | 0.461 | 0.145 | 0.751 | 46/50 |
+| SOURCE-Adam | 0.243 | 0.253 | -0.038 | 0.434 | 30/50 |
+| TrackStar | 0.241 | 0.239 | -0.037 | 0.454 | 36/50 |
+| SOURCE | 0.221 | 0.212 | -0.089 | 0.421 | 27/50 |
+| TrackStar+Adam | 0.189 | 0.190 | -0.107 | 0.471 | 23/50 |

examples/compare_wikitext/ekfac.yaml

@@ -0,0 +1,73 @@
+# Pipeline: compute EK-FAC scores for the magic-trained model against
+# test[0:50] (one score column per query), then run leave-k-out validation on
+# those scores against the retrained models from magic.yaml.
+
+# Run with `bergson examples/compare_wikitext/ekfac.yaml`
+
+run_path: runs/compare_wikitext/ekfac
+steps:
+  - ekfac:
+      index_cfg:
+        run_path: runs/compare_wikitext/ekfac
+        model: runs/compare_wikitext/random/retrained/base
+        overwrite: true
+        # GPT-2's max context length is 1024.
+        token_batch_size: 1024
+        # Cap batch size in docs so the bin-packer doesn't produce
+        # large batches of very few tokens. This matters for EK-FAC
+        # because gradients aren't compressed.
+        max_batch_size: 64
+        # GPT-2 uses an nn.Linear for its LM head.
+        filter_modules: "lm_head"
+        distributed:
+          nproc_per_node: 8
+          nnode: 1
+        data:
+          dataset: EleutherAI/bergson-wikitext-512-chunks
+          split: "train"
+          chunk_length: 0
+
+      hessian_cfg:
+        method: kfac
+        ev_correction: true
+
+      score_cfg:
+        batch_size: 1024
+        query_batch_size: 64
+
+      preprocess_cfg:
+        unit_normalize: false
+
+      hessian_pipeline_cfg:
+        query:
+          dataset: EleutherAI/bergson-wikitext-512-chunks
+          split: "test[0:50]"
+          chunk_length: 0
+        # One score column per query, for per-query LDS.
+        query_aggregation: none
+        inversion_cfg:
+          damping_factor: 0.1
+
+  - validate:
+      run_path: runs/compare_wikitext/validate_ekfac_random
+      model: gpt2
+      overwrite: true
+
+      scores: runs/compare_wikitext/ekfac/scores
+      retrained_dir: runs/compare_wikitext/random
+
+      data:
+        dataset: EleutherAI/bergson-wikitext-512-chunks
+        split: "train"
+        chunk_length: 0
+
+      query:
+        dataset: EleutherAI/bergson-wikitext-512-chunks
+        split: "test[0:50]"
+        chunk_length: 0
+
+      distributed:
+        nproc_per_node: 8
+        nnode: 1
+
+      batch_size: 64

examples/compare_wikitext/magic.yaml

@@ -0,0 +1,54 @@
+# Train GPT-2 on WikiText, computing per-query MAGIC scores (one backward per
+# query, 50 queries) and leave-k-out validation over 100 random 1%-drop
+# subsets. bs 256 / 4 epochs / eps_root 1e-8 scored per-query LDS 0.93-0.98
+# in the 2026-08-03 batch-size sweep (runs/bs_eps1e8). The retrained models
+# are saved so every other yaml in this directory validates against the same
+# family, and the final Adam second moments are kept for trackstar_adam.yaml.
+
+# Run with `bergson examples/compare_wikitext/magic.yaml`
+
+run_path: runs/compare_wikitext/magic
+steps:
+  - magic:
+      run_path: runs/compare_wikitext/random
+      model: gpt2
+      overwrite: true
+      # Resume: keep finished per_query/q<i>.pt scores and fast-forward
+      # training from the last checkpoint after an interruption.
+      resume: true
+
+      data:
+        dataset: EleutherAI/bergson-wikitext-512-chunks
+        split: "train"
+        chunk_length: 0
+
+      query:
+        dataset: EleutherAI/bergson-wikitext-512-chunks
+        split: "test[0:50]"
+        chunk_length: 0
+
+      distributed:
+        nproc_per_node: 8
+        nnode: 1
+
+      batch_size: 256
+      # Bound double-backward memory; no trajectory change without dropout.
+      # Without the cap the retained graph OOMs 48GB GPUs (vocab-sized
+      # tensors per sequence).
+      grad_accum_steps: 2
+      double_backward_batch_size: 8
+      num_epochs: 4
+      lr_schedule:
+        lr_scheduler_type: polynomial
+        lr: 0.0008
+        lr_start: 1e-6
+        lr_end: 0.00008
+        warmup_steps: 0.25
+
+      skip_validation: false
+      num_subsets: 100
+      subset_fraction: 0.01
+
+      save_models: true
+      save_optimizer_state: last
+      wandb_project: magic

examples/compare_wikitext/source.yaml

@@ -0,0 +1,102 @@
+# SOURCE (approximate unrolling, Bae et al. 2024) scores for the
+# magic-trained model against test[0:50], validated against the retrained
+# models from magic.yaml. Plain variant: no optimizer preconditioning; see
+# source_adam.yaml for the Adam variant. Shares its evenly spaced interval
+# checkpoints (steps 12..72; sqrt mode never saves the final step) — the
+# train step resumes instantly when runs/compare_wikitext/interval already
+# exists.
+
+# Run with `bergson examples/compare_wikitext/source.yaml`
+
+run_path: runs/compare_wikitext/source
+steps:
+  - train:
+      run_path: runs/compare_wikitext/interval
+      model: gpt2
+      resume: true
+
+      data:
+        dataset: EleutherAI/bergson-wikitext-512-chunks
+        split: "train"
+        chunk_length: 0
+
+      distributed:
+        nproc_per_node: 8
+        nnode: 1
+
+      batch_size: 256
+      grad_accum_steps: 2
+      num_epochs: 4
+      lr_schedule:
+        lr_scheduler_type: polynomial
+        lr: 0.0008
+        lr_start: 1e-6
+        lr_end: 0.00008
+        warmup_steps: 0.25
+
+      save_mode: interval
+      save_interval: 12
+      save_optimizer_state: all
+
+  - approxunrolling:
+      index_cfg:
+        run_path: runs/compare_wikitext/source
+        model: runs/compare_wikitext/random/retrained/base
+        overwrite: true
+        precision: fp32
+        # GPT-2's max context length is 1024.
+        token_batch_size: 1024
+        # GPT-2 uses an nn.Linear for its LM head.
+        filter_modules: "lm_head"
+        distributed:
+          nproc_per_node: 8
+          nnode: 1
+        data:
+          dataset: EleutherAI/bergson-wikitext-512-chunks
+          split: "train"
+          chunk_length: 0
+
+      hessian_cfg:
+        method: kfac
+        hessian_dtype: fp32
+        ev_correction: true
+
+      approx_unrolling_cfg:
+        checkpoints:
+          - runs/compare_wikitext/interval/checkpoints/step_12.ckpt
+          - runs/compare_wikitext/interval/checkpoints/step_24.ckpt
+          - runs/compare_wikitext/interval/checkpoints/step_36.ckpt
+          - runs/compare_wikitext/interval/checkpoints/step_48.ckpt
+          - runs/compare_wikitext/interval/checkpoints/step_60.ckpt
+          - runs/compare_wikitext/interval/checkpoints/step_72.ckpt
+        segments: 3
+        query:
+          dataset: EleutherAI/bergson-wikitext-512-chunks
+          split: "test[0:50]"
+          chunk_length: 0
+        # One score column per query, for per-query LDS.
+        query_aggregation: none
+
+  - validate:
+      run_path: runs/compare_wikitext/validate_source_random
+      model: gpt2
+      overwrite: true
+
+      scores: runs/compare_wikitext/source/scores
+      retrained_dir: runs/compare_wikitext/random
+
+      data:
+        dataset: EleutherAI/bergson-wikitext-512-chunks
+        split: "train"
+        chunk_length: 0
+
+      query:
+        dataset: EleutherAI/bergson-wikitext-512-chunks
+        split: "test[0:50]"
+        chunk_length: 0
+
+      distributed:
+        nproc_per_node: 8
+        nnode: 1
+
+      batch_size: 64