7.0 KiB

Raw Blame History

ternary_optim — Wiring it into a training loop

This is the operational guide: the materialize → forward → backward → update cycle, how gradients must be provided, the batch-counter contract, and where the recipe knobs live. For the mechanism see HOW_IT_WORKS.md; for exact signatures see API.md. A runnable end-to-end demo is ../example.c (make run).

1. The one-matrix cycle

The optimizer owns exactly one weight matrix's q state. Per step:

ternary_optim_materialize(opt, W);   // 1. q -> ternary W  (caller buffer)
// 2. forward:  use W (and optional scales/bias) to compute activations
// 3. backward: produce grad_sum — dL/dW SUMMED over the batch
ternary_optim_step(opt, grad_sum, batch_size, counter, target_step);   // 4.
counter++;                            // 5. exactly one increment per step

Materialize turns q into the ternary weights {-1,0,+1} your forward math consumes. Do it every step — q moves under the update.
Forward is entirely yours (dense matmul, BLAS, conv, whatever). The optimizer never sees your activations.
Backward is yours too; it must produce grad_sum, the per-weight gradient summed over the batch samples (see §2).
Step integrates grad_sum into q.
Increment the counter (see §3).

2. How gradients must be provided

grad_sum is dL/dW summed over the batch, in the same row-major layout as the weight matrix (element (r,c) at r*cols + c), at least numel floats.

Summed, not averaged. The optimizer divides by batch_size internally (g = grad_sum / batch_size) and then renormalizes per-matrix (s = target_step / (mean|g| + 1e-12)). Because of that per-matrix normalization the update is invariant to the overall gradient scale — but the batch_size you pass must match the number of samples actually summed, or the mean is wrong (this matters for a partial final batch).
It is the gradient w.r.t. the ternary weights (W), i.e. the same W you got from materialize. If you fold per-column scales into the forward (Z = scale ⊙ (W·x) + bias), the scale must be pushed back into the delta before the weight-grad matmul so grad_sum reflects the scaled forward — see §5.
No gradient of the ternarization itself. SR4 is straight-through in spirit: you compute the gradient as if W were the real weights; the accumulator q absorbs it. There is no dW/dq term.

grad_sum is read-only to the optimizer and stays caller-owned; reuse the same buffer every step (zero it before accumulating).

3. The batch-counter contract

batch_counter is a uint64_t global update index, starting at 0, that drives the stateless dither. Rules:

Start at 0, increment by exactly 1 per logical update step. Think updates, not epochs — the counter is the fundamental clock of the optimizer.
All matrices in one step share the same counter value. If your model has several ternary_optim handles (one per weight matrix), call step on each with the same batch_counter, then increment once:
```
ternary_optim_step(W1_opt, g1, bs, counter, ts);
ternary_optim_step(W2_opt, g2, bs, counter, ts);   // same counter
counter++;                                          // one increment for the step
```
Give each handle a distinct idx_base (non-overlapping [idx_base, idx_base+numel) ranges) so their stochastic draws are independent even at the same counter. E.g. W1 uses idx_base = 0, W2 uses idx_base = W1.numel.
Reproducibility. Same seed, same idx_base, same counter sequence, same gradients ⇒ bit-identical q. The counter must therefore be deterministic across a resume (part of what a future checkpoint must persist).

4. The recipe lives outside the optimizer

The optimizer is just the int4 update. The accuracy recipe lives where it belongs:

Knob	Where it lives	How
Batch size	training loop	dominant accuracy lever; bigger batch = cleaner per-step gradient = tighter seed band. Pass the true `batch_size` to `step`.
Cosine step anneal	training loop → `target_step` arg	compute the annealed value each step and pass it as `target_step` (pass `<= 0` to use the config default). Reference: cosine from `target_step` to `target_step/10` over all updates.
Label smoothing	loss	fold into your CE targets; the softmax gradient stays `p - t`. Zero optimizer involvement.
Annealed dropout	training loop	apply your own mask to activations in forward and to `dL/dh` in backward with identical keys; nothing to store if you use a counter-hash mask. Reference recipe: linear anneal of `p` to 0 over the first ~25% of training.
Per-column scales	the layer	see §5.

Matched-updates thinking: when you change batch size, hold the number of updates (hence the counter range and the cosine schedule length) fixed, not the number of epochs.

5. Optional per-column output scales

Ternary weights lose magnitude; a small FP32 per-output-column scale vector restores dynamic range. The vector is owned by your layer, not the optimizer. Init to all-ones (exact identity) and train by plain SGD at a dedicated small LR (reference: SCALE_LR = 0.001, ~10× below the bias LR).

Forward fold — Z = scale ⊙ Z_raw + bias:

ternary_optim_materialize(opt, W);
// Z_raw = X @ W            (your matmul)
memcpy(Z_raw_saved, Z_raw, rows*cols*sizeof(float));  // keep for the scale grad
ternary_apply_col_scale(Z, scale, rows, cols);        // Z *= scale (per column)
// add bias, activation, ...

Backward — take the bias/scale grads from the unscaled delta, then push the scale back into the delta before the weight-grad matmul:

// D = dL/dZ  (unscaled upstream delta)
ternary_col_scale_grad(D, Z_raw_saved, scale_grad, rows, cols);  // dL/dscale
// bias_grad = column-sum of D  (your code)
ternary_apply_col_scale(D, scale, rows, cols);   // D *= scale -> dL/dZ_raw
// grad_sum = X^T @ D        (your weight-grad matmul, feeds ternary_optim_step)

// SGD updates (your code):
for (c) scale[c] -= SCALE_LR * scale_grad[c] / batch_size;
for (c) bias[c]  -= BIAS_LR  * bias_grad[c]  / batch_size;

The optimizer never touches scale/bias; they are ordinary FP32 SGD parameters you maintain alongside q.

6. Evaluation / deployment

At eval time there is no update and no dropout — just materialize the ternary weights, run the forward with the scales/biases, and read predictions. The deployed artifact is ternarize(q) (the ternary matrix) plus the FP32 scales and biases; nothing else needs to persist.

7. Checkpoint / resume (currently stubbed)

ternary_optim_save/_load return TERNARY_OPTIM_ENOTIMPL today. Until they land, to snapshot state yourself, persist: the packed q (ternary_optim_q(opt), numel bytes — 4 b/w of real information), the shape, the config, your scale/bias vectors, and the batch_counter. The dither/masks need not be saved — they regenerate from (seed, counter, index).

7.0 KiB Raw Blame History Unescape Escape