98 lines
5.5 KiB
Markdown
98 lines
5.5 KiB
Markdown
# 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.*
|