Lightning training

Lightning Training — empirical

Trainer note: pre-fix, needs reassessment.

Most quantitative claims in this file were produced before the Wk / Wv / bias gradient fix in commit 1c858c0 (2026-04-23). The experiment ideas and failure modes remain useful, but the reported PPL rankings should be treated as provisional until rechecked on the corrected trainer.

Design doc: notes/trainer/lightning-training.md.

Hypothesis

Stochastic variable-depth subtree sampling (L3 mass-weighted walk) can match or beat the deterministic per-root-child sweep at matched total optimizer steps. L3's stop-probability p_stop controls the expected sample depth: small p_stop → deep/narrow samples, large p_stop → shallow/broad samples.

Reference baselines

d=16 global-radix won't fit KV cache (9.5GB). Initial sweeps live at d=8 where one full run takes ~5 sec. d=16 work will require adopting the per-subtree file format (deferred — each subtree view is self-contained so L3 within a subtree still works).

Experiments

Result: matched-budget p_stop sweep (commit 20e8121)

All Lightning configs lose to the deterministic baseline by ~2.1 PPL at matched step count. Consistent with the design doc's hypothesis: stochastic sampling at matched budget without LR retuning underperforms.

Observations:

• Variance anti-correlates with p_stop. Smaller p_stop = deeper, lower-mass samples → flakier stochastic corpus coverage.
• All runs train (no divergence). Infrastructure verified.
• Best single run (p_stop=0.3, seed=43): 18.40 PPL — within 0.4 of baseline. Suggests LR tuning or a higher step budget could close the gap.

Result: extended step-budget + L1/L2 comparators (same commit)

Key findings:

• L2 is the closest Lightning variant to baseline (19.54 vs 17.99). L2 samples root-children uniformly with replacement — essentially the deterministic sweep with stochastic noise. Expected.
• L1 is disastrous (28.55). Uniform sampling across 845k radix nodes skews heavily toward deep, low-mass leaves.
• More L3 steps at the same LR makes PPL WORSE, not better. 2× steps → 25.83 (from 20.09 at matched), 4× → 26.13. At fixed lr=3e-3 the extra stochastic steps are over-rotating the weights. LR has to come down.
• L3 was design-motivated by mass weighting; at this LR it lands between L1 (bad) and L2 (close), but can't beat baseline.

Next direction: LR sweep for L3 at higher step budgets — the "many small steps need a smaller lr" story from earlier bigram work predicts lr≈3e-4..1e-3 at 260 steps/SE. If L3 can't close the gap even with tuned LR, the claim "mass-weighted stochastic sampling beats deterministic uniform-over-root-children" is probably wrong.

Result: per-sample mass-lr scaling (new --lightning-mass-lr flag)

Scales each Lightning sample's step_lr by compress(subtree_mass)/mean so high-mass samples move weights proportionally more. This is a LR scale, not a gradient scale — necessary because under RMSProp/Adam the weight cancels out in the adaptive divisor, so gradient scaling would be a no-op at steady state.

L3 p_stop=0.3, 65×3 matched budget:

Key finding: mass-off wins; two runs hit 17.69 / 17.74 — matching baseline within GPU-reduction noise. Mass-lr scaling, even with log compression, hurts. Read: RMSProp's second-moment accumulator already handles the "low-mass = noisy gradient" problem; an external LR multiplier just perturbs the running averages without adding signal. Linear and sqrt blow up as predicted (one huge-mass sample dominates with ratio ≥ 100×).

So #4 from the design discussion (mass-weighted step) is theoretically sound but empirically unhelpful on top of RMSProp. Flag is kept — may matter for SGD where gradient scaling and LR scaling ARE equivalent and there's no adaptive divisor to absorb the imbalance.

Result: step-budget × LR sweep (beats baseline on individual runs)

L3 p_stop=0.3, mass-off, 3 SE:

Two breakthroughs:

1. s=650 lr=1e-4 seed=43 hits 17.94 PPL — beats baseline. First individual Lightning run to cross that line.
2. s=260 lr=3e-4 has 0.27 PPL spread across 3 seeds — Lightning's stochastic-variance problem largely goes away at the right LR.

Rule of thumb from this sweep:

• Baseline budget (65 samples/SE): lr=3e-3
• 4× budget (260 samples/SE): lr=3e-4 (10× smaller)
• 10× budget (650 samples/SE): lr=1e-4 (30× smaller)

Stronger-than-sqrt LR reduction — consistent with "many stochastic steps need a smaller lr to absorb noise." RMSProp β=0.999 alone doesn't absorb the full variance at matched step counts; LR compensates.

Next: push to s=1300 or higher with matching LR cuts, or run enough seeds at s=260 lr=3e-4 to statistically verify the gap to baseline.

Result: d=16 Lightning via per-subtree files

Wired Lightning into run_per_subtree_training so d=16 becomes tractable: each root-child subtree file is loaded with its own bounded KV cache (~150 MB - 1.3 GB), samples are bucketed across root-children weighted by total_edge_chars. Peak memory use is the single largest root-child's KV, not the 9.5 GB global cache that would need to fit for direct global-radix Lightning at d=16.

L3 p_stop=0.3 mass-off, 3 SE:

Same pattern as d=8: individual best runs reach baseline (14.60 matches within noise, 14.72 is close), mean is ~1-1.5 PPL worse. s=260 lr=3e-4 again has tightest variance (0.64). s=650 lr=1e-4 has best mean AND highest ceiling individual run. The step-budget×LR trade mirrors d=8 — further LR sweeping at d=16 s=260 might close more of the gap.

Result: finer LR sweep at d=16 s=260

lr=2e-4 is the d=16 sweet spot: gap to baseline narrows from ~1.5 (at lr=3e-4) down to ~0.7. Classic U-shape — lr=5e-4 too aggressive (noise overwhelms signal), lr=5e-5 too slow (hasn't converged). Best run 14.78 is within 0.2 of baseline.

Result: BF16 K/V cache (halves memory)

KV cache stored as bf16 instead of fp32, with conversion at scatter/gather boundaries. Packed attention buffers remain fp32.

• d=8 Lightning: 791 MB (was 1582 MB). PPL 16.97 on smoke test — within noise of prior fp32 runs (mean 18.71, min 17.69).
• d=16 GLOBAL radix: 4.77 GB (was 9.5 GB, was REFUSED before). First time global-radix d=16 Lightning runs. Eliminates the per-subtree-file workaround — can go either way now.
• Next: d=32 global radix becomes borderline-feasible at 19 GB bf16 via unified memory paging (was 38 GB fp32).

Result: mass=1 compact cache (skips 83-96% of char positions)

Radix edges with edge_mass == 1 are leaves (a single corpus position). Nothing ever attends to their K/V as an ancestor, so they don't need cache storage. Own-edge K/V for the current query reads directly from the fresh d_k/d_v forward buffer (no round-trip through cache). Cache size drops dramatically.

d=32 is barely larger than d=16 because the extra depth is all mass=1 tails. Empirical confirmation of the theoretical claim that branching structure saturates around d≈16 for Shakespeare.

Result: d=32 global-radix becomes tractable — new project best PPL

With the compact cache, d=32 global-radix Lightning fits comfortably. 3 SE × 260 samples, lr=2e-4, L3 p_stop=0.3, mass-off:

d=32 Lightning beats:

• paper's d=32 baseline (13.17, per-subtree 3 SE × 195 steps)
• previously recorded project best (13.36, d=32 linear mass-weight)

Timing: d=32 ≈ 220 s/run (3 SE), d=16 ≈ 24 s/run. 7× slower but unlocks deeper branching supervision.

Run

rnd/lightning-training/run_d8_pstop_sweep.sh

Each cell dumps a config JSON (*.json), per-run logs (*_r{N}.log), and prints label run N PPL = X time = Ys to stdout.