@seanhogg/builderforce-memory-engine 2026.6.39 → 2026.7.11

package/src/codec/video_rvq.ts

@@ -0,0 +1,496 @@
+/**
+ * video_rvq.ts — VideoRVQCodec: the pixels ⇄ tokens bottleneck that lets Evermind
+ * generate video without changing the generator.
+ *
+ * Evermind is a discrete-token autoregressive SSM. To make it a video model you
+ * do NOT touch the model — you give it a codec that maps a clip to a stream of
+ * discrete tokens and back. This is the visual analogue of the BPE tokenizer:
+ * where BPE turns text ⇄ ids, VideoRVQCodec turns frames ⇄ ids.
+ *
+ * Design — TEMPORAL residual vector quantization (the "video-from-the-start"
+ * choice):
+ *
+ *   • Each frame is cut into non-overlapping p×p patches; each patch is a vector
+ *     of length p·p·C (the "latent").
+ *   • KEYFRAMES (frame 0 and every `keyframeInterval`-th frame) are quantized
+ *     directly against the INTRA codebook bank.
+ *   • Every other frame is INTER: we quantize the *residual against the previous
+ *     RECONSTRUCTED frame* against the INTER bank. Video is mostly temporally
+ *     redundant, so these deltas are small and cheap — this is exactly why an
+ *     SSM (linear-time in sequence length) is the right generator for the long
+ *     token streams video produces.
+ *   • Quantization at each bank is RESIDUAL VQ over `levels` codebooks: level 0
+ *     picks the nearest entry, level 1 quantizes what level 0 missed, and so on.
+ *     More levels ⇒ finer reconstruction ⇒ more tokens per patch.
+ *
+ * The encoder runs closed-loop (it references its own reconstruction, never the
+ * ground-truth previous frame) so the decoder — which only ever has
+ * reconstructions — stays exactly in sync.
+ *
+ * Codebooks start random (a cold codec is lossy, like any untrained neural
+ * codec) and are learned by {@link VideoRVQCodec.fit} (greedy per-level k-means).
+ * Reaching photoreal fidelity needs training on a real corpus with real compute
+ * — that is the one genuine blocker on shipping generated video, not the wiring.
+ *
+ * Pure CPU, zero deps, deterministic under a seed — same conventions as the rest
+ * of the engine. A WGSL/WebGPU acceleration is a future drop-in with these shapes.
+ */
+
+import { SeededRng } from "../utils/rng.js";
+import { MultimodalVocab, VIDEO_BANK_INTRA, VIDEO_BANK_INTER } from "./multimodal_vocab.js";
+
+/** A single frame: length `height·width·channels`, layout `((y·W)+x)·C + ch`, values in [0,1]. */
+export type Frame = Float32Array;
+/** A clip: T frames, all the same shape. */
+export type Video = Frame[];
+
+export interface VideoRVQConfig {
+  height: number;
+  width: number;
+  /** Colour channels. Default 3 (RGB). */
+  channels?: number;
+  /** Square patch size; `height` and `width` must be divisible by it. Default 4. */
+  patch?: number;
+  /** Residual-VQ depth (codes per patch). Default 2. */
+  levels?: number;
+  /** Codebook entries per level per bank. Default 16. */
+  codebookSize?: number;
+  /** Emit a keyframe every N frames (frame 0 is always a keyframe). Default 12. */
+  keyframeInterval?: number;
+  /** Size of the text region this codec's tokens sit above; 0 ⇒ pure-video vocab. Default 0. */
+  textVocabSize?: number;
+  /** Deterministic seed for codebook init. */
+  seed?: number;
+}
+
+const VRQ_MAGIC = 0x56525130; // "VRQ0"
+
+const DEFAULTS = {
+  channels: 3,
+  patch: 4,
+  levels: 2,
+  codebookSize: 16,
+  keyframeInterval: 12,
+  textVocabSize: 0,
+  seed: 0x56524651, // "VRFQ"
+} as const;
+
+export class VideoRVQCodec {
+  readonly height: number;
+  readonly width: number;
+  readonly channels: number;
+  readonly patch: number;
+  readonly levels: number;
+  readonly codebookSize: number;
+  readonly keyframeInterval: number;
+
+  /** Patch latent dimension = patch·patch·channels. */
+  readonly latentDim: number;
+  /** Patches per frame = (H/patch)·(W/patch). */
+  readonly patchesPerFrame: number;
+  /** Code tokens emitted per frame = patchesPerFrame·levels. */
+  readonly tokensPerFrame: number;
+
+  /** The unified text+video vocabulary; feed `.vocab.size` to EvermindLM. */
+  readonly vocab: MultimodalVocab;
+
+  /** Two banks × levels codebooks; each codebook is `codebookSize × latentDim` row-major. */
+  private readonly banks: [Float32Array[], Float32Array[]];
+
+  constructor(config: VideoRVQConfig) {
+    const channels = config.channels ?? DEFAULTS.channels;
+    const patch = config.patch ?? DEFAULTS.patch;
+    if (config.height <= 0 || config.width <= 0) throw new Error("VideoRVQCodec: height/width must be > 0");
+    if (config.height % patch !== 0 || config.width % patch !== 0) {
+      throw new Error(`VideoRVQCodec: height/width must be divisible by patch (${patch})`);
+    }
+    this.height = config.height;
+    this.width = config.width;
+    this.channels = channels;
+    this.patch = patch;
+    this.levels = config.levels ?? DEFAULTS.levels;
+    this.codebookSize = config.codebookSize ?? DEFAULTS.codebookSize;
+    this.keyframeInterval = Math.max(1, config.keyframeInterval ?? DEFAULTS.keyframeInterval);
+    this.latentDim = patch * patch * channels;
+    this.patchesPerFrame = (this.height / patch) * (this.width / patch);
+    this.tokensPerFrame = this.patchesPerFrame * this.levels;
+
+    this.vocab = new MultimodalVocab({
+      textVocabSize: config.textVocabSize ?? DEFAULTS.textVocabSize,
+      levels: this.levels,
+      codebookSize: this.codebookSize,
+    });
+
+    const rng = new SeededRng((config.seed ?? DEFAULTS.seed) >>> 0 || 1);
+    const mkBank = (): Float32Array[] =>
+      Array.from({ length: this.levels }, () => {
+        const cb = new Float32Array(this.codebookSize * this.latentDim);
+        // Small-magnitude init: intra entries hover near mid-grey, inter (delta) near 0.
+        for (let i = 0; i < cb.length; i++) cb[i] = (rng.next() - 0.5) * 0.1;
+        return cb;
+      });
+    this.banks = [mkBank(), mkBank()];
+    // Bias the intra bank toward the [0,1] pixel range so a cold codec is not black.
+    for (const cb of this.banks[VIDEO_BANK_INTRA]) for (let i = 0; i < cb.length; i++) cb[i] += 0.5;
+  }
+
+  /** Total vocabulary size for this codec's model. */
+  get vocabSize(): number {
+    return this.vocab.size;
+  }
+
+  // ── Serialize (so a trained codec ships alongside the model artifact) ───────────
+
+  /**
+   * Serialize config + learned codebooks to a compact "VRQ0" binary. A generated
+   * video model is only servable if its codec travels with it (the decoder needs
+   * these codebooks), so this is the video analogue of EvermindLM.exportWeights.
+   */
+  serialize(): ArrayBuffer {
+    const headerEls = 9; // magic, h, w, channels, patch, levels, codebookSize, keyframeInterval, textVocabSize
+    const total = 2 * this.levels * this.codebookSize * this.latentDim;
+    const buf = new ArrayBuffer(headerEls * 4 + total * 4);
+    const head = new Uint32Array(buf, 0, headerEls);
+    head[0] = VRQ_MAGIC;
+    head[1] = this.height;
+    head[2] = this.width;
+    head[3] = this.channels;
+    head[4] = this.patch;
+    head[5] = this.levels;
+    head[6] = this.codebookSize;
+    head[7] = this.keyframeInterval;
+    head[8] = this.vocab.textVocabSize;
+    const body = new Float32Array(buf, headerEls * 4, total);
+    let o = 0;
+    for (const bank of this.banks) {
+      for (const cb of bank) {
+        body.set(cb, o);
+        o += cb.length;
+      }
+    }
+    return buf;
+  }
+
+  /** Reconstruct a codec (config + codebooks) from a "VRQ0" binary. */
+  static deserialize(buffer: ArrayBuffer): VideoRVQCodec {
+    const head = new Uint32Array(buffer, 0, 9);
+    if (head[0] !== VRQ_MAGIC) throw new Error("VideoRVQCodec.deserialize: bad magic (not a VRQ0 blob)");
+    const codec = new VideoRVQCodec({
+      height: head[1]!,
+      width: head[2]!,
+      channels: head[3]!,
+      patch: head[4]!,
+      levels: head[5]!,
+      codebookSize: head[6]!,
+      keyframeInterval: head[7]!,
+      textVocabSize: head[8]!,
+    });
+    const total = 2 * codec.levels * codec.codebookSize * codec.latentDim;
+    const body = new Float32Array(buffer.slice(36, 36 + total * 4));
+    let o = 0;
+    for (const bank of codec.banks) {
+      for (const cb of bank) {
+        cb.set(body.subarray(o, o + cb.length));
+        o += cb.length;
+      }
+    }
+    return codec;
+  }
+
+  // ── Encode ───────────────────────────────────────────────────────────────────
+
+  /** Encode a clip to a self-delimiting token stream: `<vid> (marker codes…)… </vid>`. */
+  encode(video: Video): number[] {
+    this.assertShape(video);
+    const out: number[] = [this.vocab.bosVideo];
+    let prevRecon: Float32Array[] | null = null;
+
+    for (let t = 0; t < video.length; t++) {
+      const patches = this.toPatches(video[t]!);
+      const isKey = prevRecon === null || t % this.keyframeInterval === 0;
+      const bankId = isKey ? VIDEO_BANK_INTRA : VIDEO_BANK_INTER;
+      out.push(isKey ? this.vocab.frameKey : this.vocab.frameDelta);
+
+      const reconPatches: Float32Array[] = [];
+      for (let p = 0; p < this.patchesPerFrame; p++) {
+        const target = isKey ? patches[p]! : sub(patches[p]!, prevRecon![p]!);
+        const { codes, recon } = this.rvqEncode(bankId, target);
+        for (let l = 0; l < this.levels; l++) out.push(this.vocab.codeToken(bankId, l, codes[l]!));
+        reconPatches.push(isKey ? recon : add(prevRecon![p]!, recon));
+      }
+      prevRecon = reconPatches;
+    }
+
+    out.push(this.vocab.eosVideo);
+    return out;
+  }
+
+  // ── Decode ───────────────────────────────────────────────────────────────────
+
+  /**
+   * Decode a token stream back to frames. Tolerant by design: it skips text /
+   * stray tokens, treats a frame whose codes run short as zero-padded, and stops
+   * at `</vid>` or end-of-stream — so a stream sampled from an under-trained
+   * generator still yields a valid (if noisy) clip instead of throwing.
+   */
+  decode(tokens: number[]): Video {
+    const frames: Video = [];
+    let prevRecon: Float32Array[] | null = null;
+    let i = 0;
+
+    // Advance to the first frame marker (tolerate a leading prompt / <vid>).
+    while (i < tokens.length && !this.vocab.isVideoMarker(tokens[i]!)) {
+      if (tokens[i] === this.vocab.eosVideo) return frames;
+      i++;
+    }
+
+    while (i < tokens.length) {
+      const marker = tokens[i]!;
+      if (!this.vocab.isVideoMarker(marker)) break; // </vid>, text, or garbage → end
+      i++;
+      const isKey = marker === this.vocab.frameKey || prevRecon === null;
+      const bankId = isKey ? VIDEO_BANK_INTRA : VIDEO_BANK_INTER;
+
+      // Pull up to tokensPerFrame code tokens; stop early on the next marker/EOS.
+      const codes: number[] = [];
+      while (codes.length < this.tokensPerFrame && i < tokens.length && this.vocab.isCode(tokens[i]!)) {
+        codes.push(tokens[i]!);
+        i++;
+      }
+
+      const reconPatches: Float32Array[] = [];
+      for (let p = 0; p < this.patchesPerFrame; p++) {
+        const recon = new Float32Array(this.latentDim);
+        for (let l = 0; l < this.levels; l++) {
+          const slot = p * this.levels + l;
+          const tok = slot < codes.length ? codes[slot]! : -1;
+          if (tok < 0) continue; // missing → contribute zero (graceful degradation)
+          const { level, code } = this.vocab.decodeCode(tok);
+          addRowInto(recon, this.banks[bankId]![level]!, code, this.latentDim);
+        }
+        reconPatches.push(isKey ? recon : add(prevRecon![p]!, recon));
+      }
+
+      frames.push(this.fromPatches(reconPatches));
+      prevRecon = reconPatches;
+    }
+
+    return frames;
+  }
+
+  // ── Learn codebooks ────────────────────────────────────────────────────────────
+
+  /**
+   * Learn both codebook banks from a set of clips (greedy per-level k-means over
+   * residuals — the standard way to train residual VQ). Returns the mean
+   * reconstruction MSE over the training clips after fitting. This is the codec's
+   * "training"; the generator is trained separately on the resulting token streams.
+   */
+  fit(videos: Video[], opts: { iterations?: number; seed?: number } = {}): number {
+    const iterations = opts.iterations ?? 8;
+    const rng = new SeededRng((opts.seed ?? 0x66697421) >>> 0 || 1); // "fit!"
+    for (const v of videos) this.assertShape(v);
+
+    // Gather intra patches and inter deltas (open-loop deltas are a fine proxy for fitting).
+    const intra: Float32Array[] = [];
+    const inter: Float32Array[] = [];
+    for (const video of videos) {
+      let prev: Float32Array[] | null = null;
+      for (let t = 0; t < video.length; t++) {
+        const patches = this.toPatches(video[t]!);
+        const isKey = prev === null || t % this.keyframeInterval === 0;
+        if (isKey) for (const p of patches) intra.push(p);
+        else for (let p = 0; p < patches.length; p++) inter.push(sub(patches[p]!, prev![p]!));
+        prev = patches;
+      }
+    }
+
+    this.fitBank(VIDEO_BANK_INTRA, intra, iterations, rng);
+    this.fitBank(VIDEO_BANK_INTER, inter, iterations, rng);
+
+    // Report closed-loop reconstruction error over the corpus.
+    let se = 0;
+    let n = 0;
+    for (const video of videos) {
+      const recon = this.decode(this.encode(video));
+      for (let t = 0; t < video.length; t++) {
+        const a = video[t]!;
+        const b = recon[t]!;
+        for (let k = 0; k < a.length; k++) {
+          const d = a[k]! - b[k]!;
+          se += d * d;
+        }
+        n += a.length;
+      }
+    }
+    return n > 0 ? se / n : 0;
+  }
+
+  /** Fit one bank's `levels` codebooks greedily over residuals. */
+  private fitBank(bankId: number, vectors: Float32Array[], iterations: number, rng: SeededRng): void {
+    if (vectors.length === 0) return;
+    let residuals = vectors.map((v) => Float32Array.from(v));
+    for (let l = 0; l < this.levels; l++) {
+      const centroids = kmeans(residuals, this.codebookSize, this.latentDim, iterations, rng);
+      this.banks[bankId]![l] = centroids;
+      residuals = residuals.map((r) => {
+        const idx = nearestRow(centroids, r, this.codebookSize, this.latentDim);
+        const out = Float32Array.from(r);
+        const off = idx * this.latentDim;
+        for (let d = 0; d < this.latentDim; d++) out[d] = out[d]! - centroids[off + d]!;
+        return out;
+      });
+    }
+  }
+
+  // ── Residual VQ over one bank ───────────────────────────────────────────────────
+
+  private rvqEncode(bankId: number, vector: Float32Array): { codes: number[]; recon: Float32Array } {
+    const residual = Float32Array.from(vector);
+    const recon = new Float32Array(this.latentDim);
+    const codes: number[] = [];
+    for (let l = 0; l < this.levels; l++) {
+      const cb = this.banks[bankId]![l]!;
+      const idx = nearestRow(cb, residual, this.codebookSize, this.latentDim);
+      codes.push(idx);
+      const off = idx * this.latentDim;
+      for (let d = 0; d < this.latentDim; d++) {
+        recon[d] = recon[d]! + cb[off + d]!;
+        residual[d] = residual[d]! - cb[off + d]!;
+      }
+    }
+    return { codes, recon };
+  }
+
+  // ── Patch <-> frame ─────────────────────────────────────────────────────────────
+
+  private toPatches(frame: Frame): Float32Array[] {
+    const { patch, channels, width } = this;
+    const pw = this.width / patch;
+    const ph = this.height / patch;
+    const out: Float32Array[] = [];
+    for (let py = 0; py < ph; py++) {
+      for (let px = 0; px < pw; px++) {
+        const vec = new Float32Array(this.latentDim);
+        for (let dy = 0; dy < patch; dy++) {
+          for (let dx = 0; dx < patch; dx++) {
+            const y = py * patch + dy;
+            const x = px * patch + dx;
+            const src = ((y * width) + x) * channels;
+            const dst = ((dy * patch) + dx) * channels;
+            for (let c = 0; c < channels; c++) vec[dst + c] = frame[src + c]!;
+          }
+        }
+        out.push(vec);
+      }
+    }
+    return out;
+  }
+
+  private fromPatches(patches: Float32Array[]): Frame {
+    const { patch, channels, width } = this;
+    const pw = this.width / patch;
+    const frame = new Float32Array(this.height * this.width * channels);
+    for (let p = 0; p < patches.length; p++) {
+      const py = Math.floor(p / pw);
+      const px = p - py * pw;
+      const vec = patches[p]!;
+      for (let dy = 0; dy < patch; dy++) {
+        for (let dx = 0; dx < patch; dx++) {
+          const y = py * patch + dy;
+          const x = px * patch + dx;
+          const dst = ((y * width) + x) * channels;
+          const src = ((dy * patch) + dx) * channels;
+          for (let c = 0; c < channels; c++) frame[dst + c] = clamp01(vec[src + c]!);
+        }
+      }
+    }
+    return frame;
+  }
+
+  private assertShape(video: Video): void {
+    const expect = this.height * this.width * this.channels;
+    for (let t = 0; t < video.length; t++) {
+      if (video[t]!.length !== expect) {
+        throw new Error(`VideoRVQCodec: frame ${t} has length ${video[t]!.length}, expected ${expect}`);
+      }
+    }
+  }
+}
+
+// ── free helpers ─────────────────────────────────────────────────────────────────
+
+function sub(a: Float32Array, b: Float32Array): Float32Array {
+  const out = new Float32Array(a.length);
+  for (let i = 0; i < a.length; i++) out[i] = a[i]! - b[i]!;
+  return out;
+}
+
+function add(a: Float32Array, b: Float32Array): Float32Array {
+  const out = new Float32Array(a.length);
+  for (let i = 0; i < a.length; i++) out[i] = a[i]! + b[i]!;
+  return out;
+}
+
+function addRowInto(acc: Float32Array, codebook: Float32Array, row: number, dim: number): void {
+  const off = row * dim;
+  for (let d = 0; d < dim; d++) acc[d] = acc[d]! + codebook[off + d]!;
+}
+
+function clamp01(v: number): number {
+  return v < 0 ? 0 : v > 1 ? 1 : v;
+}
+
+/** Nearest codebook row to `v` by squared L2. `codebook` is `count × dim` row-major. */
+function nearestRow(codebook: Float32Array, v: Float32Array, count: number, dim: number): number {
+  let best = 0;
+  let bestD = Infinity;
+  for (let r = 0; r < count; r++) {
+    const off = r * dim;
+    let d = 0;
+    for (let k = 0; k < dim; k++) {
+      const diff = v[k]! - codebook[off + k]!;
+      d += diff * diff;
+      if (d >= bestD) break;
+    }
+    if (d < bestD) {
+      bestD = d;
+      best = r;
+    }
+  }
+  return best;
+}
+
+/** Lloyd's k-means → `count × dim` centroids. Empty clusters reseed to a random sample. */
+function kmeans(vectors: Float32Array[], count: number, dim: number, iterations: number, rng: SeededRng): Float32Array {
+  const centroids = new Float32Array(count * dim);
+  // Init: distinct random samples.
+  for (let r = 0; r < count; r++) {
+    const src = vectors[Math.floor(rng.next() * vectors.length) % vectors.length]!;
+    centroids.set(src, r * dim);
+  }
+  const sums = new Float32Array(count * dim);
+  const counts = new Int32Array(count);
+  for (let it = 0; it < iterations; it++) {
+    sums.fill(0);
+    counts.fill(0);
+    for (const v of vectors) {
+      const r = nearestRow(centroids, v, count, dim);
+      counts[r] = counts[r]! + 1;
+      const off = r * dim;
+      for (let d = 0; d < dim; d++) sums[off + d] = sums[off + d]! + v[d]!;
+    }
+    for (let r = 0; r < count; r++) {
+      const off = r * dim;
+      if (counts[r]! === 0) {
+        // Reseed the dead centroid onto a random sample so it can win points next round.
+        const src = vectors[Math.floor(rng.next() * vectors.length) % vectors.length]!;
+        centroids.set(src, off);
+        continue;
+      }
+      const inv = 1 / counts[r]!;
+      for (let d = 0; d < dim; d++) centroids[off + d] = sums[off + d]! * inv;
+    }
+  }
+  return centroids;
+}

package/src/index.ts

@@ -33,6 +33,8 @@ export type {
     MoEEpochResult,
     EvermindModelManifest,
     EvermindModelCard,
+    EvermindModelType,
+    EvermindModality,
     PackageMeta,
     ValidationResult,
 } from './moe/index.js';
@@ -62,6 +64,26 @@ export {
 export { BPETokenizer } from './tokenizer/bpe.js';
 export type { BPEEncodeOptions, PadSide, HuggingFaceTokenizerSpec, SpecialTokenOverrides } from './tokenizer/bpe.js';
 
+// ── Modality codecs (media ⇄ tokens; lets EvermindLM generate video) ────────────
+
+export {
+    MultimodalVocab,
+    VIDEO_BANK_INTRA,
+    VIDEO_BANK_INTER,
+    VideoRVQCodec,
+    ImageRVQCodec,
+    buildVideoSequence,
+    generateVideo,
+    generateImage,
+} from './codec/index.js';
+export type {
+    MultimodalVocabConfig,
+    TokenKind,
+    VideoRVQConfig,
+    Frame,
+    Video,
+} from './codec/index.js';
+
 // ── Model export (the publishing step: ONNX / safetensors / GGUF / HF repo) ─────
 
 export {

package/src/kernels/selective_scan.ts

@@ -60,12 +60,15 @@ fn softplus(x: f32) -> f32 {
 //   B_bar = (A_bar - 1) / A * B  ≈  Δ * B  (first-order for simplicity)
 fn discretise_A(delta_val: f32, a_log: f32) -> f32 {
     // A is stored as -exp(a_log) to ensure A_bar < 1 (stable)
-    let a_cont = -exp(a_log);
+    // Clamp log-decay so -exp(a_log) can't overflow to -Inf (A_bar→0, state
+    // death) nor collapse toward 0 (A_bar→1, no decay). Keeps A_bar strictly in
+    // (0,1) across repeated adapts. Belt-and-suspenders: WSLA also freezes A_log.
+    let a_cont = -exp(clamp(a_log, -10.0, 5.0));
     return exp(delta_val * a_cont);
 }
 
 fn discretise_B(delta_val: f32, a_log: f32, b_val: f32) -> f32 {
-    let a_cont  = -exp(a_log);
+    let a_cont  = -exp(clamp(a_log, -10.0, 5.0));
     let a_bar   = exp(delta_val * a_cont);
     // (A_bar - 1) / A_cont * B
     let b_bar   = (a_bar - 1.0) / a_cont * b_val;
@@ -264,7 +267,10 @@ fn softplus_grad(x: f32) -> f32 {
 }
 
 fn discretise_A(delta_val: f32, a_log: f32) -> f32 {
-    let a_cont = -exp(a_log);
+    // Clamp log-decay so -exp(a_log) can't overflow to -Inf (A_bar→0, state
+    // death) nor collapse toward 0 (A_bar→1, no decay). Keeps A_bar strictly in
+    // (0,1) across repeated adapts. Belt-and-suspenders: WSLA also freezes A_log.
+    let a_cont = -exp(clamp(a_log, -10.0, 5.0));
     return exp(delta_val * a_cont);
 }
 

package/src/kernels/ssd.ts

@@ -50,7 +50,10 @@ struct SsdParams {
 @group(0) @binding(9) var<storage, read_write> state_carry : array<f32>; // [n_chunks+1,B,H,N,d_head]
 
 fn softplus(x: f32) -> f32 {
-    return log(1.0 + exp(x));
+    // Numerically stable: max(x,0) + log1p(exp(-|x|)). The naive log(1+exp(x))
+    // overflows to +Inf for x ≳ 88 (f32 exp range) — and a downstream
+    // exp(-Inf*0) then yields NaN, poisoning the state. This form never overflows.
+    return max(x, 0.0) + log(1.0 + exp(-abs(x)));
 }
 
 // Workgroup: one chunk × one head × one batch item
@@ -186,7 +189,10 @@ struct SsdParams {
 @group(0) @binding(14) var<storage, read_write> dD_vec     : array<f32>;
 
 fn softplus(x: f32) -> f32 {
-    return log(1.0 + exp(x));
+    // Numerically stable: max(x,0) + log1p(exp(-|x|)). The naive log(1+exp(x))
+    // overflows to +Inf for x ≳ 88 (f32 exp range) — and a downstream
+    // exp(-Inf*0) then yields NaN, poisoning the state. This form never overflows.
+    return max(x, 0.0) + log(1.0 + exp(-abs(x)));
 }
 fn d_softplus(x: f32) -> f32 {
     return 1.0 / (1.0 + exp(-x));

package/src/kernels/weight_update.ts

@@ -19,6 +19,10 @@ struct AdamParams {
     weight_decay   : f32,   // default 0.01
     beta1_t        : f32,   // beta1^t  (precomputed bias correction term)
     beta2_t        : f32,   // beta2^t
+    max_delta      : f32,   // trust region: max |Δθ| per step (0 ⇒ unbounded)
+    _pad0          : f32,   // pad struct to a 16-byte multiple (uniform layout)
+    _pad1          : f32,
+    _pad2          : f32,
 };
 
 @group(0) @binding(0) var<uniform>             adam     : AdamParams;
@@ -52,9 +56,24 @@ fn adamw_update(
     let m_hat = m_new / (1.0 - adam.beta1_t);
     let v_hat = v_new / (1.0 - adam.beta2_t);
 
-    // Weight decay (decoupled) + gradient step
-    param[i] = p * (1.0 - adam.lr * adam.weight_decay) -
-               adam.lr * m_hat / (sqrt(v_hat) + adam.eps);
+    // Adam step.
+    var step = adam.lr * m_hat / (sqrt(v_hat) + adam.eps);
+
+    // Numerical guard: never write a non-finite step into a weight. A NaN or Inf
+    // here (from a bad gradient, a zero v_hat, an overflow) would permanently
+    // poison the parameter and every future forward — the "model dies" failure.
+    // NaN fails self-comparison; treat ±Inf-magnitude as non-finite too.
+    if (step != step || step > 3.4e38 || step < -3.4e38) { step = 0.0; }
+
+    // Trust region: bound how far ONE step can move a weight. With write-through
+    // adaptation running repeatedly, an unbounded step (even from a noisy
+    // gradient) compounds across executions and blows the weights up; clamping
+    // the per-element delta keeps every adapt small and reversible. max_delta==0
+    // disables the bound (full-training callers that don't set it).
+    if (adam.max_delta > 0.0) { step = clamp(step, -adam.max_delta, adam.max_delta); }
+
+    // Weight decay (decoupled) + bounded gradient step
+    param[i] = p * (1.0 - adam.lr * adam.weight_decay) - step;
 }
 `;
 

package/src/model/mamba_model.ts

@@ -413,6 +413,33 @@ export class HybridMambaModel {
         return params;
     }
 
+    /**
+     * The parameters a `learn()` / `adapt()` step is allowed to update.
+     *
+     * Under WSLA (write-through / narrow adaptation — the Evermind "update ==
+     * replace" path) this is ONLY the per-layer narrow subset (the input
+     * projections that produce Δ/B/C). The backbone — token embedding, final
+     * norm, and every block's `A_log` / conv / output projection — stays FROZEN.
+     * Freezing `A_log` is the key stability guarantee: the state-transition
+     * decay can never drift into a degenerate regime (state death / no-decay)
+     * across repeated adapts, which is what made the model "die after several
+     * executions". Outside WSLA every parameter trains (full fine-tune / distill).
+     *
+     * Param names are namespaced `layer{i}.{name}` to match {@link parameters},
+     * so a name-keyed optimizer (see MambaTrainer) reuses the same Adam moments
+     * whether it's stepping the full set or just the WSLA subset.
+     */
+    getTrainableParams(): LayerParam[] {
+        if (!this._wslaMode) return this.parameters();
+        const params: LayerParam[] = [];
+        for (let i = 0; i < this.layers.length; i++) {
+            for (const p of this.layers[i]!.getTrainableParams()) {
+                params.push({ ...p, name: `layer${i}.${p.name}` });
+            }
+        }
+        return params;
+    }
+
     setWSLAMode(enabled: boolean): void {
         for (const layer of this.layers) layer.setWSLAMode(enabled);
         this._wslaMode = enabled;

package/src/moe/index.ts

@@ -18,6 +18,7 @@ export type {
   EvermindModelManifest,
   EvermindModelCard,
   EvermindModelType,
+  EvermindModality,
   PackageMeta,
   ValidationResult,
 } from "./moe_package.js";