[luciaquirke]

Adds a conjugate-gradient solver to ASTRA (astra.solver: cg) and replaces the leaderboard's EK-FAC + ASTRA row with three runs on current main.

The old row (LDS 0.643) used Kronecker factors fit before #495, a sum over the 4,608 training documents. With the document-averaged factors the published config (relative damping 0.1, lr 0.1) reaches NaN by step 50 on every query.

Leaderboard (GPT-2 / WikiText, 50 test queries, 100-subset bank)

Method	ASTRA Hessian scale	QLD	LDS	queries p<.05	per-token QLD
EK-FAC + ASTRA (CG)	1	0.077 [0.066, 0.089]	0.685 [0.664, 0.702]	50/50	1.014 [0.986, 1.043]
EK-FAC + ASTRA (paper scaling)	1/511	0.071 [0.060, 0.084]	0.629 [0.607, 0.647]	50/50
EK-FAC + ASTRA (consistent scaling)	1	0.064 [0.051, 0.077]	0.494 [0.468, 0.519]	44/50
EK-FAC		0.070 [0.058, 0.082]	0.454 [0.426, 0.479]	49/50	0.863 [0.833, 0.894]

• Paper scaling (ekfac_astra_paper.yaml): the settings of examples/replicate_bae_approx_unrolling_source/wikitext_gpt2_astra.yaml (absolute damping 0.011, astra.loss_reduction: mean, lr 0.011, momentum 0.9, batch 16, 300 steps, lr x0.9 every 100). The factors use the per-document summed loss and the Hessian-vector products the per-token mean, so ASTRA's Hessian is 1/511 of the factors' scale.
• Consistent scaling (ekfac_astra_consistent.yaml): the same settings without the loss_reduction override. The step size follows the paper's rule, the lowest final objective over test[0:2] from 1e-2 to 1e-8: 1e-2, 1e-3 and 1e-4 diverge, 1e-5 is lowest. Queries 2, 15 and 44 diverge at 1e-5.
• CG (ekfac_astra_cg.yaml): relative damping 10, 20 conjugate-gradient iterations on one batch of 256 documents.

CG is the README row; the leaderboard lists all three. The leaderboard.rst ASTRA note gives each variant's Hessian scale. The per-token README row and paired differences are recomputed for CG (MAGIC − CG 0.361 [0.342, 0.379], CG − EK-FAC 0.151 [0.138, 0.165], each 50/50).

Eigenvalue-corrected Shampoo + ASTRA (0.625) was produced the same way as the old EK-FAC row and is unchanged here.

Paper replication

examples/replicate_bae_approx_unrolling_source/wikitext_gpt2_astra.yaml on main (e741b8c), unmodified, all 481 validation queries, 5-seed retrain banks; the consistent row drops loss_reduction: mean and uses lr 1e-5.

Method	LDS, 481 queries	LDS, first 64	#494 (first 64)	Paper
EK-FAC (damping 0.011)	0.437 ± 0.014	0.435 ± 0.045	0.435	~0.44
ASTRA, paper scaling	0.498 ± 0.014	0.467 ± 0.046	0.467	~0.45
ASTRA, consistent scaling	0.441 ± 0.020	0.398 ± 0.060	0.397

Tests

test_cg_reaches_the_damped_solution: on the full batch with an identity preconditioner, CG matches (H + D)^-1 q to 1e-6.

🤖 Generated with Claude Code

https://claude.ai/code/session_018bMcJTgvvdNptVguk887JW