@seanhogg/builderforce-memory-engine 2026.10.1 → 2026.10.2

package/src/ml/classical/kmeans.ts

@@ -57,11 +57,17 @@ export function fitKMeans(samples: readonly MlSample[], opts: KMeansOptions = {}
     const total = nearest.reduce((a, b) => a + b, 0);
     let pick = 0;
     if (total > 0) {
+      // Walk only points with weight: a zero-distance point (already a centroid) can
+      // never be picked, even when rnd() returns 0 and r starts at 0.
       let r = rnd() * total;
-      for (pick = 0; pick < X.length - 1; pick++) {
+      let last = -1;
+      for (pick = 0; pick < X.length; pick++) {
+        if (!(nearest[pick]! > 0)) continue;
+        last = pick;
         r -= nearest[pick]!;
         if (r <= 0) break;
       }
+      if (pick >= X.length) pick = last; // float round-off left r > 0: take the last weighted point
     } else {
       pick = centroids.length; // every point coincides with a centroid; take the next index
     }
@@ -71,8 +77,8 @@ export function fitKMeans(samples: readonly MlSample[], opts: KMeansOptions = {}
   }
 
   const assign = new Int32Array(X.length).fill(-1);
-  const maxIt = opts.maxIterations ?? 100;
-  for (let it = 0; it < maxIt; it++) {
+  /** Assign every point to its nearest current centroid; true when any assignment moved. */
+  const reassign = (): boolean => {
     let changed = false;
     const model = new KMeansModel(centroids, 0);
     X.forEach((x, i) => {
@@ -82,7 +88,15 @@ export function fitKMeans(samples: readonly MlSample[], opts: KMeansOptions = {}
         changed = true;
       }
     });
-    if (!changed) break;
+    return changed;
+  };
+  const maxIt = opts.maxIterations ?? 100;
+  let converged = false;
+  for (let it = 0; it < maxIt; it++) {
+    if (!reassign()) {
+      converged = true;
+      break;
+    }
     const sums = centroids.map(() => new Float64Array(d));
     const counts = new Array<number>(k).fill(0);
     X.forEach((x, i) => {
@@ -96,6 +110,10 @@ export function fitKMeans(samples: readonly MlSample[], opts: KMeansOptions = {}
     });
   }
 
+  // Stopped on maxIterations: the centroids moved after the last assignment, so
+  // assign once more — inertia (and the stored model) must describe the same state.
+  if (!converged) reassign();
+
   let inertia = 0;
   X.forEach((x, i) => (inertia += squaredDistance(x, centroids[assign[i]!]!)));
   return new KMeansModel(centroids, inertia);

package/src/ml/classical/knn.ts

@@ -11,7 +11,10 @@ import { classCount, cosine, featureDim, squaredDistance, topProbability, type V
 export type KnnMetric = "cosine" | "euclidean";
 
 export interface KnnOptions {
-  /** Neighbours that vote. Default 5 (capped at the sample count). */
+  /**
+   * Neighbours that vote. Default min(5, ⌊√n⌋), at least 1: with a small set a fixed
+   * 5 makes EVERY point vote, and the answer collapses to the class prior.
+   */
   k?: number;
   metric?: KnnMetric;
   /** Number of classes (default: inferred from the labels, at least 2). */
@@ -92,7 +95,8 @@ export class KnnModel implements MlPredictor {
 export function fitKnn(samples: readonly MlSample[], opts: KnnOptions = {}): KnnModel {
   featureDim(samples.map((s) => s.x));
   const classes = opts.classes ?? classCount(samples.map((s) => s.y));
-  const model = new KnnModel(Math.min(opts.k ?? 5, samples.length), opts.metric ?? "cosine", classes);
+  const k = opts.k ?? Math.min(5, Math.max(1, Math.floor(Math.sqrt(samples.length))));
+  const model = new KnnModel(Math.min(k, samples.length), opts.metric ?? "cosine", classes);
   for (const s of samples) model.add(s.x, s.y, s.weight ?? 1);
   return model;
 }

package/src/ml/classical/logistic.ts

@@ -6,6 +6,7 @@
  * features on a stable scale, so a weight means the same thing on every machine —
  * which is what lets trained weights FEDERATE (`exportParams` / `importParams`).
  */
+import { mlEvalLoss } from "../eval_loss.js";
 import type { MlPrediction, MlPredictor, MlSample } from "../types.js";
 import { classCount, featureDim, softmax, topProbability } from "./math.js";
 
@@ -50,16 +51,9 @@ export class LogisticModel implements MlPredictor {
     return { scores, confidence: topProbability(scores) };
   }
 
-  /** Weighted mean cross-entropy over `samples` (the eval-gate loss). */
+  /** Weighted mean cross-entropy over `samples` — exactly the eval-gate loss ({@link mlEvalLoss}); 0 for an empty set. */
   loss(samples: readonly MlSample[]): number {
-    let sum = 0;
-    let wSum = 0;
-    for (const s of samples) {
-      const w = s.weight ?? 1;
-      sum += -w * Math.log(Math.max(this.proba(s.x)[s.y] ?? 0, 1e-12));
-      wSum += w;
-    }
-    return wSum > 0 ? sum / wSum : 0;
+    return mlEvalLoss("classify", this, samples) ?? 0;
   }
 
   /** One full-batch gradient step; returns the pre-step mean loss. */
@@ -115,6 +109,12 @@ export function fitLogistic(
 ): LogisticModel {
   const d = featureDim(samples.map((s) => s.x));
   const k = opts.classes ?? init?.classes ?? classCount(samples.map((s) => s.y));
+  // A label at or above `classes` has no row of weights: its probability would be
+  // read as undefined and every loss and gradient would turn NaN, silently.
+  const top = classCount(samples.map((s) => s.y));
+  if (top > k && samples.some((s) => s.y >= k)) {
+    throw new Error(`fitLogistic: labels go up to ${top - 1} but the model has ${k} classes`);
+  }
   const model = init ?? new LogisticModel(d, k);
   if (model.dim !== d || model.classes !== k) throw new Error(`fitLogistic: init is ${model.dim}×${model.classes}, data is ${d}×${k}`);
   const lr = opts.learningRate ?? 0.5;

package/src/ml/classical/math.ts

@@ -1,93 +1,93 @@
-/**
- * ml/classical/math.ts — the small vector helpers every classical algorithm shares.
- * Pure, allocation-light, deterministic.
- */
-import { SeededRng } from "../../utils/rng.js";
-
-/** The engine's ONE sigmoid / softmax (`training/cpu_ops.ts`). */
-export { sigmoid, softmax } from "../../training/cpu_ops.js";
-
-export type Vec = ArrayLike<number>;
-
-export function dot(a: Vec, b: Vec): number {
-  const n = Math.min(a.length, b.length);
-  let s = 0;
-  for (let i = 0; i < n; i++) s += a[i]! * b[i]!;
-  return s;
-}
-
-export function norm(a: Vec): number {
-  return Math.sqrt(dot(a, a));
-}
-
-/** Cosine similarity in [-1, 1]; 0 when either vector is zero. */
-export function cosine(a: Vec, b: Vec): number {
-  const na = norm(a);
-  const nb = norm(b);
-  return na > 0 && nb > 0 ? dot(a, b) / (na * nb) : 0;
-}
-
-export function squaredDistance(a: Vec, b: Vec): number {
-  const n = Math.min(a.length, b.length);
-  let s = 0;
-  for (let i = 0; i < n; i++) {
-    const d = a[i]! - b[i]!;
-    s += d * d;
-  }
-  return s;
-}
-
-/** Median of a numeric list (copy-sorted). NaN for an empty list. */
-export function median(values: readonly number[]): number {
-  if (values.length === 0) return Number.NaN;
-  const s = [...values].sort((a, b) => a - b);
-  const m = s.length >> 1;
-  return s.length % 2 ? s[m]! : (s[m - 1]! + s[m]!) / 2;
-}
-
-/** Index of the largest value (first on ties). */
-export function argmax(values: ArrayLike<number>): number {
-  let best = 0;
-  for (let i = 1; i < values.length; i++) if (values[i]! > values[best]!) best = i;
-  return best;
-}
-
-/** Number of distinct classes implied by integer labels 0..k-1. */
-export function classCount(labels: readonly number[]): number {
-  let k = 0;
-  for (const y of labels) {
-    if (!Number.isInteger(y) || y < 0) throw new Error(`class label must be a non-negative integer, got ${y}`);
-    k = Math.max(k, y + 1);
-  }
-  return Math.max(k, 2);
-}
-
-/** Feature dimension of a non-empty sample matrix; throws on ragged rows. */
-export function featureDim(X: readonly Vec[]): number {
-  if (X.length === 0) throw new Error("no samples");
-  const d = X[0]!.length;
-  for (const x of X) if (x.length !== d) throw new Error(`ragged features: expected ${d}, got ${x.length}`);
-  return d;
-}
-
-/** Confidence of a probability vector: its top probability. */
-export function topProbability(p: readonly number[]): number {
-  return p.length ? Math.max(...p) : 0;
-}
-
-/** A deterministic [0,1) source for algorithms that sample — the engine's ONE seeded RNG. */
-export function seededUniform(seed: number): () => number {
-  const rng = new SeededRng(seed);
-  return () => rng.next();
-}
-
-/**
- * Average path length of an unsuccessful BST search over n points — the
- * isolation-forest normaliser c(n).
- */
-export function averagePathLength(n: number): number {
-  if (n <= 1) return 0;
-  if (n === 2) return 1;
-  const harmonic = Math.log(n - 1) + 0.5772156649;
-  return 2 * harmonic - (2 * (n - 1)) / n;
-}
+/**
+ * ml/classical/math.ts — the small vector helpers every classical algorithm shares.
+ * Pure, allocation-light, deterministic.
+ */
+import { SeededRng } from "../../utils/rng.js";
+
+/** The engine's ONE sigmoid / softmax (`training/cpu_ops.ts`). */
+export { sigmoid, softmax } from "../../training/cpu_ops.js";
+
+export type Vec = ArrayLike<number>;
+
+export function dot(a: Vec, b: Vec): number {
+  const n = Math.min(a.length, b.length);
+  let s = 0;
+  for (let i = 0; i < n; i++) s += a[i]! * b[i]!;
+  return s;
+}
+
+export function norm(a: Vec): number {
+  return Math.sqrt(dot(a, a));
+}
+
+/** Cosine similarity in [-1, 1]; 0 when either vector is zero. */
+export function cosine(a: Vec, b: Vec): number {
+  const na = norm(a);
+  const nb = norm(b);
+  return na > 0 && nb > 0 ? dot(a, b) / (na * nb) : 0;
+}
+
+export function squaredDistance(a: Vec, b: Vec): number {
+  const n = Math.min(a.length, b.length);
+  let s = 0;
+  for (let i = 0; i < n; i++) {
+    const d = a[i]! - b[i]!;
+    s += d * d;
+  }
+  return s;
+}
+
+/** Median of a numeric list (copy-sorted). NaN for an empty list. */
+export function median(values: readonly number[]): number {
+  if (values.length === 0) return Number.NaN;
+  const s = [...values].sort((a, b) => a - b);
+  const m = s.length >> 1;
+  return s.length % 2 ? s[m]! : (s[m - 1]! + s[m]!) / 2;
+}
+
+/** Index of the largest value (first on ties). */
+export function argmax(values: ArrayLike<number>): number {
+  let best = 0;
+  for (let i = 1; i < values.length; i++) if (values[i]! > values[best]!) best = i;
+  return best;
+}
+
+/** Number of distinct classes implied by integer labels 0..k-1. */
+export function classCount(labels: readonly number[]): number {
+  let k = 0;
+  for (const y of labels) {
+    if (!Number.isInteger(y) || y < 0) throw new Error(`class label must be a non-negative integer, got ${y}`);
+    k = Math.max(k, y + 1);
+  }
+  return Math.max(k, 2);
+}
+
+/** Feature dimension of a non-empty sample matrix; throws on ragged rows. */
+export function featureDim(X: readonly Vec[]): number {
+  if (X.length === 0) throw new Error("no samples");
+  const d = X[0]!.length;
+  for (const x of X) if (x.length !== d) throw new Error(`ragged features: expected ${d}, got ${x.length}`);
+  return d;
+}
+
+/** Confidence of a probability vector: its top probability. */
+export function topProbability(p: readonly number[]): number {
+  return p.length ? Math.max(...p) : 0;
+}
+
+/** A deterministic [0,1) source for algorithms that sample — the engine's ONE seeded RNG. */
+export function seededUniform(seed: number): () => number {
+  const rng = new SeededRng(seed);
+  return () => rng.next();
+}
+
+/**
+ * Average path length of an unsuccessful BST search over n points — the
+ * isolation-forest normaliser c(n).
+ */
+export function averagePathLength(n: number): number {
+  if (n <= 1) return 0;
+  if (n === 2) return 1;
+  const harmonic = Math.log(n - 1) + 0.5772156649;
+  return 2 * harmonic - (2 * (n - 1)) / n;
+}

package/src/ml/classical/naive_bayes.ts

@@ -4,6 +4,10 @@
  * Per class: a weighted prior, and an independent Gaussian per feature. Variance
  * is smoothed by a fraction of the largest feature variance (as scikit-learn's
  * `var_smoothing`), so a constant feature never divides by zero.
+ *
+ * A class with no training weight has prior −∞ and probability exactly 0: its
+ * Gaussian is a placeholder (mean 0, variance ≈ the smoothing floor), which would
+ * otherwise out-score every real class on an input near zero.
  */
 import type { MlPrediction, MlPredictor, MlSample } from "../types.js";
 import { classCount, featureDim, softmax, topProbability } from "./math.js";
@@ -21,11 +25,12 @@ export class NaiveBayesModel implements MlPredictor {
     readonly variance: Float64Array[],
   ) {}
 
-  /** Per-class joint log-likelihood log P(c) + Σ log N(x_i | μ, σ²). */
+  /** Per-class joint log-likelihood log P(c) + Σ log N(x_i | μ, σ²); −∞ for an unseen class. */
   jointLogLikelihood(x: ArrayLike<number>): number[] {
     return this.mean.map((mu, c) => {
       const v = this.variance[c]!;
       let ll = this.logPrior[c]!;
+      if (ll === -Infinity) return ll;
       for (let i = 0; i < mu.length; i++) {
         const d = (x[i] ?? 0) - mu[i]!;
         ll += -0.5 * Math.log(2 * Math.PI * v[i]!) - (d * d) / (2 * v[i]!);
@@ -71,7 +76,9 @@ export function fitNaiveBayes(samples: readonly MlSample[], opts: NaiveBayesOpti
   for (let c = 0; c < k; c++) for (let i = 0; i < d; i++) variance[c]![i] = variance[c]![i]! + eps;
 
   const total = wSum.reduce((a, b) => a + b, 0);
-  // An unseen class keeps a vanishing (not -∞) prior so softmax stays defined.
-  const logPrior = Float64Array.from(wSum, (w) => Math.log(Math.max(w, 1e-12) / total));
+  if (!(total > 0)) throw new Error("fitNaiveBayes: the samples carry no weight");
+  // An unseen class is masked out (−∞, so softmax gives it exactly 0). At least one
+  // class has weight, so softmax always has a finite maximum.
+  const logPrior = Float64Array.from(wSum, (w) => (w > 0 ? Math.log(w / total) : -Infinity));
   return new NaiveBayesModel(logPrior, mean, variance);
 }

package/src/ml/classical/robust_z.ts

@@ -4,7 +4,10 @@
  * Per feature: z = |x − median| / (1.4826 · MAD). Median and MAD are not dragged
  * by the very outliers being looked for, unlike mean and standard deviation. When
  * MAD is 0 (more than half the history identical) the mean absolute deviation
- * stands in, then a floor, so a constant history still scores a change.
+ * stands in, then two floors: a fraction of |median| (relative) and `minScale`
+ * (absolute, in the feature's own units). Without the absolute one a history that
+ * is all zeros — no tool failures, ever — gets a scale of ~1e-9, and the first
+ * small non-zero value scores astronomically.
  *
  * Prediction: `scores[0] = max_i z_i / threshold` — ≥ 1 is anomalous.
  */
@@ -21,6 +24,11 @@ export interface RobustZOptions {
   threshold?: number;
   /** Smallest scale per feature, as a fraction of |median| (plus 1e-9). Default 0.01. */
   minScaleFraction?: number;
+  /**
+   * Smallest scale per feature in absolute feature units. Default 0 — the algorithm
+   * cannot know the units; a task on a fixed scale sets it (`run.anomaly` does).
+   */
+  minScale?: number;
 }
 
 export class RobustZModel implements MlPredictor {
@@ -50,6 +58,7 @@ export function fitRobustZ(samples: readonly MlSample[], opts: RobustZOptions =
   const X = samples.map((s) => s.x);
   const d = featureDim(X);
   const minFrac = opts.minScaleFraction ?? 0.01;
+  const minAbs = opts.minScale ?? 0;
   const center = new Float64Array(d);
   const scale = new Float64Array(d);
   for (let i = 0; i < d; i++) {
@@ -59,7 +68,7 @@ export function fitRobustZ(samples: readonly MlSample[], opts: RobustZOptions =
     let s = MAD_TO_SD * median(dev);
     if (!(s > 0)) s = MEANAD_TO_SD * (dev.reduce((a, b) => a + b, 0) / dev.length);
     center[i] = med;
-    scale[i] = Math.max(s, minFrac * Math.abs(med) + 1e-9);
+    scale[i] = Math.max(s, minFrac * Math.abs(med) + 1e-9, minAbs);
   }
   return new RobustZModel(center, scale, opts.threshold ?? 3.5);
 }

package/src/ml/eval_loss.ts

@@ -0,0 +1,27 @@
+/**
+ * ml/eval_loss.ts — the ONE held-out loss for catalog models.
+ *
+ * What the trainer reports as `holdoutLoss`, what a federated merge is gated on, and
+ * what `LogisticModel.loss` returns. A leaf (types only), so a classical algorithm can
+ * use it without importing the trainer that imports the algorithms.
+ */
+import type { MlPredictor, MlSample, MlTaskKind } from "./types.js";
+
+/**
+ * Mean loss of `predictor` on `samples` for a task of `kind` — classify: weighted
+ * cross-entropy of the true class; regress: weighted ½·squared error. Null for
+ * kinds without a supervised loss, or an empty (or zero-weight) set.
+ */
+export function mlEvalLoss(kind: MlTaskKind, predictor: MlPredictor, samples: readonly MlSample[]): number | null {
+  if (samples.length === 0 || (kind !== "classify" && kind !== "regress")) return null;
+  let sum = 0;
+  let wSum = 0;
+  for (const s of samples) {
+    const w = s.weight ?? 1;
+    const p = predictor.predict(s.x);
+    const loss = kind === "classify" ? -Math.log(Math.max(p.scores[s.y] ?? 0, 1e-12)) : 0.5 * (p.scores[0]! - s.y) ** 2;
+    sum += w * loss;
+    wSum += w;
+  }
+  return wSum > 0 ? sum / wSum : null;
+}

package/src/ml/federate.ts

@@ -8,62 +8,26 @@
  * `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.
+ *
+ * The head's manifest entry and checkpoint format live with the artifact
+ * (`evermind/ml_head.ts`); this module only imports evermind/, never the reverse.
  */
-import { buildDeltaLearnPayload, decodeDeltaPayload, deltaUnusableReason, mlHeadTarget, type DeltaLearnPayload } from "../evermind/delta_wire.js";
+import {
+  buildDeltaLearnPayload,
+  decodeDeltaPayload,
+  deltaTargetMismatch,
+  deltaUnusableReason,
+  mlHeadTarget,
+  type DeltaLearnPayload,
+} from "../evermind/delta_wire.js";
 import { mergeCheckpointDiffsGated, type GatedMergeResult, type MergeGateOptions } from "../evermind/merge.js";
+import { headCheckpoint, headParams, type EvermindMlHead } from "../evermind/ml_head.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;
-}
+export { headCheckpoint, headParams, mlHeadUnusableReason } from "../evermind/ml_head.js";
+export type { EvermindMlHead } from "../evermind/ml_head.js";
 
 /** A manifest entry for a head's current parameters. */
 export function mlHeadEntry(
@@ -121,6 +85,11 @@ export interface MlHeadMergeOutcome {
  * 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.
+ *
+ * Fails CLOSED: `requireEval` defaults to true here, so a round `score` cannot judge
+ * (null — e.g. `mlEvalLoss` for an anomaly/cluster head or an empty holdout — or a
+ * non-finite loss) is refused. A caller that federates an unscorable head on purpose
+ * passes `{ requireEval: false }` explicitly.
  */
 export function mergeMlHeadDeltas(
   base: EvermindMlHead,
@@ -134,7 +103,8 @@ export function mergeMlHeadDeltas(
   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}` });
+    const wrongTarget = deltaTargetMismatch(p.target, target);
+    if (wrongTarget) return void rejected.push({ index, reason: wrongTarget });
     if (p.baseVersion !== base.version) return void rejected.push({ index, reason: `stale base version ${p.baseVersion} (head is ${base.version})` });
     let diff: ArrayBuffer;
     try {
@@ -148,7 +118,7 @@ export function mergeMlHeadDeltas(
     weights.push(p.weight ?? 1);
   });
 
-  const merge = mergeCheckpointDiffsGated(baseCp, diffs, weights, (cp) => score(headParams(cp)), opts);
+  const merge = mergeCheckpointDiffsGated(baseCp, diffs, weights, (cp) => score(headParams(cp)), { ...opts, requireEval: opts.requireEval ?? true });
   if (!merge.accepted) return { head: base, merge, rejected };
   const head: EvermindMlHead = {
     ...base,

package/src/ml/head_trainer.ts

@@ -5,8 +5,9 @@
  *
  * 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).
+ * no device it runs the kernel's CPU mirror, `adamwUpdateInPlace` (`optim/adamw.ts`,
+ * NaN/Inf guard included), so training works everywhere (CI, Node without
+ * @webgpu/node, a laptop with no GPU). Neither path bounds the step (maxDelta 0).
  *
  * The head owns its forward/backward and its loss (MSE for regression, BCE /
  * softmax cross-entropy for classification); this file owns only the loop:
@@ -23,6 +24,7 @@ import {
   cdiv,
 } from "../utils/gpu_utils.js";
 import { WEIGHT_UPDATE_WGSL } from "../kernels/weight_update.js";
+import { adamwUpdateInPlace, type AdamWKernelStep } from "../optim/adamw.js";
 
 /** A named trainable tensor (flat row-major Float32Array). */
 export interface HeadParam {
@@ -63,29 +65,19 @@ interface AdamMoment {
   v: Float32Array;
 }
 
-interface AdamHyper {
-  learningRate: number;
-  weightDecay: number;
-  beta1: number;
-  beta2: number;
-  eps: number;
-  beta1_t: number;
-  beta2_t: number;
-}
+/** Byte size of WEIGHT_UPDATE_WGSL's `AdamParams` uniform: 1 u32 + 8 f32 + 3 f32 pad = 48. */
+const ADAM_PARAMS_BYTES = 48;
 
-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);
+/**
+ * Pack `AdamParams` exactly as WEIGHT_UPDATE_WGSL declares it:
+ * `num_elements: u32, lr, beta1, beta2, eps, weight_decay, beta1_t, beta2_t,
+ * max_delta, _pad0, _pad1, _pad2` (all f32 after the first) — 48 bytes.
+ */
+function packAdamParams(numElements: number, hp: AdamWKernelStep): ArrayBuffer {
+  const buf = new ArrayBuffer(ADAM_PARAMS_BYTES);
   new Uint32Array(buf, 0, 1).set([numElements]);
-  new Float32Array(buf, 4, 7).set([lr, beta1, beta2, eps, weightDecay, beta1_t, beta2_t]);
+  // Pads stay 0 (fresh ArrayBuffer).
+  new Float32Array(buf, 4, 8).set([hp.lr, hp.beta1, hp.beta2, hp.eps, hp.weightDecay, hp.beta1_t, hp.beta2_t, hp.maxDelta]);
   return buf;
 }
 
@@ -159,7 +151,7 @@ export class HeadTrainer<S> {
       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 };
+      const hp: AdamWKernelStep = { lr: learningRate, weightDecay, beta1, beta2, eps, beta1_t, beta2_t, maxDelta: 0 };
       if (this.gpuTraining) {
         await this._adamwStepGpu(grads, hp);
       } else {
@@ -181,30 +173,20 @@ export class HeadTrainer<S> {
     return total / samples.length;
   }
 
-  private _adamwStepCpu(grads: HeadParam[], hp: AdamHyper): void {
+  /** AdamW on the CPU via the kernel's shared mirror (`adamwUpdateInPlace`). */
+  private _adamwStepCpu(grads: HeadParam[], hp: AdamWKernelStep): 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);
-      }
+      adamwUpdateInPlace(params[pi]!.data, grads[pi]!.data, mom.m, mom.v, hp);
     }
   }
 
   /** AdamW on the GPU via the shared WEIGHT_UPDATE_WGSL kernel. Awaited per step. */
-  private async _adamwStepGpu(grads: HeadParam[], hp: AdamHyper): Promise<void> {
+  private async _adamwStepGpu(grads: HeadParam[], hp: AdamWKernelStep): 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]!;
@@ -213,10 +195,7 @@ export class HeadTrainer<S> {
       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 uni = createUniformBuffer(device, packAdamParams(p.numel, hp));
       const bg = createBindGroup(device, pipeline, [uni, paramBuf, gradBuf, mBuf, vBuf]);
       dispatchKernel(device, pipeline, bg, [cdiv(p.numel, 256), 1, 1]);