10 KiB

Raw Permalink Blame History

How SR4 Works

A complete mechanistic walkthrough of the SR4 optimizer as implemented in this library (ternary_optim.c), adapted from the SR4 research prototype's design notes. Read ../README.md first for the one-paragraph version and the motivation.

Everything below describes the exact running code — constants and formulas are lifted from the research prototype and preserved bit-for-bit in ternary_optim_step() / ternary_optim_materialize().

1. The state: one 4-bit integer per weight

Each weight i has a single signed integer accumulator q_i:

q_i ∈ {-7, -6, ..., -1, 0, +1, ..., +6, +7}      (15 states → a 4-bit nibble)

That accumulator is the entire persistent per-weight training state. No FP32 master copy, no momentum, no Adam m/v. (Here q is stored one-per-byte for convenience; the information is 4 bits — 15 states pack into a nibble, which is the memory number that matters: 4 bits/weight.)

Two config constants govern it (ternary_optim_config):

q_clamp = 7     # symmetric saturation:  -7 ≤ q ≤ +7
theta   = 2     # dead-zone threshold (see §2)

Initialization (ternary_optim_create). q starts at {-3, 0, +3} with probabilities {33.5%, 33%, 33.5%} — ~67% of weights active, signs uniform, magnitude 3 = theta + 1 giving every active weight one LSB of dead-zone margin. The draw uses an internal xorshift seeded from (seed ^ idx_base) so distinct matrices in a model get distinct inits.

2. Forward: ternarize the accumulator (`ternary_optim_materialize`)

The deployed weight is a pure function of q:

w_i = 0            if |q_i| < theta        (i.e. q_i ∈ {-1, 0, +1})
      +1           if q_i ≥ +theta         (i.e. q_i ∈ {+2, ..., +7})
      -1           if q_i ≤ -theta         (i.e. q_i ∈ {-2, ..., -7})

So the deployed weights are ternary {-1, 0, +1}.

The theta = 2 dead zone is the key behavioral detail. A weight does not fire the instant its accumulator moves off zero — it must build conviction (|q| ≥ 2) before it turns on, and stays on until the accumulator decays back below 2. That hysteresis lets a low-resolution integer state produce stable ternary weights instead of chattering.

3. Update: stochastic-rounded integer integration (`ternary_optim_step`)

Once per batch, per weight matrix, q is nudged toward the negative gradient. This is the whole optimizer. Given grad_sum[i] (summed over the batch) and batch size bs:

3a. Per-matrix step normalization

g_i     = grad_sum[i] / bs                       # mean gradient for weight i
mean|g| = (1/N) · Σ_i |g_i|                       # mean gradient magnitude, this matrix
s       = target_step / (mean|g| + 1e-12)         # scalar, recomputed every batch

s rescales the update so the typical weight (|g_i| ≈ mean|g|) proposes a step of ~target_step LSB, regardless of the raw gradient scale. target_step is small (~1e-3), so per update q moves by a tiny fraction of an integer — it integrates gradient direction over many updates. s is a continuous scalar recomputed each batch; nothing is stored. (The training loop may pass an annealed target_step — e.g. cosine — as the argument; if it passes <= 0 the config default is used.)

3b. Propose a fixed-point step

Snapped to 8.8 fixed point (256 sub-LSB units per integer) — the single float→discrete boundary in the whole update:

prop     = -g_i · s · 256                         # negative gradient, in 1/256-LSB units
prop     = clamp(prop, -4096, +4096)              # defensive: ±16 LSB saturates q anyway
step_fx  = lrintf(prop)                            # int32, fixed-point step

3c. Stochastic rounding: split into integer + fractional parts

k = step_fx >> 8        # whole LSBs (arithmetic shift = floor, correct for negatives)
f = step_fx & 255       # fractional part, in [0, 255]  (i.e. f/256)

The integer part k is applied deterministically. The fractional part f/256 is applied stochastically as an extra +1 with probability f/256:

sr_bit = 1  if  dither_u < f  else  0            # dither_u ∈ [0, 255], uniform

so the expected increment is exactly k + f/256 = step_fx/256 — the update is unbiased. Stochastic rounding is what recovers sub-integer resolution from a 4-bit state.

3d. Saturating add

q_i = clamp(q_i + k + sr_bit, -q_clamp, +q_clamp)

The ±7 rails are legitimate fully-committed states (never wraparound). That's the entire update: a saturating integer add whose fractional part is stochastically rounded.

4. The dither: where `dither_u` comes from (stateless)

dither_u is not a stored RNG stream. It's a pure function of position from a counter-based hash (SplitMix64, sr_hash in the source):

sr_hash(seed, counter, global_idx):
    x = mix(seed  ^ C0)
    x = mix(x     ^ counter)
    x = mix(x     ^ global_idx)
    return x

Stateless / nothing stored — regenerated from (seed, counter, weight-index) on demand; zero persistent memory.
Parallelizable — every weight computes its own draw independently.
Reproducible & flag-immune — pure integer arithmetic, bit-identical across -O levels and -ffast-math.

global_idx is idx_base + i, so giving each matrix a distinct idx_base decorrelates their draws.

Low-discrepancy (variance-reduced) rounding — `use_dither = 1` (default)

Plain i.i.d. stochastic rounding is unbiased but noisy — the realized fire-rate random-walks around f/256. SR4 uses a stratified Weyl (low-discrepancy) dither, still unbiased but far lower variance:

stratum  = counter >> 8                 # re-key the per-weight phase every 256 batches
j        = counter & 255                # position within the stratum
h_i      = sr_hash(seed, stratum, idx_base+i) & 255   # per-weight uniform phase
dither_u = (h_i + 159 · j) mod 256      # near-golden (256/φ) Weyl rotation

Because 159 is odd, j → (159·j mod 256) is a bijection on {0..255}: over a full stratum a constant fraction f fires exactly f times, O(1) discrepancy instead of random-walk √-noise. Same expectation, tighter variance — which directly tightens the seed-to-seed spread of trained models. Setting use_dither = 0 falls back to plain i.i.d. dither_u = h_i (sr_hash keyed on the raw counter).

5. Co-trained per-column output scales (optional adjunct)

Ternary weights {-1,0,+1} throw away magnitude. To give the network back its dynamic range, each layer carries a small vector of FP32 per-output-column scales, folded into the forward pass:

Z = scale ⊙ Z_raw + bias        # Z_raw = ternary-weight matmul output

The scales train by plain SGD (at a dedicated small SCALE_LR, ~10× below the bias LR) alongside q. Memory cost is tiny — ~266 FP32 scalars over ~203k weights ≈ 0.042 b/w — so the total still rounds to 4.0 b/w. This library provides the two decoupled primitives (ternary_apply_col_scale, ternary_col_scale_grad) generalized to any (rows, cols); the scale vector itself is owned by the caller (the layer), NOT the optimizer. The optimizer works without it, just with less dynamic range.

6. Putting one training step together

for each batch:
    materialize ternary weights:  w = ternarize(q)          # §2
    forward:   Z_raw = w · x ;   Z = scale ⊙ Z_raw + bias   # §5
    loss + backward → grad_sum (per weight), scale/bias grads
    ternary_optim_step(opt, grad_sum, bs, counter, target_step)  # §3–§4
    scale/bias ← SGD step
    counter += 1

The deployed model at any time is just ternarize(q) plus the scales.

7. Memory accounting (the whole point)

Component	Per-weight cost	Notes
`q` accumulator	4 bits	15 states, nibble-packable — the entire integer state
per-column scales	~0.042 b/w	~266 FP32 over ~203k weights
biases	negligible	per-output-unit
Total persistent training state	≈ 4.0 b/w
FP32 master weight	0	— none —
momentum / Adam m,v	0	— none —
SR dither / dropout masks	0	regenerated from the counter-hash on demand

At 3.6B parameters: 4.0 b/w × 3.6e9 / 8 = ~1.8 GB of total training state.

8. Why it's stochastic at every step

Unlike a deterministic optimizer — where the seed only touches init and shuffle order — SR4's stochastic rounding threads the seed through every weight on every update via the counter-hash. Consequences:

There is a real seed band (spread of final accuracy across seeds); results are reported as 3+-seed bands, never single runs.
Variance-reducers compound. Bigger batch, low-discrepancy dither, and dropout each shrink the band.
The stochastic rounding is how the 4-bit state works at all — not incidental noise, it's the resolution mechanism.

9. Design principles carried from the research

Training-state precision ≠ deployed-weight precision.
q is a conviction integrator. Slow integer integration + a dead-zone gives stable ternary weights from a low-resolution state.
Stochastic rounding = the missing bits. Low-discrepancy dithering makes it low-variance.
Stateless-by-construction randomness. Counter-hash, not stored RNG.
Think in updates, not epochs. The invariant is the number of weight updates (the batch_counter), not passes over the data.
Regime matters. A lever's value depends on its base; tune in the target regime.

10. Open questions for the library / scale work

Token-efficiency vs Adam at scale — SR4 is a slow integrator; does it need more updates/tokens than Adam to reach the same loss?
Generalization of the scales & dither to arbitrary weight-tensor shapes.
State serialization & optimizer-state reset — the packed q is the thing to persist (currently stubbed here).
Where the recipe lives once modularized: the integer update → the optimizer; batch/label-smoothing/dropout → loss + loop; scales → the layer.

10 KiB Raw Permalink Blame History Unescape Escape