YMRohit/icml17708-grok-grokking-repro

To Grok Grokking — independent reproduction

This package independently reproduces and audits the five requested claims of To Grok Grokking: Provable Grokking in Ridge Regression by Mingyue Xu, Gal Vardi, and Itay Safran (ICML 2026 paper 17708; OpenReview 5nNNVY8NW4; arXiv 2601.19791, v3). Reproduction date: 2026-07-16.

Outcome first

The ridge mechanism and its hyperparameter predictions reproduce strongly, but the package does not support every claim without qualification. Theorem 4.1 states end-to-end bounds for all three zero-teacher phases; a numerical mechanism check gives t1=210, t2=22,469, and a 22,259-step delay. Theorem 4.2 and Theorems 4.4–4.6 have the advertised scope and decomposition, but their printed optimization conditions need repair: a two-eigenvalue empirical operator diverges under a step size that satisfies the stated bound. At the paper's Figure 2 dimensions, 200 exact trajectory evaluations reproduce the predicted t2 ∝ lambda^-1 relation and the sample-size direction. In the nonlinear extension, the selected Figure 4 sweeps reproduce on an L4, while the stated independent nonzero-teacher Figure 3 protocol never reaches the 0.01 test-MSE threshold and is therefore not reproduced under the documented assumptions.

The machine-readable version of this table is `evidence_manifest.csv`. Explicit permitted and forbidden inferences are in `data/audit_negative_cases.json`.

Primary-source and code-availability audit

• —Paper: https://arxiv.org/abs/2601.19791
• —OpenReview: https://openreview.net/forum?id=5nNNVY8NW4
• —Archived v3 PDF, text, source tarball, and OpenReview metadata: `references/`
• —Unpacked TeX used for line-level theorem and experiment audits: `ridge_agent/paper_source/`

No official experiment code or raw experiment data were linked from the OpenReview submission or included in the arXiv v3 source archive, and no author repository was found in exact-title/source searches as of 2026-07-16. The implementations in this package were reconstructed independently from the paper's equations, prose, captions, plot axes, and released plot PNGs. They must not be described as the authors' code.

Main findings

Theorem audit: a concrete step-size gap

Theorems 4.2 and 4.4 print eta < 1 / (lambda + lambda_min_plus(Phi^T Phi / n)). That controls the smallest positive eigenvalue, not the largest eigenvalue governing quadratic-GD stability. `scripts/audit_theorems.py` tests a concrete two-eigenvalue empirical operator with eigenvalues 0.99 and 0.01, takes lambda=0.001, and tests eta=3:

• —the printed upper bound is 90.909, so eta=3 satisfies it;
• —the update's actual spectral radius is 1.973;
• —the training loss grows 11,986,302.7× in 12 steps;
• —the standard largest-eigenvalue stability bound is eta < 2.0182;
• —a stable control with eta=0.5 reduces training loss to 0.00887× its initial value over the same 12 steps.

Theorem 4.6 also says GD converges to the ridge minimizer without stating a step-size restriction. The bounded-feature, sample-complexity, and statistical parts were not invalidated by this audit; the finding isolates the optimization condition. It also does not imply that the paper's Gaussian simulations with eta=1 diverge—the reproduced trajectories are stable.

Figure 2: paper-scale ridge result

The ridge solver evaluates discrete full-batch GD exactly in the eigensystem; it is mathematically the same integer-step trajectory, not a shortened-training proxy. The T4 campaign used the paper dimensions and selected plot values: m=1000, default n=100, 20 deterministic trials, five decays 2^-19,…,2^-15, and n∈{10,20,40,80,160}. This yields 200 trajectory rows (100 decay + 100 sample-size evaluations) from 120 sampled/eigendecomposed problems. Mean grokking time rises from 74,981 to 1,202,915 across the decay range. For sample size, it is 23,478.45 at n=10 versus 22,237.55 at n=160; the main effect is the predicted rise in training time (t1 mean 17.55 to 355.10). All 200 rows satisfy both literal natural-log Equation (8) bounds.

Figures 3 and 4: corrected post-hoc interpretation

The selected Figure 3 campaign uses the stated random-feature dimensions d=100, m=10,000, the paper initialization and independent unit ReLU teacher, Gaussian inputs, an 8,192-example Gaussian test set, three seeds, and 18 runs spanning three decay and three sample-size values. Every run reaches the training threshold; all 18 generalization crossings are censored. In the decay sweep, final test MSE medians are 0.20670, 0.20705, and 0.20914, far above 0.01.

The selected Figure 4 campaign uses explicit simultaneous full-batch GD for both layers at d=50, m=1000, a zero teacher, a 4,096-example test set, two seeds, and 18 runs. All t2 values are observed at 250-step resolution and both tested directions reproduce.

The raw GPU artifacts contain a pre-fix claim5_verdict.json whose all_tested_trends_supported field counted only assessable trends and could therefore be true while Figure 3 was entirely censored. The frozen CSVs are the source of truth; `outputs/jobs/claim5_posthoc_verdict.json` is the corrected campaign-level verdict and supersedes that boolean.

A matched, one-seed zero-target sensitivity holds all Figure 3 dimensions, test size, initialization, thresholds, and decay values fixed while changing only the teacher target. It gives t2={14,694; 146,943; 1,469,430} for lambda={1e-4; 1e-5; 1e-6}, with t1={664; 703; 707}. This clean 1/lambda response shows that the implementation can express the predicted mechanism and that teacher realizability is material. It is diagnostic only: Figure 3 states an independent nonzero ReLU teacher, so the sensitivity does not establish which hidden protocol the authors used and is not counted as a Figure 3 reproduction.

Compute, hardware, scale, and cost

Estimated GPU total: $0.1885. GPU durations are the difference between Job creation and Trackio artifact creation and therefore include environment setup and artifact upload. The Figure 2 Job reports Linux 6.12, Python 3.12.12, PyTorch 2.13.0+cu130, and a Tesla T4. Local verified smoke output records WSL2, Python 3.13.11, and PyTorch 2.10.0+cu130. The L4 runs resolved torch and trackio through uv --with; the exact executed source/config hashes and raw results are preserved, but this campaign did not freeze a package lockfile.

Exact commands

Install the two small dependency sets from the repository root:

python -m pip install -r repro_17708/ridge_agent/requirements.txt
python -m pip install -r repro_17708/relu_agent/requirements.txt

Run the local checks and deterministic report build:

python repro_17708/scripts/verify_zero_teacher.py
python repro_17708/scripts/audit_theorems.py
python repro_17708/ridge_agent/test_closed_form.py
python repro_17708/ridge_agent/reproduce_figure2.py \
  --trials 2 --device cpu \
  --output-dir repro_17708/outputs/ridge_smoke_verified \
  --trackio-run-name ridge-local-smoke-verified
python repro_17708/relu_agent/relu_grokking.py \
  --config repro_17708/relu_agent/configs/smoke.json \
  --output-dir repro_17708/outputs/relu_smoke_verified \
  --device cpu --trackio-project icml17708-relu
python repro_17708/relu_agent/relu_grokking.py \
  --config repro_17708/relu_agent/configs/fig3_zero_matched.json \
  --output-dir repro_17708/outputs/fig3_zero_matched \
  --device cpu --trackio-project icml17708-relu-zero-sensitivity
python repro_17708/scripts/build_report_assets.py

After the folder is finalized, log the complete publication bundle (code, configs, raw outputs, manifests, figures, and primary-source archive) as one Trackio artifact:

HF_HOME=/home/tihor/ICML/.hf TRACKIO_LOGBOOK_AUTONOTE=0 \
  python repro_17708/scripts/log_reproduction_bundle.py \
  --project icml17708-reproduction-package
HF_HOME=/home/tihor/ICML/.hf \
HF_TOKEN_PATH=/home/tihor/.cache/huggingface/token \
  trackio sync --project icml17708-reproduction-package \
  --space-id YMRohit/icml17708-grokking-metrics --force

These are the three GPU commands captured by the logbook:

hf jobs uv run --flavor t4-small --timeout 30m --detach \
  --secrets HF_TOKEN \
  --env TRACKIO_SPACE_ID=YMRohit/icml17708-grokking-metrics \
  --label paper=5nNNVY8NW4 --label claim=4 \
  --with torch --with trackio \
  https://huggingface.co/datasets/YMRohit/icml17708-grok-grokking-repro/resolve/main/ridge_agent/reproduce_figure2.py \
  --trials 20 --device cuda --dtype float32 \
  --output-dir outputs/claim4-gpu \
  --trackio-project icml17708-ridge \
  --trackio-run-name ridge-figure2-t4-20seed

hf jobs uv run --flavor l4x1 --timeout 60m --detach \
  --secrets HF_TOKEN \
  --env TRACKIO_SPACE_ID=YMRohit/icml17708-grokking-metrics \
  --label paper=5nNNVY8NW4 --label claim=5-fig3 \
  --with torch --with trackio \
  https://huggingface.co/datasets/YMRohit/icml17708-grok-grokking-repro/resolve/main/relu_agent/relu_grokking.py \
  --config https://huggingface.co/datasets/YMRohit/icml17708-grok-grokking-repro/resolve/main/relu_agent/configs/job_fig3_selected.json \
  --output-dir outputs/claim5-fig3-gpu --device cuda \
  --trackio-project icml17708-relu

hf jobs uv run --flavor l4x1 --timeout 2h --detach \
  --secrets HF_TOKEN \
  --env TRACKIO_SPACE_ID=YMRohit/icml17708-grokking-metrics \
  --label paper=5nNNVY8NW4 --label claim=5-fig4 \
  --with torch --with trackio \
  https://huggingface.co/datasets/YMRohit/icml17708-grok-grokking-repro/resolve/main/relu_agent/relu_grokking.py \
  --config https://huggingface.co/datasets/YMRohit/icml17708-grok-grokking-repro/resolve/main/relu_agent/configs/job_fig4_selected.json \
  --output-dir outputs/claim5-fig4-gpu --device cuda \
  --trackio-project icml17708-relu

The commands above used /resolve/main/, but all three were launched after dataset commit `8b373bfa6dcd08342989fb1cfa3e2a7bcba3fa37` and before any later dataset commit. The exact files downloaded from that immutable revision are under `job_snapshots/revision_8b373bfa/`. For an immutable rerun, replace resolve/main in each URL with resolve/8b373bfa6dcd08342989fb1cfa3e2a7bcba3fa37.

Artifact and provenance map

Each GPU output directory was downloaded from its Trackio artifact. Artifact version names and manifest digests are recorded in outputs/jobs/jobs.json. The working ReLU script was subsequently patched so censored trends cannot be counted as successes; the unmodified Job source remains in the snapshot, and the corrected interpretation is derived only from the frozen results.csv files.

Public resources

Scope and limitations

1. 1.The theory assessment is a direct source audit plus executable sanity and counterexample checks; it is not a machine-checked formal proof.
2. 2.No author code or raw data were available. The ridge setup is well specified, but the ReLU section omits its input distribution, test-set size, random seeds/repeats, and threshold-interpolation rule. We declare Gaussian inputs, held-out sizes, seeds, and evaluation resolution in every config.
3. 3.Figure 3 and Figure 4 retain the paper dimensions for the tested lambda and sample-size panels but do not rerun every width/initialization panel. Three and two seeds, respectively, are directional evidence rather than exhaustive statistical replication.
4. 4.Random-feature and ridge trajectories are evaluated analytically at exact integer GD steps after eigendecomposition. Figure 4 uses explicit updates; its test crossing is resolved only to 250 steps.
5. 5.The Figure 3 negative result is scoped to the stated independent nonzero teacher plus our documented missing-protocol choices. It does not establish what code or hidden settings generated the authors' plot.
6. 6.The Dataset card's MIT license covers the reproduction code and original package materials. Archived paper/source assets remain the authors' work and retain their original terms.
7. 7.Two independent local review passes audited the claim evidence and poster; their findings were resolved in the final package. The requested Claude Code Fable 5 review was not run because explicit approval to transmit the evidence to that external service was not received.