Foundry/ternary_optim/docs/OPERATING_GUIDE.md

5.5 KiB
Raw Permalink Blame History

SR4 Operating Guide — how it behaves, and how to train with it well

This is the operational knowledge about the SR4 optimizer — distilled from the SR4 MNIST research campaign (reached the FP32 ceiling, memory-neutral) and the first RWKV-x LM runs. HOW_IT_WORKS.md covers the mechanism; this covers how to actually get good results with it and how to read its behavior. If you are training with SR4, read this first.

The one root fact everything follows from

SR4 is stochastic on every update. The int4 accumulator q is integrated by stochastic rounding — so the per-update gradient noise is amplified into the weights, every step, every weight. A deterministic optimizer (SGD/Adam) only touches randomness at init + data order; SR4 threads it through all ~10⁵ updates. Almost every operating rule below is a consequence of this.

Rule 1 — Batch size is the DOMINANT lever. Never run tiny batches.

Bigger batch → lower per-update gradient variance → the stochastic rounding integrates a cleaner signal. This is the single most important knob.

  • On MNIST, going batch 128 → 2048 was the path from "SR4 works" to "SR4 matches full-precision-dense." Every variance-reducer we found was subsumed by big batch at the top.
  • batch=1 is pathological — maximum per-update noise, and SR4 will plateau while Adam (batch-insensitive due to per-coordinate adaptivity) descends right past it. If you see that pattern, the fix is batch, not learning rate.
  • Diagnostic signature of "SR4 starved of batch": noisy/spiky train loss + an early plateau, while a reference optimizer keeps improving.

Rule 2 — SR4 has its own learning rate, and it is NOT Adam's

The step size is TARGET_STEP (the mean proposed step in LSBs per update). It must be tuned per optimizer and per scale — do a small sweep at each new rung. Observed optima: MNIST-MLP 0.001, LM memorization 0.003, LM real-data ~0.001. Never share an LR with Adam in an A/B — that rigs the comparison. Compare each optimizer at its own best LR, at matched tokens, and report wall-clock separately.

Rule 3 — Report seed BANDS, never single runs

Because SR4 is stochastic every step, final quality varies by seed. Run ≥3 seeds and compare band means with a noise margin (MNIST band ≈ 0.25pp; a gain must clear the band to count). A single lucky/unlucky seed is not a result.

Rule 4 — Variance-reducers compound (super-additively)

Anything that lowers the noise floor lets small mean-gains become visible and stackable:

  • Low-discrepancy (Weyl) dither — built into the core; variance-reduced unbiased stochastic rounding (free, always on).
  • Batch — the big one (Rule 1).
  • Dropout — a tunable reliability dial: more dropout → tighter seed band (its cost is mean, so anneal/gentle-decay it; regime-dependent — see Rule 6). Small sub-band mechanisms that look like noise individually start stacking once the floor drops.

Rule 5 — The recipe (core vs knobs)

The core is small: int4 accumulator + stochastic-rounded integer update + ternary forward. Everything else is recipe, and it lives in different places:

  • Batch size → training loop (Rule 1, dominant).
  • Co-trained per-column FP32 scales → per layer; restore the dynamic range ternary weights throw away (Z = scale ⊙ (ternaryW·x)). ~0.04 b/w. Basically required for quality.
  • Label smoothing → loss side (biggest single regularizer on MNIST).
  • Dropout / annealed dropout → the reliability dial (Rule 4). The core got MNIST to ~97; the recipe got it to the FP32 ceiling (~98.3). Don't judge SR4 without its recipe.

Rule 6 — Regime matters

A lever's value depends on its base. Examples that actually flipped: early-release dropout was the worst schedule at batch 1024 but the best at batch 2048; a higher step helped at moderate batch and was null at large batch. Tune levers in the target regime; do not assume a lever that helped at one scale/batch transfers.

Rule 7 — Which params are ternary (the split)

SR4 ternarizes the large weight matrices (the projections — the bulk of params, the memory win). Keep small/precision-sensitive params continuous (float, plain optimizer): embeddings, the LM head, LayerNorm γ/β, LoRA/low-rank gates, biases, and the per-column scales themselves. This mirrors the BitNet/ternary-LLM convention. Each ternary matrix gets its own idx_base so the SR dither lanes stay independent.

Rule 8 — Think in updates/tokens, not epochs

The meaningful invariant is the number of weight updates (and, at LM scale, tokens) — not passes over the data. When you change batch size, hold updates constant (scale steps accordingly), or you are changing two things at once.

Quick reference — reading an SR4 run

You see It usually means Do
Spiky train loss + early plateau; Adam descends past it per-update noise too high raise batch (Rule 1)
Whole seed band wide noise floor high add a variance-reducer (batch / dither / dropout)
Loss won't drop at all LR wrong or scales missing check TARGET_STEP sweep + that per-column scales are on
Matches Adam only after the recipe expected that's the point — SR4 core + recipe ≈ full precision, at ~24× less optimizer memory

Provenance: the SR4 MNIST research campaign and the RWKV-x SR4-vs-AdamW scale test (both unpublished). Update this as behavior at LM scale is learned.