@seanhogg/builderforce-memory-engine 2026.9.34 → 2026.10.1

package/src/ml/federate.ts

@@ -0,0 +1,160 @@
+/**
+ * ml/federate.ts — trained ML heads federate like the model's own weights.
+ *
+ * A head's trained state is ONE flat f32 vector (`MlPredictor.exportParams`),
+ * stored as a CRC-trailed checkpoint so it is the same shape the delta pipeline
+ * already handles: a learner diffs its trained copy against the version it
+ * started from (`diffCheckpoints`), ships it on the existing delta wire with
+ * `target: "ml:<task>"`, and the writer folds N of them with the existing FedAvg
+ * (`mergeCheckpointDiffsGated` → `mergeCheckpointDiffs`), refusing a merge whose
+ * held-out loss regresses. No second wire, no second merge.
+ */
+import { buildDeltaLearnPayload, decodeDeltaPayload, deltaUnusableReason, mlHeadTarget, type DeltaLearnPayload } from "../evermind/delta_wire.js";
+import { mergeCheckpointDiffsGated, type GatedMergeResult, type MergeGateOptions } from "../evermind/merge.js";
+import type { EvermindModelPackage } from "../evermind/package.js";
+import { base64ToBytes, bytesToBase64 } from "../utils/base64.js";
+import { appendCrcTrailer, verifyCrcTrailer } from "../utils/crc32.js";
+import { diffCheckpoints } from "../utils/delta.js";
+
+/**
+ * One trained ML-catalog head carried in a model's manifest (`mlHeads`): small, so
+ * it rides in the manifest, and versioned, so it federates — a delta against it is
+ * `target: "ml:<task>"`, `baseVersion: version` on the delta wire.
+ */
+export interface EvermindMlHead {
+  /** Catalog task id, e.g. `run.local-success`. */
+  task: string;
+  /** Engine key, e.g. `local:classical:logistic`. */
+  engine: string;
+  /** Integer head version — what a federated delta for it is diffed against. */
+  version: number;
+  /** What a loader needs to rebuild the predictor (e.g. `{ dim, classes }`). */
+  shape: Record<string, number>;
+  /** base64 of the CRC-trailed f32 checkpoint of the head's flat parameters. */
+  params: string;
+  /** Held-out loss of this version, when measured. */
+  evalLoss?: number;
+}
+
+/** A head's flat parameters as a CRC-trailed f32 checkpoint. */
+export function headCheckpoint(params: ArrayLike<number>): ArrayBuffer {
+  const body = new ArrayBuffer(params.length * 4);
+  new Float32Array(body).set(Array.from(params));
+  return appendCrcTrailer(body);
+}
+
+/** The flat parameters of a head checkpoint. Throws on a failed CRC. */
+export function headParams(checkpoint: ArrayBuffer): Float32Array {
+  const crc = verifyCrcTrailer(checkpoint);
+  if (!crc.ok) throw new Error("headParams: checkpoint failed its CRC check");
+  if (crc.body.byteLength % 4 !== 0) throw new Error("headParams: checkpoint is not f32");
+  return new Float32Array(crc.body);
+}
+
+/** Why a manifest head entry is unusable, or null (the package's validate() calls this). */
+export function mlHeadUnusableReason(h: EvermindMlHead): string | null {
+  if (!h || typeof h.task !== "string" || typeof h.engine !== "string") return "task and engine are required";
+  if (!Number.isInteger(h.version) || h.version < 0) return "version must be a non-negative integer";
+  try {
+    const crc = verifyCrcTrailer(base64ToBytes(h.params));
+    if (!crc.hasTrailer || !crc.ok) return "params checkpoint failed its CRC check";
+    if (crc.body.byteLength % 4 !== 0) return "params checkpoint is not f32";
+  } catch {
+    return "params is not base64";
+  }
+  return null;
+}
+
+/** A manifest entry for a head's current parameters. */
+export function mlHeadEntry(
+  task: string,
+  engine: string,
+  version: number,
+  params: ArrayLike<number>,
+  shape: Record<string, number> = {},
+  evalLoss?: number,
+): EvermindMlHead {
+  return {
+    task,
+    engine,
+    version,
+    shape,
+    params: bytesToBase64(headCheckpoint(params)),
+    ...(evalLoss != null && Number.isFinite(evalLoss) ? { evalLoss } : {}),
+  };
+}
+
+/** The entry's flat parameters. */
+export function mlHeadParams(entry: EvermindMlHead): Float32Array {
+  return headParams(base64ToBytes(entry.params));
+}
+
+/** The head a package carries for `task`, if any. */
+export function packageMlHead(pkg: EvermindModelPackage, task: string): EvermindMlHead | undefined {
+  return pkg.manifest.mlHeads?.find((h) => h.task === task);
+}
+
+/** `pkg` with `entry` as its head for `entry.task` (replacing the old one). */
+export function withPackageMlHead(pkg: EvermindModelPackage, entry: EvermindMlHead): EvermindModelPackage {
+  return pkg.withMlHeads([...(pkg.manifest.mlHeads ?? []).filter((h) => h.task !== entry.task), entry]);
+}
+
+/**
+ * The delta-wire payload for a locally trained head: its new parameters diffed
+ * against the version it started from. The caller owns the size guard, as for LM deltas.
+ */
+export function buildMlHeadDelta(base: EvermindMlHead, trained: ArrayLike<number>, weight?: number, label?: string): DeltaLearnPayload {
+  const diff = diffCheckpoints(base64ToBytes(base.params), headCheckpoint(trained));
+  return buildDeltaLearnPayload(diff, base.version, weight, label ?? `ml head ${base.task}`, mlHeadTarget(base.task));
+}
+
+export interface MlHeadMergeOutcome {
+  /** The head after the round: version + 1 when accepted, the base unchanged when refused. */
+  head: EvermindMlHead;
+  merge: GatedMergeResult;
+  /** Contributions left out of the batch, by index, with why. */
+  rejected: { index: number; reason: string }[];
+}
+
+/**
+ * One federated round for a head: validate each contribution ON ITS OWN (wrong
+ * target, stale base or a malformed diff costs only itself), FedAvg the rest, and
+ * keep the result only if `score` (held-out loss of a parameter vector, lower is
+ * better) does not regress.
+ */
+export function mergeMlHeadDeltas(
+  base: EvermindMlHead,
+  payloads: readonly DeltaLearnPayload[],
+  score: (params: Float32Array) => number | null,
+  opts: MergeGateOptions = {},
+): MlHeadMergeOutcome {
+  const baseCp = base64ToBytes(base.params);
+  const target = mlHeadTarget(base.task);
+  const diffs: ArrayBuffer[] = [];
+  const weights: number[] = [];
+  const rejected: { index: number; reason: string }[] = [];
+  payloads.forEach((p, index) => {
+    if (p.target !== target) return void rejected.push({ index, reason: `target ${p.target ?? "(model)"} is not ${target}` });
+    if (p.baseVersion !== base.version) return void rejected.push({ index, reason: `stale base version ${p.baseVersion} (head is ${base.version})` });
+    let diff: ArrayBuffer;
+    try {
+      diff = decodeDeltaPayload(p);
+    } catch (error) {
+      return void rejected.push({ index, reason: `diff is not base64: ${error instanceof Error ? error.message : String(error)}` });
+    }
+    const why = deltaUnusableReason(diff, baseCp);
+    if (why) return void rejected.push({ index, reason: why });
+    diffs.push(diff);
+    weights.push(p.weight ?? 1);
+  });
+
+  const merge = mergeCheckpointDiffsGated(baseCp, diffs, weights, (cp) => score(headParams(cp)), opts);
+  if (!merge.accepted) return { head: base, merge, rejected };
+  const head: EvermindMlHead = {
+    ...base,
+    version: base.version + 1,
+    params: bytesToBase64(merge.checkpoint),
+    ...(merge.newLoss != null ? { evalLoss: merge.newLoss } : {}),
+  };
+  return { head, merge, rejected };
+}

package/src/ml/head_trainer.ts

@@ -0,0 +1,268 @@
+/**
+ * ml/head_trainer.ts — HeadTrainer: the ONE gradient trainer for a small head over
+ * an embedding (extracted from the limbic affect trainer, which is now a thin
+ * registration over it).
+ *
+ * Full-batch gradient descent with AdamW. When a GPUDevice is supplied the
+ * optimiser step runs on the GPU via the shared WEIGHT_UPDATE_WGSL kernel; with
+ * no device it uses a numerically-equivalent CPU AdamW, so training works
+ * everywhere (CI, Node without @webgpu/node, a laptop with no GPU).
+ *
+ * The head owns its forward/backward and its loss (MSE for regression, BCE /
+ * softmax cross-entropy for classification); this file owns only the loop:
+ * zero → accumulate per sample → average → clip → AdamW step.
+ */
+
+import {
+  createUniformBuffer,
+  createStorageBuffer,
+  createComputePipeline,
+  createBindGroup,
+  dispatchKernel,
+  readBuffer,
+  cdiv,
+} from "../utils/gpu_utils.js";
+import { WEIGHT_UPDATE_WGSL } from "../kernels/weight_update.js";
+
+/** A named trainable tensor (flat row-major Float32Array). */
+export interface HeadParam {
+  name: string;
+  data: Float32Array;
+  numel: number;
+}
+
+/** What a head must expose to be trained by {@link HeadTrainer}. */
+export interface TrainableHead<S> {
+  /** Trainable parameters, in a stable order. */
+  parameters(): HeadParam[];
+  /** Gradient buffers, index-aligned with {@link parameters}. */
+  gradients(): HeadParam[];
+  zeroGrad(): void;
+  /** Accumulate gradients for one sample into {@link gradients}; returns its loss. */
+  accumulate(sample: S): number;
+  /** Forward-only loss of one sample. */
+  loss(sample: S): number;
+}
+
+export interface HeadTrainOptions {
+  learningRate?: number;
+  epochs?: number;
+  weightDecay?: number;
+  beta1?: number;
+  beta2?: number;
+  eps?: number;
+  /** Max global gradient L2 norm before the optimiser step. Default 1.0. */
+  maxGradNorm?: number;
+  /** Per-epoch hook. `gradNorm` is the pre-clip gradient L2 norm this epoch (the
+   *  instability early-warning) — computed CPU-side so it's always available. */
+  onEpochEnd?: ((epoch: number, loss: number, gradNorm?: number) => void) | null;
+}
+
+interface AdamMoment {
+  m: Float32Array;
+  v: Float32Array;
+}
+
+interface AdamHyper {
+  learningRate: number;
+  weightDecay: number;
+  beta1: number;
+  beta2: number;
+  eps: number;
+  beta1_t: number;
+  beta2_t: number;
+}
+
+function packAdamParams(
+  numElements: number,
+  lr: number,
+  beta1: number,
+  beta2: number,
+  eps: number,
+  weightDecay: number,
+  beta1_t: number,
+  beta2_t: number,
+): ArrayBuffer {
+  const buf = new ArrayBuffer(32);
+  new Uint32Array(buf, 0, 1).set([numElements]);
+  new Float32Array(buf, 4, 7).set([lr, beta1, beta2, eps, weightDecay, beta1_t, beta2_t]);
+  return buf;
+}
+
+export class HeadTrainer<S> {
+  readonly head: TrainableHead<S>;
+  readonly device: GPUDevice | null;
+  private _moments: AdamMoment[] | null = null;
+  private _step = 0;
+  private readonly _adamwPipeline: GPUComputePipeline | null;
+
+  constructor(
+    head: TrainableHead<S>,
+    device: GPUDevice | null = null,
+    /** Name used in error messages (the registration that owns this trainer). */
+    private readonly label: string = "HeadTrainer",
+  ) {
+    this.head = head;
+    this.device = device;
+    this._adamwPipeline = device
+      ? createComputePipeline(device, WEIGHT_UPDATE_WGSL, "adamw_update")
+      : null;
+  }
+
+  /** Whether the optimiser step runs on the GPU. */
+  get gpuTraining(): boolean {
+    return this.device != null && this._adamwPipeline != null;
+  }
+
+  private _initMoments(): void {
+    if (this._moments) return;
+    this._moments = this.head.parameters().map((p) => ({
+      m: new Float32Array(p.numel),
+      v: new Float32Array(p.numel),
+    }));
+  }
+
+  /**
+   * Train on `samples` for `epochs` passes; returns the per-epoch mean loss.
+   * Full-batch: grads accumulate over every sample, are averaged, clipped, then
+   * applied once per epoch.
+   */
+  async train(samples: S[], opts: HeadTrainOptions = {}): Promise<number[]> {
+    if (samples.length === 0) throw new Error(`${this.label}.train: no samples`);
+    const {
+      learningRate = 0.05,
+      epochs = 50,
+      weightDecay = 0.0,
+      beta1 = 0.9,
+      beta2 = 0.999,
+      eps = 1e-8,
+      maxGradNorm = 1.0,
+      onEpochEnd = null,
+    } = opts;
+
+    this._initMoments();
+    const losses: number[] = [];
+
+    for (let epoch = 0; epoch < epochs; epoch++) {
+      this.head.zeroGrad();
+      let epochLoss = 0;
+      for (const s of samples) epochLoss += this.head.accumulate(s);
+
+      const grads = this.head.gradients();
+      const invN = 1 / samples.length;
+      for (const g of grads) {
+        for (let i = 0; i < g.data.length; i++) g.data[i]! *= invN;
+      }
+
+      const gradNorm = clipGradients(grads, maxGradNorm);
+
+      this._step++;
+      const beta1_t = Math.pow(beta1, this._step);
+      const beta2_t = Math.pow(beta2, this._step);
+      const hp: AdamHyper = { learningRate, weightDecay, beta1, beta2, eps, beta1_t, beta2_t };
+      if (this.gpuTraining) {
+        await this._adamwStepGpu(grads, hp);
+      } else {
+        this._adamwStepCpu(grads, hp);
+      }
+
+      const avg = epochLoss / samples.length;
+      losses.push(avg);
+      if (onEpochEnd) onEpochEnd(epoch + 1, avg, gradNorm);
+    }
+    return losses;
+  }
+
+  /** Mean loss over samples (no weight update). 0 for an empty set. */
+  evaluate(samples: S[]): number {
+    if (samples.length === 0) return 0;
+    let total = 0;
+    for (const s of samples) total += this.head.loss(s);
+    return total / samples.length;
+  }
+
+  private _adamwStepCpu(grads: HeadParam[], hp: AdamHyper): void {
+    const params = this.head.parameters();
+    const { learningRate: lr, weightDecay: wd, beta1, beta2, eps, beta1_t, beta2_t } = hp;
+    for (let pi = 0; pi < params.length; pi++) {
+      const p = params[pi]!.data;
+      const g = grads[pi]!.data;
+      const mom = this._moments![pi]!;
+      for (let i = 0; i < p.length; i++) {
+        const gi = g[i]!;
+        mom.m[i] = beta1 * mom.m[i]! + (1 - beta1) * gi;
+        mom.v[i] = beta2 * mom.v[i]! + (1 - beta2) * gi * gi;
+        const mHat = mom.m[i]! / (1 - beta1_t);
+        const vHat = mom.v[i]! / (1 - beta2_t);
+        p[i] = p[i]! * (1 - lr * wd) - (lr * mHat) / (Math.sqrt(vHat) + eps);
+      }
+    }
+  }
+
+  /** AdamW on the GPU via the shared WEIGHT_UPDATE_WGSL kernel. Awaited per step. */
+  private async _adamwStepGpu(grads: HeadParam[], hp: AdamHyper): Promise<void> {
+    const device = this.device!;
+    const pipeline = this._adamwPipeline!;
+    const params = this.head.parameters();
+    const { learningRate: lr, weightDecay: wd, beta1, beta2, eps, beta1_t, beta2_t } = hp;
+
+    for (let pi = 0; pi < params.length; pi++) {
+      const p = params[pi]!;
+      const mom = this._moments![pi]!;
+      const paramBuf = createStorageBuffer(device, p.data, true);
+      const gradBuf = createStorageBuffer(device, grads[pi]!.data, false);
+      const mBuf = createStorageBuffer(device, mom.m, true);
+      const vBuf = createStorageBuffer(device, mom.v, true);
+      const uni = createUniformBuffer(
+        device,
+        packAdamParams(p.numel, lr, beta1, beta2, eps, wd, beta1_t, beta2_t),
+      );
+      const bg = createBindGroup(device, pipeline, [uni, paramBuf, gradBuf, mBuf, vBuf]);
+      dispatchKernel(device, pipeline, bg, [cdiv(p.numel, 256), 1, 1]);
+
+      p.data.set((await readBuffer(device, paramBuf, p.numel * 4)).subarray(0, p.numel));
+      mom.m.set((await readBuffer(device, mBuf, p.numel * 4)).subarray(0, p.numel));
+      mom.v.set((await readBuffer(device, vBuf, p.numel * 4)).subarray(0, p.numel));
+
+      paramBuf.destroy();
+      gradBuf.destroy();
+      mBuf.destroy();
+      vBuf.destroy();
+      uni.destroy();
+    }
+  }
+}
+
+/** Clip grads to `maxNorm` in place; returns the PRE-clip L2 norm (the signal). */
+function clipGradients(grads: HeadParam[], maxNorm: number): number {
+  let normSq = 0;
+  for (const g of grads) for (let i = 0; i < g.data.length; i++) normSq += g.data[i]! * g.data[i]!;
+  const norm = Math.sqrt(normSq);
+  if (norm > maxNorm && norm > 0) {
+    const scale = maxNorm / norm;
+    for (const g of grads) for (let i = 0; i < g.data.length; i++) g.data[i]! *= scale;
+  }
+  return norm;
+}
+
+/** Concatenate a head's parameters into one flat vector (the federation unit). */
+export function flattenParams(params: readonly HeadParam[]): Float32Array {
+  const out = new Float32Array(params.reduce((n, p) => n + p.numel, 0));
+  let o = 0;
+  for (const p of params) {
+    out.set(p.data, o);
+    o += p.numel;
+  }
+  return out;
+}
+
+/** Write a flat vector back into a head's parameters (inverse of {@link flattenParams}). */
+export function assignParams(params: readonly HeadParam[], flat: ArrayLike<number>): void {
+  const total = params.reduce((n, p) => n + p.numel, 0);
+  if (flat.length !== total) throw new Error(`assignParams: expected ${total} values, got ${flat.length}`);
+  let o = 0;
+  for (const p of params) {
+    for (let i = 0; i < p.numel; i++) p.data[i] = flat[o + i]!;
+    o += p.numel;
+  }
+}

package/src/ml/index.ts

@@ -0,0 +1,57 @@
+/**
+ * ml — the ML task catalog (docs/ml-task-catalog.md): small, purpose-built models
+ * for bounded questions, answered on the person's own machine without a GPU or an
+ * LLM call. ONE registry, ONE runner, ONE trainer seam; trained heads federate on
+ * the existing delta wire. Pure and portable: no vscode, no DOM, no network.
+ */
+export type {
+  MlTask,
+  MlTaskKind,
+  MlEngineRef,
+  ClassicalAlgorithm,
+  InferenceMode,
+  MlPrediction,
+  MlPredictor,
+  FeatureSpec,
+  MlSample,
+  TrainSpec,
+  OutputSpec,
+} from "./types.js";
+export { engineKey, ML_TASK_ID_RE } from "./types.js";
+
+export { MlTaskRegistry } from "./registry.js";
+export {
+  runMlTask,
+  candidateEngines,
+  InMemoryMlModels,
+  MlDecisionLog,
+  MlEngineUnavailableError,
+} from "./run.js";
+export type { MlRunContext, MlResult, MlModelStore, MlCloudPort, MlDecision, MlDecisionSink, MlSkipReason } from "./run.js";
+
+export { trainMlTask, mlEvalLoss } from "./train.js";
+export type { MlTrainExample, MlTrainOptions, MlTrainResult } from "./train.js";
+
+export { HeadTrainer, flattenParams, assignParams } from "./head_trainer.js";
+export type { HeadParam, TrainableHead, HeadTrainOptions } from "./head_trainer.js";
+export { DenseHead } from "./dense_head.js";
+export type { DenseHeadConfig } from "./dense_head.js";
+
+export { extractFeatures, extractFeatureMatrix, logScale, clampRatio } from "./features.js";
+
+export * from "./classical/index.js";
+
+export {
+  headCheckpoint,
+  headParams,
+  mlHeadEntry,
+  mlHeadParams,
+  mlHeadUnusableReason,
+  packageMlHead,
+  withPackageMlHead,
+  buildMlHeadDelta,
+  mergeMlHeadDeltas,
+} from "./federate.js";
+export type { EvermindMlHead, MlHeadMergeOutcome } from "./federate.js";
+
+export * from "./tasks/index.js";

package/src/ml/registry.ts

@@ -0,0 +1,44 @@
+/**
+ * ml/registry.ts — the ONE registry of ML tasks, keyed by task id.
+ *
+ * A new task is a new row, not a new module. The registry only stores and looks
+ * up rows; running lives in `run.ts`, training in `train.ts`.
+ */
+import { ML_TASK_ID_RE, type MlTask } from "./types.js";
+
+export class MlTaskRegistry {
+  private readonly _tasks = new Map<string, MlTask<any, any>>();
+
+  /** Register a row. Throws on a malformed id, a duplicate id, or a row with no engine. */
+  register<I, O>(task: MlTask<I, O>): MlTask<I, O> {
+    if (!ML_TASK_ID_RE.test(task.id)) throw new Error(`MlTaskRegistry: malformed task id '${task.id}'`);
+    if (this._tasks.has(task.id)) throw new Error(`MlTaskRegistry: task '${task.id}' is already registered`);
+    if (task.engines.length === 0) throw new Error(`MlTaskRegistry: task '${task.id}' lists no engine`);
+    if (task.kind === "classify" && !(task.output.labels && task.output.labels.length >= 2)) {
+      throw new Error(`MlTaskRegistry: classify task '${task.id}' needs at least two output labels`);
+    }
+    this._tasks.set(task.id, task);
+    return task;
+  }
+
+  /** The row for `id`, or undefined. */
+  get<I = unknown, O = unknown>(id: string): MlTask<I, O> | undefined {
+    return this._tasks.get(id) as MlTask<I, O> | undefined;
+  }
+
+  /** The row for `id`; throws when it is not registered. */
+  require<I = unknown, O = unknown>(id: string): MlTask<I, O> {
+    const t = this.get<I, O>(id);
+    if (!t) throw new Error(`MlTaskRegistry: unknown task '${id}'`);
+    return t;
+  }
+
+  has(id: string): boolean {
+    return this._tasks.has(id);
+  }
+
+  /** Every row, in registration order. */
+  list(): MlTask<any, any>[] {
+    return [...this._tasks.values()];
+  }
+}