260 lines
10 KiB
Markdown
260 lines
10 KiB
Markdown
# 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`](../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.
|