catt-bench/mt5-small-ftrace-retrain

mT5-small FTRACE-TREx models: base run and remove-and-retrain replicates

From the paper Elucidating the Design Space of LM Data Attribution.

Archive of every mT5-small model trained for the FTRACE-TREx attribution experiments. Each run is its own subfolder with final_model/ (step 60,000), its training_config.json, manifest.json (checkpoints and training curve) and train_log.jsonl.

• sd1001/: the base model, fine-tuned on the full corpus (seed 1001), with intermediate checkpoints at steps 5,000, 10,000 and 30,000. All attribution scores are computed on it.
• cf50_{budget}_{method}_sd{seed}/: remove-and-retrain replicates. The run is retrained from google/mt5-small with the rows in excluded_rows.npy removed: the union, over 50 evaluation facts, of the training rows a given attribution method ranks highest. budget is the number of removed rows (b2000, b5000, b11522; no tag means the per-fact top-500 union), method the ranking (e.g. ekfac_ce, ekfac_margin_top10, trak_w, graddot_margin, bm25, gt for ground-truth proponents), and seed the retraining seed. emb marks removal sets whose TRAK rankings use the trak_emb features. cf50_full_sd* are the no-removal controls, retrained the same way.

Load a run with

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
repo = "catt-bench/mt5-small-ftrace-retrain"
model = AutoModelForSeq2SeqLM.from_pretrained(repo, subfolder="cf50_trak_w_sd2001/final_model")
tokenizer = AutoTokenizer.from_pretrained(repo, subfolder="cf50_trak_w_sd2001/final_model")