# 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](HOW_IT_WORKS.md); for
exact signatures see [API.md](API.md). A runnable end-to-end demo is
[`../example.c`](../example.c) (`make run`).

---

## 1. The one-matrix cycle

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

```c
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
```

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

  ```c
  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`:

```c
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:

```c
// 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)`.