/* SPDX-License-Identifier: AGPL-3.0-or-later */ /* Copyright (C) 2026 David (Corey) Herrera */ /** * @file ternary_optim.h * @brief SR4 — 4-bit Stochastic-Rounded Latent Accumulator (ternary optimizer). * * SR4 trains a network whose DEPLOYED weights are ternary ({-1, 0, +1}) while * carrying an ENTIRE persistent per-weight training state of a single 4-bit * signed integer accumulator `q` — no FP32 master weight, no momentum, no Adam * moments. Stochastic rounding (driven by a stateless counter-hash) supplies the * sub-integer resolution that a low-bit state would otherwise lack. * * This header is the CLEAN, MNIST-decoupled public API of the optimizer core, * generalized to an arbitrary weight matrix of shape (rows x cols). It knows * nothing about MNIST, the 784-256-10 MLP, data loading, or the training loop — * those live in the caller. * * The numerics (stochastic rounding, low-discrepancy dither, per-matrix step * normalization, saturating add, the SplitMix64 counter-hash) are lifted * faithfully and bit-for-bit from the SR4 research prototype's `sr_*` core. * * SCOPE (honest): this is a first-pass CPU REFERENCE library of the core * mechanism, validated on MNIST-MLP. See README.md for what is stubbed/deferred * (state serialization, GPU kernels, foundry tensor/optim integration, * scale-generalization testing). * * Threading / determinism: every stochastic draw is a pure function of * (seed, batch_counter, global weight index) via an integer counter-hash — no * stored RNG stream. Results are reproducible per seed and bit-identical across * compiler optimization / fast-math flags. */ #ifndef TERNARY_OPTIM_H #define TERNARY_OPTIM_H #include #include #ifdef __cplusplus extern "C" { #endif /* --------------------------------------------------------------------------- * Return codes * ------------------------------------------------------------------------- */ enum { TERNARY_OPTIM_OK = 0, /**< success */ TERNARY_OPTIM_EINVAL = -1, /**< invalid argument */ TERNARY_OPTIM_ENOMEM = -2, /**< allocation failure */ TERNARY_OPTIM_ENOTIMPL = -3 /**< not implemented (stubbed feature) */ }; /* --------------------------------------------------------------------------- * Configuration * ------------------------------------------------------------------------- */ /** * @brief Immutable configuration for one optimizer instance. * * Populate with ternary_optim_config_default() and then override fields as * needed. All fields are copied into the optimizer at create time. */ typedef struct { /** Average proposed step, in LSB, for a weight whose |gradient| equals the * matrix mean |gradient|. Small (~0.001): `q` INTEGRATES gradient * direction over many updates rather than jumping. The training loop may * additionally anneal an effective per-step value (e.g. cosine) and pass * it to ternary_optim_step(); this default is used when a step call passes * a non-positive target_step. */ float target_step; /** Symmetric saturation of the accumulator: q is clamped to [-q_clamp, * +q_clamp]. Reference value 7 (15 states -> a 4-bit nibble). */ int q_clamp; /** Dead-zone threshold: the deployed weight is 0 while |q| < theta, else * sign(q). Reference value 2 (a weight must build conviction |q|>=2 before * it fires; hysteresis prevents chatter). */ int theta; /** 1 = stratified Weyl low-discrepancy dither (variance-reduced, the * reference/default); 0 = plain i.i.d. counter-hash dither (unbiased but * noisier). Both are unbiased; LD tightens the seed-to-seed band. */ int use_dither; /** Seed threaded through EVERY weight on EVERY update via the counter-hash. * Also seeds the initial q distribution. Two instances that must diverge * should differ in either `seed` or `idx_base`. */ uint64_t seed; /** Global weight-index base for this matrix in the counter-hash lane. When * a model has several matrices, give each a distinct, non-overlapping * [idx_base, idx_base+numel) range so their stochastic draws are * independent. A single-matrix caller may leave this 0. */ int64_t idx_base; } ternary_optim_config; /** * @brief Fill @p cfg with the reference defaults. * * target_step=0.001, q_clamp=7, theta=2, use_dither=1, seed=42, idx_base=0. * * @param cfg out: configuration to initialize (must be non-NULL). */ void ternary_optim_config_default(ternary_optim_config *cfg); /* --------------------------------------------------------------------------- * Optimizer handle * ------------------------------------------------------------------------- */ /** Opaque optimizer / param-state handle. One instance owns the `q` state for * exactly one (rows x cols) weight matrix. */ typedef struct ternary_optim ternary_optim; /** * @brief Create an optimizer for a (rows x cols) row-major weight matrix. * * Allocates the int4 accumulator state (one signed byte per weight; the * INFORMATION content is 4 bits/weight) and initializes it to {-3, 0, +3} at * {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 initial * distribution is drawn from an internal xorshift RNG seeded from * (cfg->seed ^ cfg->idx_base), so distinct matrices get distinct inits. * * @param rows matrix rows (> 0). * @param cols matrix columns (> 0). * @param cfg configuration; copied. If NULL, defaults are used. * @return owning handle, or NULL on invalid args / allocation failure. * Free with ternary_optim_free(). */ ternary_optim *ternary_optim_create(int rows, int cols, const ternary_optim_config *cfg); /** * @brief Release an optimizer and its state. Safe to call with NULL. * @param opt handle to free (may be NULL). */ void ternary_optim_free(ternary_optim *opt); /* --------------------------------------------------------------------------- * Core operations * ------------------------------------------------------------------------- */ /** * @brief Materialize the ternary deployed weights from `q` into a caller buffer. * * out[i] = 0 if |q[i]| < theta * +1 if q[i] >= +theta * -1 if q[i] <= -theta * * Call this before each forward pass; the result is the actual weight matrix * the forward math consumes. * * @param opt optimizer (non-NULL). * @param out_weights out: caller-owned buffer of at least numel floats, * filled row-major with values in {-1, 0, +1}. */ void ternary_optim_materialize(const ternary_optim *opt, float *out_weights); /** * @brief One SR4 update step: integrate the batch gradient into `q`. * * Given the per-weight gradient SUMMED over the batch (same row-major layout as * the weight matrix) and the batch size, this: * 1. normalizes the step per-matrix: s = target_step / (mean|g| + 1e-12), * where g = grad_sum / batch_size — so a typical weight proposes ~ * target_step LSB regardless of raw gradient scale; * 2. snaps the proposed step to 8.8 fixed point and splits it into an integer * part (applied deterministically) and a fractional part (applied * stochastically as an occasional +1), so E[increment] is exact (unbiased); * 3. draws the stochastic bit from the counter-hash — either the stratified * Weyl low-discrepancy dither (use_dither=1) or plain i.i.d. (use_dither=0); * 4. saturating-adds into q, clamped to [-q_clamp, +q_clamp]. * * The only side effect is the integer add into the internal q state. Nothing is * stored beyond q: the dither is regenerated from (seed, batch_counter, index). * * @param opt optimizer (non-NULL). * @param grad_sum per-weight gradient summed over the batch (>= numel * floats, row-major, matching the weight layout). * @param batch_size number of samples summed into grad_sum (> 0). Used for * the /batch_size mean; may vary (e.g. a partial final * batch). * @param batch_counter monotonically increasing global update index, starting * at 0. All matrices updated within the SAME step use the * SAME counter value; the caller increments it by 1 AFTER * stepping every matrix. Drives the stateless dither; see * docs/USAGE.md. * @param target_step effective per-step target (e.g. a cosine-annealed * value). If <= 0, the config default is used. */ void ternary_optim_step(ternary_optim *opt, const float *grad_sum, int batch_size, uint64_t batch_counter, float target_step); /* --------------------------------------------------------------------------- * Accessors (diagnostics) * ------------------------------------------------------------------------- */ /** @return number of weights (rows*cols), or 0 if opt is NULL. */ int ternary_optim_numel(const ternary_optim *opt); /** @return matrix rows, or 0 if opt is NULL. */ int ternary_optim_rows(const ternary_optim *opt); /** @return matrix cols, or 0 if opt is NULL. */ int ternary_optim_cols(const ternary_optim *opt); /** * @brief Read-only view of the raw int4 accumulator (one signed byte/weight, * clamped to [-q_clamp, +q_clamp]). For diagnostics/fingerprints only — * the update rule never reads it back. Valid until the optimizer is * freed. * @return pointer to numel int8 values, or NULL if opt is NULL. */ const int8_t *ternary_optim_q(const ternary_optim *opt); /* --------------------------------------------------------------------------- * Optional per-column output scales (recipe adjunct, decoupled) * * Ternary weights throw away magnitude. These free helpers restore per-output- * column dynamic range via an FP32 scale vector folded into the forward pass: * Z = scale (*) Z_raw + bias. * The scale vector is NOT owned by the optimizer (it belongs to the layer); * these operate on caller buffers over an arbitrary (rows x cols) matrix. The * caller allocates the scale vector (length cols, init 1.0) and trains it by * plain SGD at a small dedicated learning rate. See docs/USAGE.md. * ------------------------------------------------------------------------- */ /** * @brief In place, multiply each column c of a row-major (rows x cols) matrix * by scale[c]. Exact identity when scale == all-ones. * @param Z in/out: row-major matrix (rows*cols floats). * @param scale length-cols per-column scale vector. * @param rows matrix rows. * @param cols matrix cols. */ void ternary_apply_col_scale(float *Z, const float *scale, int rows, int cols); /** * @brief Per-column scale gradient: out[c] = sum_r D[r,c] * Z_raw[r,c]. * * D is the UNSCALED upstream delta (dL/dZ before folding the scale back in) and * Z_raw the pre-scale, pre-bias matmul output. `out` is OVERWRITTEN (not * accumulated). Apply to the scale vector by plain SGD at a dedicated small LR. * * @param D upstream delta, row-major (rows*cols floats). * @param Z_raw pre-scale/pre-bias matmul output, row-major (rows*cols floats). * @param out out: length-cols gradient (overwritten). * @param rows matrix rows. * @param cols matrix cols. */ void ternary_col_scale_grad(const float *D, const float *Z_raw, float *out, int rows, int cols); /* --------------------------------------------------------------------------- * State serialization (STUBBED — TODO: checkpoint / resume) * * The thing to persist is the packed `q` array (4 bits/weight) plus the shape * and config; the dither and any masks are regenerated from the counter-hash so * they never need saving. These are intentionally NOT implemented in this * first-pass reference library: both return TERNARY_OPTIM_ENOTIMPL. See * README.md "Deferred / stubbed". * ------------------------------------------------------------------------- */ /** * @brief TODO/STUB: save optimizer state (packed q + shape + config) to @p path. * @return TERNARY_OPTIM_ENOTIMPL always (not yet implemented). */ int ternary_optim_save(const ternary_optim *opt, const char *path); /** * @brief TODO/STUB: load optimizer state previously written by * ternary_optim_save(). * @param path path to a saved state file. * @param out out: receives a newly created handle on success. * @return TERNARY_OPTIM_ENOTIMPL always (not yet implemented); *out set to NULL. */ int ternary_optim_load(const char *path, ternary_optim **out); #ifdef __cplusplus } /* extern "C" */ #endif #endif /* TERNARY_OPTIM_H */