|
|
||
|---|---|---|
| .. | ||
| docs | ||
| Makefile | ||
| README.md | ||
| example.c | ||
| ternary_optim.c | ||
| ternary_optim.h | ||
README.md
ternary_optim — SR4 4-bit Stochastic-Rounded Latent Accumulator
A small, standalone C reference library for the SR4 ternary optimizer: a
from-scratch training optimizer whose entire persistent per-weight state is a
single 4-bit integer — no FP32 master weight, no momentum, no Adam moments —
which trains a network whose deployed weights are ternary ({-1, 0, +1}).
This library is the extraction step: the SR4 core lifted out of the SR4 research
prototype (an MNIST training harness) into a reusable, MNIST-decoupled API
generalized to arbitrary (rows x cols) weight matrices. The numerics
(stochastic rounding, low-discrepancy dither, per-matrix step normalization,
saturating add, counter-hash) are bit-for-bit faithful to the research
prototype.
ternary_optim.h/.c— the library.example.c— a self-contained demo (recover a ternary teacher via regression).docs/OPERATING_GUIDE.md— how SR4 behaves and how to train with it well (batch is the dominant lever, seed bands, the recipe, regime-matters). READ THIS BEFORE TRAINING.docs/HOW_IT_WORKS.md— the mechanism, step by step, with the math.docs/API.md— every public function: signature, semantics, ownership.docs/USAGE.md— how to wire it into a training loop.
Why it exists — the mission
Train large models on commodity hardware by making the training state almost free. Standard optimizers carry a lot of state per weight:
| Optimizer | Persistent training state / weight | @ 3.6B params |
|---|---|---|
| Adam / AdamW (mixed precision) | FP32 master (4B) + m (4B) + v (4B) = 12 B | ~43 GB |
| 8-bit Adam | master (4B) + m (1B) + v (1B) = 6 B | ~21 GB |
| SGD + momentum | master (4B) + momentum (4B) = 8 B | ~29 GB |
| SR4 | one int4 accumulator = 0.5 B (4 bits) | ~1.8 GB |
~24x less training-state memory than Adam, and the weights come out ternary (cheap inference too).
The core idea
Separate two things people usually conflate:
- Deployed-weight precision — how precise the weights you ship are: here
ternary (
{-1, 0, +1}, ~1.58 bits). - Training-state precision — how precise the state you must hold during
training is: here 4 bits — a tiny signed integer accumulator
qper weight.
Classic straight-through training keeps an FP32 shadow weight (32 bits of training state) and quantizes it on the forward pass. SR4 keeps the same idea of a latent accumulator, but at 4 bits instead of 32. Stochastic rounding supplies the sub-integer resolution the missing 28 bits used to provide. No FP32 shadow weights anywhere.
How it works at a glance
Each weight has one integer accumulator q ∈ [-7, +7] (15 states → a nibble):
- Forward (ternarize):
w = 0 if |q| < 2 else sign(q). A weight must build conviction (|q| ≥ 2) before it fires — a built-in hysteresis / dead-zone. - Update (integrate): each batch, nudge
qtoward the negative gradient by a tiny amount, using stochastic rounding so fractional-LSB steps are realized as an occasional ±1.qis a slow integer integrator of gradient direction.
The stochastic rounding is driven by a stateless counter-hash (SplitMix64
keyed on (seed, counter, weight-index)) — nothing is stored, reproducible,
parallelizable, bit-identical across compiler flags. See
docs/HOW_IT_WORKS.md.
Quick start
#include "ternary_optim.h"
ternary_optim_config cfg;
ternary_optim_config_default(&cfg); /* target_step=0.001, q_clamp=7,
theta=2, dither on, seed=42 */
ternary_optim *opt = ternary_optim_create(rows, cols, &cfg);
float *W = malloc(rows*cols*sizeof(float)); /* caller-owned scratch */
float *grad_sum = malloc(rows*cols*sizeof(float));
uint64_t counter = 0;
for (int step = 0; step < n_steps; step++) {
ternary_optim_materialize(opt, W); /* q -> ternary weights */
/* ... forward with W, backward -> fill grad_sum (summed over the batch) ...*/
ternary_optim_step(opt, grad_sum, batch_size, counter, /*target_step*/ -1.0f);
counter++; /* one increment per update */
}
ternary_optim_free(opt);
Build and run the demo:
make run
What's validated (and what isn't)
Validated on MNIST (784-256-10 MLP), memory-neutral at 4 b/w:
- Ternary SR4 reaches ~98.3, essentially the FP32 non-ternary ceiling (98.36) of the same architecture.
- Beat the prior V1 gold standard (INT16 accumulator, 16 b/w) at a quarter of the memory.
Not yet validated: LM / transformer scale. Adam earns its state most on hard, large-scale loss surfaces; whether SR4 holds parity there is the open question.
The recipe (optimizer core vs. the knobs)
The core — the int4 accumulator + stochastic-rounded integer update + ternary forward + low-discrepancy dither — is what this library ships. Everything else is a recipe that lives outside the optimizer, in the training loop or the loss:
- Batch size — the dominant accuracy lever. Training loop.
- Co-trained per-column output scales — restore ternary dynamic range
(~0.042 b/w). Provided here as optional decoupled helpers
(
ternary_apply_col_scale,ternary_col_scale_grad); the scale vector belongs to the layer, not the optimizer. - Label smoothing — loss-side regularizer. Loss.
- Annealed dropout — variance/reliability control. Training loop.
- Cosine step-size anneal — schedule; pass the annealed value as the
target_stepargument toternary_optim_step().
Deferred / stubbed (honest scope)
This is a first-pass CPU reference of the core mechanism. Explicitly NOT done here:
- State serialization —
ternary_optim_save/ternary_optim_loadare stubs returningTERNARY_OPTIM_ENOTIMPL. The thing to persist is the packedq(4 b/w) + shape + config; the format is a TODO. - GPU kernels — CPU scalar reference only.
- foundry integration — not wired into dcherrera/foundry's tensor + optim framework; that is the follow-up.
- Scale-generalization testing — the per-column scale helpers are generalized
to
(rows, cols)but have only been validated in the MNIST MLP shape. - Multi-layer convenience — one handle owns one matrix; a model wires several
handles (give each a distinct
idx_base). No container type yet.
Provenance
Mechanism source of truth: the sr_* core of the SR4 research prototype (not
published) and its design notes, which docs/HOW_IT_WORKS.md adapts. The update
arithmetic here is a faithful, decoupled lift of that code.