SEoS: data for "Stochastic Edge of Stability"

Training logs, EoS/EoC analysis outputs and figure inputs for the paper. No model weights are stored here. On 2026-09-29 all checkpoints (checkpoints/step_N.pth) and all Lanczos eigenvector files (eos_eig/ckpt-K/eig_step*.pt) were deleted. The eigenvalues are still in the JSON/JSONL files next to them. Checkpoints of the new-code (eos-mup) runs remain on GCS (gs://tml-data-personal-us-south1/willcai/enuma-elis/ckpt/eos-mup/<run>/).

Folders

folder	code	content	used in the paper
WSD-cases/	old (Stochastic-EoS)	GPT-extra-small, 8k vocab, WSD, peak LR 0.0005 and 0.032	Figure 8 (a), (c)
LR-scheduler/	old	GPT-extra-small WSD baseline, peak LR 0.004 (the WSD baseline of the LR sweep, stored in this folder)	Figure 8 (b), Figure 17 (training summary)
decay-length/, decay-ratio/	old	GPT-extra-small linear-decay ablations (peak 0.004, decay from step 2000)	Figure 12
eosBSZ-power-rule/, eosBSZ-sqrt-rule/, criticalBSZ-8k/	old	batch-size sweeps on arxiv / c4 / dclm / openwebtext / wikipedia (power-rule) and FineWeb (sqrt-rule), CBZ runs	supporting data; Figure 4 is built from other runs (see below)
scale20TPP8k/	new (eos-mup)	GPT-small/medium/large, 8k vocab, 20 tokens per parameter, 8 peak LRs, 10 stable-phase checkpoints each	Figures 5, 14–16
TPP20Plots/	new	dense-checkpoint runs at peak LR 1.25e-3 (small, medium, large)	Figure 2
fig1-tiny/	new	GPT-extra-small (tiny) WSD 0.002 run used for Figure 1 drafts (uploaded by t-eos2)	the final Figure 1 data are in eosmup-runs/eos-fig1b-*
eosmup-runs/	new	all new-code runs and analyses of T-EoS without weights: GPT-extra-small (tiny) batch-size runs (Figure 9), multistep and full-cosine reruns (Figures 3, 11), LR 1e-4 baseline, Figure 1 insets (eos-fig1b-eoc: λ_P = 0 simulations, eos-fig1b-true: true AdamW loss on the same batches), 50k-vocab CompleteP/muP runs, large-model EoS passes, code snapshots (source-*)	Figures 1, 3, 9, 11
paper-figures/	new	figure builders and their inputs (TPP20 panels, legend variants, heatmap redraws, Figure 1 script)	Figures 1–3, 8–16
eosmup-data-meta/	new	small metadata of the tokenized data (sidecars, receipts, tokenizer); the shards themselves are on GCS	—
archive-2026-09/	older studies	reproduction (162M), NTK (43M), AdamW-prox and paired studies from 2026-09-14 to 09-20, logs and results only	not in the paper

Definitions used by the paper

LR factor = simulated LR / training LR. The critical LR factor q is the smallest LR factor in the grid whose compensated local simulation explodes (displacement norm reaches 1e5); a checkpoint is at EoS when 0.9 ≤ q ≤ 1.1. Some summaries in scale20TPP8k/ and TPP20Plots/ report the largest bounded factor with a 1e3 cutoff instead (round-1 views); recompute q from the trial logs for the paper's numbers.

Known issues (not fixed in the data)

• Compensation fallback (old code). When the λ_P search returned 0, eos_al ran the trials at the top of the search bracket instead: decay-length wsd-8000 (ckpts 5000–6000) and wsd-10000 (ckpts 7200–7600) at 9500, and one or two columns of the Figure 8 runs. In these columns the λ panel shows 0 while the heatmap used the bracket top.
• Saturated searches. WSD-cases lr 0.032: the λ_P search did not find a window at 20 of 49 checkpoints (λ_P pinned at 312.5 = 10/η).
• Last checkpoint not analysed in the old-code groups (49 of 50 columns, 9 of 10 in eosBSZ-*).
• eosBSZ- schedules*: the names say wsd-T-(T/10)-(2T)-(2.5T), but training stops at T, so these are warmup + constant-LR runs. Both groups use the same batch-size rule; they differ in dataset.
• decay-length names: wsd-----; the 4th number is where the decay ends (decay length 3000 = wsd-8000-…-5000, 5000 = wsd-10000-…-7000, 8000 = wsd-13000-…-10000).
• Warmup of the 20-TPP runs (scale20TPP8k/, TPP20Plots/) is 6.2% of the steps (247/3990, 407/6563, 621/10019).
• Figure 4 (EBZ vs CBZ) is built from coauthor runs (run_tiny_*, OpenWebMath / C4 / DCLM / FineWeb / OpenWebText, batch sizes 64–1024) that are not in this repo.

Removed on 2026-09-29

drop-PeakLR/ and drop-ratio/ (GPT-small, 50k vocab, unused, eos_tl ignored the LR factor), BSZ-cases/ (superseded fixed-LR batch-size runs), the LR-scheduler runs …_cos-10000-400-8000-10000-0_260117-073412 (WSD with a cosine tail, unused) and …_multis-10000-400-3000-10000-0.7-0_… (unused, broken analysis), and 12 duplicate eosBSZ-power-rule/*_260125-093236 runs (twins of *_260125-095909). Git history keeps them.