@seanhogg/builderforce-memory-engine 2026.10.1 → 2026.10.3

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,21 @@ 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: one value for every
+   * feature, or one per feature (a missing entry is 0). Default 0 — the algorithm
+   * cannot know the units; a task on a fixed scale sets it (`run.anomaly` does,
+   * per feature, because each of its metrics is normalised by a different reference).
+   */
+  minScale?: MinScale;
+}
+
+/** An absolute scale floor: one for every feature, or one per feature. */
+export type MinScale = number | readonly number[];
+
+/** The floor for feature `i` (a per-feature array's missing entry is 0). */
+export function minScaleAt(minScale: MinScale, i: number): number {
+  return typeof minScale === "number" ? minScale : (minScale[i] ?? 0);
 }
 
 export class RobustZModel implements MlPredictor {
@@ -50,6 +68,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: MinScale = opts.minScale ?? 0;
   const center = new Float64Array(d);
   const scale = new Float64Array(d);
   for (let i = 0; i < d; i++) {
@@ -59,7 +78,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, minScaleAt(minAbs, i));
   }
   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,13 @@
  *
  * 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).
+ *
+ * A non-finite gradient element is skipped by both paths — its moments and weight
+ * are left as they were, so the next finite step trains it again — and counted in
+ * {@link HeadTrainer.skippedNonFinite}, so a caller can see a run that diverged.
  *
  * 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 +28,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 +69,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;
 }
 
@@ -94,6 +90,7 @@ export class HeadTrainer<S> {
   readonly device: GPUDevice | null;
   private _moments: AdamMoment[] | null = null;
   private _step = 0;
+  private _skippedNonFinite = 0;
   private readonly _adamwPipeline: GPUComputePipeline | null;
 
   constructor(
@@ -109,6 +106,19 @@ export class HeadTrainer<S> {
       : null;
   }
 
+  /**
+   * Parameter elements whose gradient was non-finite (NaN/±Inf) and whose update was
+   * therefore skipped, summed over every step this trainer has taken. 0 on a healthy run.
+   */
+  get skippedNonFinite(): number {
+    return this._skippedNonFinite;
+  }
+
+  /** Whether any step had to skip a non-finite gradient element. */
+  get diverged(): boolean {
+    return this._skippedNonFinite > 0;
+  }
+
   /** Whether the optimiser step runs on the GPU. */
   get gpuTraining(): boolean {
     return this.device != null && this._adamwPipeline != null;
@@ -159,7 +169,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,42 +191,31 @@ 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`), which counts its skips. */
+  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);
-      }
+      this._skippedNonFinite += 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]!;
       const mom = this._moments![pi]!;
+      // The kernel skips these elements too; it cannot count, so count them here.
+      this._skippedNonFinite += countNonFinite(grads[pi]!.data);
       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 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]);
 
@@ -233,16 +232,35 @@ export class HeadTrainer<S> {
   }
 }
 
-/** Clip grads to `maxNorm` in place; returns the PRE-clip L2 norm (the signal). */
+/** Elements of `a` that are NaN or ±Inf. */
+function countNonFinite(a: Float32Array): number {
+  let n = 0;
+  for (let i = 0; i < a.length; i++) if (!Number.isFinite(a[i]!)) n++;
+  return n;
+}
+
+/**
+ * Clip grads to `maxNorm` in place; returns the PRE-clip L2 norm (the signal — NaN
+ * or Inf when any element is). The clip itself uses the norm of the FINITE elements:
+ * those are skipped by the optimiser anyway, and one NaN must not exempt the rest
+ * from clipping.
+ */
 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;
+  let finiteSq = 0;
+  for (const g of grads) {
+    for (let i = 0; i < g.data.length; i++) {
+      const sq = g.data[i]! * g.data[i]!;
+      normSq += sq;
+      if (Number.isFinite(sq)) finiteSq += sq;
+    }
+  }
+  const finiteNorm = Math.sqrt(finiteSq);
+  if (finiteNorm > maxNorm && finiteNorm > 0) {
+    const scale = maxNorm / finiteNorm;
     for (const g of grads) for (let i = 0; i < g.data.length; i++) g.data[i]! *= scale;
   }
-  return norm;
+  return Math.sqrt(normSq);
 }
 
 /** Concatenate a head's parameters into one flat vector (the federation unit). */