Add YAMLs for Adam WikiText SOURCE + EK-FAC replications - #369
Merged
Conversation
Segment eigenvalues were unnormalized sums over (documents x checkpoints). EK-FAC is unaffected -- every inversion is homogeneous of degree -1 in lambda because damping is relative to mean(lambda) -- but SOURCE's eigenfunctions are not scale-invariant, and sigma enters exp(-lr*K*sigma). Eq. 1 of Bae et al. 2024 defines the empirical risk as a mean over data points, so fisher_normalization defaults to per-document; token units put the mask argument in its dead-linear range. Also makes precompute_checkpoints resumable rather than overwrite-only, and lets resolve() fall back when a checkpoint's sibling config.yaml turns out to belong to an attribution run.
…C trainer (cherry picked from commit 16c29730ff6fdadd1c61ea72dc9dda43e6a85274)
The other save modes' backward replay indexes the data stream by checkpoint, so an extra trailing snapshot runs it past the end (DataStream index out of range in test_fsdp_ddp_scores_match).
The flag only toggled allow_compile; the backward still ran (and crashed on interval-mode end-of-training snapshots). Port the dummy-zero-scores skip path so bank-only builds go straight to validation. (cherry picked from commit 13b631d3b030be312a881ceff6c31475cbceb4b5)
Companion to the cherry-picked gate: defines the MagicConfig field and the do_metagradient condition it reads.
Bae et al. 2024 App. D: 'use the diagonal Hessian approximation for computing the matrix exponential in Equation 39 and Equation 43. Note that we still use the EK-FAC factors to compute H^-1 g in Equation 43.' Evaluating the eigenfunction wholly on the diagonal drops that second half. With adam_segment_hybrid the diagonal supplies only the exponential mask (f_one_minus_exp) and a chained EK-FAC inverse supplies 1/x, applied to the query as the adjoint H^-1 @ M_mask @ q.
Paper recipe from Bae et al. 2024 App. B.1 (AdamW, lr 3e-5, wd 1e-2, batch 8, seq 512, 6 checkpoints) with train_mode enabled so GPT-2's 0.1 dropout is active, matching kronfluence's train.py. SOURCE uses the App. C preconditioned variant with the App. D Eq-43 hybrid, L=3, per-document Fisher units.
Scoring side of a published retrain bank (kronfluence ships masks.pt / losses.pt for its WikiText-2 experiment): sum each subset's member scores, Spearman against the measured subset losses, average over queries. Lets a replication be checked against published numbers without retraining.
…ment save_run_config has always wrapped the step list in a document with a metadata block, so load_training_config rejected every real run directory; its test hand-rolled the bare list it expected instead of calling save_run_config. Have the test write configs the way the writer does, and report an unparseable payload as ValueError so callers guessing at a run dir still degrade to their fallback.
Export -> EK-FAC IF -> SOURCE -> LDS against kronfluence's shipped bank. EK-FAC 0.4451 against their reported 0.44; SOURCE 0.4506, above IF as in the paper's Figure 6.
luciaquirke
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
For EK-FAC, replication refers to kronfluence's
masks.pt(100 x 4656 alpha=0.5 keep-masks) andlosses.pt(100 x 481 five-seed-averaged subset losses). We replicate the relevant run with --train_mode. For SOURCE the results come from the paper, which doesn't specify an exact number but it looks right from the figure sketch.
Configs
examples/replication/:wikitext_gpt2_train.yaml: the paper's App. B.1 model training recipewikitext_gpt2_ekfac.yaml: EK-FAC influence functions at the final checkpoint.wikitext_gpt2_source.yaml: SOURCE, L=3 segments over C=6 checkpoints.experiments/wikitext_replication/run_chain.py: runs export -> IF -> SOURCE-> LDS.
Code changes
Trainer. add
save_mode: intervalfor SOURCE-style runs that want a few evenly spaced checkpoints, and per-snapshotoptimizer.ptso the AdamW SOURCE variant can read the second moments at each checkpoint.Eq-43 hybrid. The paper's App. D says to "use the diagonal Hessian approximation for computing the matrix exponential in Equation 39 and Equation 43. Note that we still use the EK-FAC factors to compute H^-1 g in Equation 43."
adam_segment_hybridimplements exactly that: the elementwise1 - exp(-lr*K*sigma)mask on the diagonal, then the standard EK-FAC inverse, rather than approximating the whole segment operator diagonally.Fisher normalization. The segment eigenvalues are now means rather than sums, which matters for SOURCE where eigenfunctions are not scale-invariant.
Fixes.
load_training_confignever read a real run directory:save_run_configwraps the step list in a{steps, metadata}document, and the loader only handled a bare list. Its test hand-rolled the shape it expected instead ofcalling the writer, so this went unnoticed.
save_mode: interval; undersqrtit broke MAGIC's backward replay.skip_metagradientwas accepted but not honoured.New.
lds_from_precomputed_subsetsscores attribution output against a published retrain bank.