txus/lsr1-artifacts/repro_bundle/code - Storage Bucket

Reproduction: L-SR1 — Learned Symmetric-Rank-One Preconditioning

Independent reproduction of ICML 2026 paper L-SR1: Learned Symmetric-Rank-One Preconditioning (OpenReview w1fkDwiZgN, arXiv 2508.12270; Lifshitz, Zuler, Fouks, Raviv, Tel Aviv University).

Part of the HF × AlphaXiv ICML-2026 open-reproductions effort. No official code / checkpoints / data exist — everything here is re-implemented from the paper text, Algorithm 1, Eqs. 8–9 and Tables 4–6.

What L-SR1 is

A lightweight, self-supervised learned second-order optimizer. It extends the classical SR1 quasi-Newton update x_{k+1} = x_k − α_k B_k g_k with three small element-wise MLPs and a limited-memory PSD preconditioner:

• Encoder ℰ: per-coordinate state (x, p, d, g, q) → latent (5→128).
• Vector generator 𝒫: latent → SR1 vector v_k (128→1).
• LR generator 𝒢: latent → α̃_k, with α_k = γ₁·exp(γ₂·α̃_k) (128→1).

The descent direction is d_k = Σ_{v∈B_L} v (vᵀ g) over a fixed-size buffer of the last L generated vectors (Algorithm 1). A learned PSD projection shapes the vectors during meta-training by penalizing the secant residual ‖p_k − B̃_k q_k‖² with B̃_k = I + Σ v vᵀ ⪰ 0 (Eq. 8), added to the meta-loss with weight λ_sec (Eq. 9).

Files

• lsr1.py — model (Tables 4–5), limited-memory SR1 rollout (Alg. 1), meta-loss with secant penalty (Eqs. 8–9), analytic problems (quadratic / Rosenbrock / Rastrigin) + Newton direction.
• baselines.py — L-BFGS, Adam, AdaHessian; Dolan–Moré performance profile.
• train_analytic.py — meta-training (presets quad2, quad100, rosen100, rastr100), projection ablation, per-experiment hyperparameters from Table 6.
• eval_analytic.py — Newton-direction alignment, performance profiles, convergence curves.
• efficiency.py — Claim 6 mechanism microbenchmark (L-SR1 update module vs an LGD-style full-state updater): runtime + memory + params.
• jobsrc/job_driver.py — the Hugging Face GPU Job driver (trains everything, evaluates, benchmarks; writes to an HF bucket).
• plot.py — figures + CSVs from the job artifacts.

Reproduction notes / deviations

• Preconditioner B̃ = I + Σ vvᵀ for the step (not just Σ vvᵀ). SR1 starts from B₀ = I, matching the secant penalty's B̃. This was essential: with the identity baseline the learned step reaches cosine ≈1.0 with the Newton direction on N=2 quadratics; without it the optimizer must learn the whole preconditioner from scratch and fails to align. L-SR1 aligns with the Newton direction; L-BFGS still converges faster on pure quadratics (the paper's quadratic-tailored reference) — L-SR1's edge on quadratics is alignment, and its headline is the performance profile over the mixed 30-problem suite.
• Small-init generator heads. The untrained SR1 preconditioner is near-zero so the initial K-step rollout is numerically stable; standard for learned optimizers.
• Trust region (per-coordinate step clamp) on the stiff families (Rosenbrock/Rastrigin). Not in the paper; a standard learned-optimizer stabilizer needed because d=(Σvvᵀ)g is linear in the gradient and blows up on stiff problems early in training. Loose enough not to bind near convergence.
• Relative secant residual (÷‖q‖²) on the stiff families only, so a single λ_sec works across gradient scales. The quadratic Claim-2 study uses the absolute residual exactly as written.
• Scale. Meta-training uses 6K (quadratics) / 4K (Rosenbrock, Rastrigin) iterations vs the paper's 10K — a scaled but non-toy run on one L40S GPU.

See the Trackio logbook for results, figures, and the Hugging Face artifact links.