SGD vs AGPT sanity check

SGD vs AGPT Sanity Check

Status: concluded, inconclusive (reviewed 2026-09-28; this line previously read "incomplete — redo required."). The answer is in the front matter above.

Trainer note: pre-fix, needs reassessment.

Code: branch sgd-sanity-check (intended to merge to main once validated). Adds --mass-weight <mode> where mode ∈ {off, log, sqrt, linear}, replacing the old boolean --mass-weight flag. Linear mode matches SGD's frequency weighting — needed to isolate whether the weighting choice (not the aggregation) is what differentiates AGPT from SGD.

Current status

This experiment's purpose is still valid, but the restored runs in this directory were taken on the pre-fix AGPT trainer and should be treated as provisional notes, not current evidence.

Why:

• The AGPT side predates commit 1c858c0 (2026-04-23), which fixed the frozen Wk / Wv / bias bug in bin/agpt_train.
• The old SGD side also hit a separate cuBLAS / NaN issue.

This is why the directory was removed from main in commit 15323ef ("Remove pre-fix sgd-sanity-check experiment") and later restored only as a cleanup/reassessment reference.

For now, treat the restored logs here as branch artifacts to learn from while planning the redo. Do not cite their quantitative results without re-running the experiment on the post-fix trainer / stable SGD baseline.

Hypothesis

If the D-trie captures the corpus's conditional distributions faithfully, and AGPT's gradient aggregation is mathematically equivalent to "sum the gradient contributions from all corpus positions sharing the same context," then AGPT and standard SGD over corpus positions should converge to similar held-out PPL at matched compute budget.

If AGPT is significantly better at matched compute: the trie aggregation is adding something (prefix-sharing speedup, or the equal-per-context weighting is actually beneficial).

If SGD is better: either there's a bug in AGPT's gradient aggregation OR the equal-per-context weighting is actively hurting us and we should use --mass-weight.

If they match: AGPT is behaving as specified — the aggregation is mathematically sound. This is the primary sanity check.

The weighting question (why this isn't just a re-run)

SGD with uniform corpus-position sampling trains in proportion to corpus frequency — common D-grams get sampled often, rare ones rarely. Each gradient step is weighted implicitly by f(context).

AGPT's default behavior is equal weight per radix node, regardless of how many corpus positions that node represents. So:

• AGPT-default: sums ∇L across unique contexts with equal weight each.
• AGPT --mass-weight: weights by log(1 + edge_mass) — partial frequency restoration.
• SGD: weights proportional to corpus frequency (linear in count).

These are THREE different effective loss functions, even though they're training the same model on the same corpus. Comparing them cleanly is part of this experiment.

What we'll measure

1. Raw held-out PPL at each config, matched at:
   • Same wall-clock time.
   • Same number of optimizer steps.
   • Same number of gradient samples (corpus positions covered).
2. Convergence speed (PPL vs. training time curve).
3. Sensitivity to weighting scheme (AGPT default vs. AGPT mass-weighted vs. SGD uniform).

Configs

All runs use the same model architecture:

• d_model=64, n_heads=4, n_layers=2, d_ff=256, head_dim=16, vocab_size=65.
• lr=3e-3, rmsprop (β=0.999), warmup-cosine schedule.

Variants to compare

Matched-compute scheme

AGPT d=16 at 3 super-epochs = 195 optimizer steps. Baseline. AGPT d=32 at 3 super-epochs = 195 optimizer steps. SGD will be run at various step budgets (195, 1000, 5000, 20000) to see where its PPL curve sits relative to AGPT's.

Files

• run.sh — runs the matrix of configs.
• logs/ — raw training output per config.
• results/summary.md — condensed PPL comparison.

Results from the pre-fix run set

The table below is preserved only as a record of what the old branch claimed. It is not yet a trusted current conclusion.

Key reads:

• Linear weighting (≡ SGD frequency weighting) wins or ties at both depths. The "rare sequences count too much" concern is real; SGD-style weighting is the right direction.
• At d=16, log and sqrt actively hurt. They're a bad middle ground — enough frequency bias to break AGPT's equal-per-context compensating structure, not enough to match SGD's actual distribution.
• At d=32, weighting mode barely matters (except linear edges ahead). Because most cap endpoints are singletons with count=1-2, log(1+count) ≈ sqrt(count) ≈ count in that regime.
• d=32 linear = 13.36 is the new overall best mean seen in this project (tied with d=16 blending within noise).

Historical implication only: the branch was leaning toward --mass-weight linear, but that conclusion needs to be re-earned on the correctly-trained model.

Decisions made (documentable so we don't re-litigate)

• Training data: data/input.txt (Shakespeare, 1.11M chars).
• Eval: bin/perplexity --file data/input.txt --max-positions 16384 --backend openblas. Uses the same file for eval — relies on perplexity tool's held-out-by-default convention. (Worth double-checking: is the eval really held-out?)
• Model init: fresh weights per config (not loaded from a common checkpoint) — otherwise we'd be measuring fine-tuning dynamics, not the training-from-scratch behavior we want.
• Single-subtree mode: for AGPT (matches the memory's winning recipe).

Redo plan

When this experiment is revived after cleanup, keep the original question but restart from scratch:

1. Use post-fix bin/agpt_train only.
2. Use a known-good SGD reference path (likely openblas, not the old cuBLAS run).
3. Keep the comparison focused on:
   • AGPT aggregation correctness
   • weighting choice (off, log, sqrt, linear)
4. Write new results into a fresh summary instead of mixing them with the invalidated logs restored from branch history.