EleutherAI/bergson-wikitext-2-gpt2

GPT-2 fine-tuned on WikiText-2 with checkpoints and optimizer state

GPT-2 (gpt2) fine-tuned on the 4,656-document WikiText-2 training set from EleutherAI/bergson-wikitext-2-4656-chunks with the recipe of Bae et al. 2024, Training Data Attribution via Approximate Unrolled Differentiation (App. B.1): AdamW, lr 3e-5 constant, weight decay 0.01, batch size 8, 3 epochs (1,746 steps), fp32, seed 1004.

Checkpoints are saved at step 0 and every 291 steps with the full AdamW state, which is what SOURCE needs to unroll training.

What is here

path	what it is
retrained/base/	the final model in HF format
exported/checkpoint-{291,582,873,1164,1455,1746}/	each checkpoint in HF format with its optimizer.pt
checkpoints/step_*.ckpt/	the same checkpoints as saved by the bergson trainer (torch.distributed.checkpoint shards plus optimizer.pt)
checkpoints/log_history.json	learning rate and loss per step
config.yaml	the exact training configuration
source_scores/	SOURCE scores for the 481 validation queries (scores/), with the per-segment EK-FAC factors and preconditioners they were built from
if_scores/	EK-FAC influence scores at the final checkpoint (scores/), with the KFAC factors
lds/	per-query LDS of both scorers against the retrain banks (source_lds.csv, if_lds.csv, validate_*/)

Using it

from transformers import AutoModelForCausalLM
from huggingface_hub import snapshot_download

model = AutoModelForCausalLM.from_pretrained("EleutherAI/bergson-wikitext-2-gpt2", subfolder="retrained/base")

# all checkpoints, for SOURCE-style unrolling
path = snapshot_download("EleutherAI/bergson-wikitext-2-gpt2")

Ground truth

Leave-half-out retrain banks for this run, one per seed:

• EleutherAI/LDS-retrain-bank-adamw-wikitext2-N4656-bs8-seed1004
• EleutherAI/LDS-retrain-bank-adamw-wikitext2-N4656-bs8-seed1005
• EleutherAI/LDS-retrain-bank-adamw-wikitext2-N4656-bs8-seed1006
• EleutherAI/LDS-retrain-bank-adamw-wikitext2-N4656-bs8-seed1007
• EleutherAI/LDS-retrain-bank-adamw-wikitext2-N4656-bs8-seed1008

Measured against the five-seed ground truth: EK-FAC IF LDS 0.468 ± 0.015, SOURCE LDS 0.476 ± 0.015 (mean Spearman over 481 validation queries).

Replicating with bergson

Every artifact here is produced by the configs in examples/replicate_bae_approx_unrolling_source of bergson, run in this order with PYTHONPATH=$PWD python -m bergson <config>:

step	config	produces
0	prep_dataset.py	the training chunks, hosted at EleutherAI/bergson-wikitext-2-4656-chunks
1	wikitext_gpt2_train.yaml	checkpoints/, exported/, retrained/base/, config.yaml
2	wikitext_gpt2_source.yaml	source_scores/ (SOURCE scores, per-segment EK-FAC factors, preconditioners)
3	wikitext_gpt2_ekfac.yaml	if_scores/ (EK-FAC influence scores and KFAC factors)
4	wikitext_gpt2_retrain.yaml	the five retrain banks linked above
5	wikitext_gpt2_validate.yaml	lds/ (per-query LDS CSVs and the validate run configs)

The scores under source_scores/scores and if_scores/scores plus the banks' validation.csv files are enough to recompute the LDS table without a GPU; the checkpoints, factors and preconditioners let any single step be rerun in isolation. Part of the Data Attribution collection.