117 lines
4.6 KiB
C
117 lines
4.6 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* example.c — minimal, self-contained SR4 usage demo (no MNIST, no BLAS).
|
|
*
|
|
* Synthetic task: recover a fixed random TERNARY "teacher" weight matrix W*
|
|
* (rows x cols) from noiseless linear regression data. A batch of random inputs
|
|
* X (batch x rows) is mapped by the teacher to targets Y* = X @ W*. SR4 trains
|
|
* its own ternary W to minimise the mean-squared error ||X @ W - Y*||^2.
|
|
*
|
|
* It exercises the full library cycle every step:
|
|
* init -> ( materialize -> forward -> backward -> update ) -> free
|
|
* and prints the loss falling as SR4's ternary weights converge on the teacher.
|
|
*
|
|
* Build: make (produces ./example)
|
|
* Run: ./example
|
|
*/
|
|
|
|
#include "ternary_optim.h"
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <math.h>
|
|
|
|
#define ROWS 16 /* input dim (matrix rows) */
|
|
#define COLS 4 /* output dim (matrix cols) */
|
|
#define BATCH 64
|
|
#define STEPS 4000
|
|
|
|
/* tiny local xorshift so the demo is self-contained */
|
|
static uint64_t g_rng = 12345;
|
|
static float frand(void) {
|
|
uint64_t x = g_rng; x ^= x << 13; x ^= x >> 7; x ^= x << 17; g_rng = x;
|
|
return (float)(x >> 11) / (float)(1ULL << 53);
|
|
}
|
|
|
|
int main(void) {
|
|
const int numel = ROWS * COLS;
|
|
|
|
/* Fixed ternary teacher W* in {-1,0,+1}. */
|
|
float *W_star = malloc(numel * sizeof(float));
|
|
for (int i = 0; i < numel; i++) {
|
|
float r = frand();
|
|
W_star[i] = (r < 0.33f) ? -1.0f : (r < 0.66f) ? 0.0f : 1.0f;
|
|
}
|
|
|
|
/* Create the optimizer for a (ROWS x COLS) weight matrix. */
|
|
ternary_optim_config cfg;
|
|
ternary_optim_config_default(&cfg);
|
|
cfg.seed = 7;
|
|
ternary_optim *opt = ternary_optim_create(ROWS, COLS, &cfg);
|
|
if (!opt) { fprintf(stderr, "create failed\n"); return 1; }
|
|
|
|
/* Caller-owned scratch buffers (the loop owns these, not the optimizer). */
|
|
float *W = malloc(numel * sizeof(float)); /* materialized */
|
|
float *grad_sum = malloc(numel * sizeof(float)); /* dL/dW summed */
|
|
float *X = malloc(BATCH * ROWS * sizeof(float)); /* batch inputs */
|
|
float *pred = malloc(BATCH * COLS * sizeof(float));
|
|
float *target = malloc(BATCH * COLS * sizeof(float));
|
|
float *err = malloc(BATCH * COLS * sizeof(float));
|
|
|
|
printf("SR4 example: recover a %dx%d ternary teacher via linear regression\n",
|
|
ROWS, COLS);
|
|
printf("step | mse\n");
|
|
printf("-----+---------\n");
|
|
|
|
for (uint64_t step = 0; step < STEPS; step++) {
|
|
/* ---- materialize the ternary weights for this forward pass ---- */
|
|
ternary_optim_materialize(opt, W);
|
|
|
|
/* ---- draw a fresh batch and build targets from the teacher ---- */
|
|
for (int s = 0; s < BATCH; s++)
|
|
for (int r = 0; r < ROWS; r++)
|
|
X[s * ROWS + r] = frand() * 2.0f - 1.0f;
|
|
for (int s = 0; s < BATCH; s++)
|
|
for (int c = 0; c < COLS; c++) {
|
|
float acc_star = 0.0f, acc_pred = 0.0f;
|
|
for (int r = 0; r < ROWS; r++) {
|
|
acc_star += X[s * ROWS + r] * W_star[r * COLS + c];
|
|
acc_pred += X[s * ROWS + r] * W[r * COLS + c];
|
|
}
|
|
target[s * COLS + c] = acc_star; /* forward (teacher) */
|
|
pred[s * COLS + c] = acc_pred; /* forward (student) */
|
|
}
|
|
|
|
/* ---- loss + backward: err = pred - target, grad = X^T @ err ---- */
|
|
double mse = 0.0;
|
|
for (int i = 0; i < BATCH * COLS; i++) {
|
|
err[i] = pred[i] - target[i];
|
|
mse += (double)err[i] * err[i];
|
|
}
|
|
mse /= (BATCH * COLS);
|
|
|
|
for (int i = 0; i < numel; i++) grad_sum[i] = 0.0f;
|
|
for (int s = 0; s < BATCH; s++)
|
|
for (int r = 0; r < ROWS; r++)
|
|
for (int c = 0; c < COLS; c++)
|
|
grad_sum[r * COLS + c] += X[s * ROWS + r] * err[s * COLS + c];
|
|
|
|
/* ---- SR4 update: integrate the batch gradient into q ---- */
|
|
ternary_optim_step(opt, grad_sum, BATCH, step, -1.0f /* cfg default */);
|
|
|
|
if (step % 400 == 0 || step == STEPS - 1)
|
|
printf("%4llu | %.5f\n", (unsigned long long)step, mse);
|
|
}
|
|
|
|
/* Report how many ternary weights match the teacher exactly. */
|
|
ternary_optim_materialize(opt, W);
|
|
int match = 0;
|
|
for (int i = 0; i < numel; i++) if (W[i] == W_star[i]) match++;
|
|
printf("\nrecovered %d / %d teacher weights exactly\n", match, numel);
|
|
|
|
ternary_optim_free(opt);
|
|
free(W_star); free(W); free(grad_sum); free(X); free(pred); free(target); free(err);
|
|
return 0;
|
|
}
|