decidekit 0.1.1 → 0.2.0

package/CHANGELOG.md

@@ -1,5 +1,14 @@
 # Changelog
 
+## 0.2.0 (unreleased)
+
+### Training
+
+- From 8 examples per label (mean across labels), `train()` fits the head with full-batch Adam, starting from the same centroid and label-name point. Test accuracy with all data: BANKING77 90.6% to 93.1%, CLINC150 94.6% to 96.1%, HWU64 90.15% to 91.17% (one seed, on cached embeddings; through train() in eval/results-h14.json, BANKING77 93.1% and CLINC150 96.3%). At 5 per label nothing changes. At 10 per label: BANKING77 84.4% to 87.5%, CLINC150 92.2% to 92.5%, HWU64 85.50% to 85.38%. Adam's settings come from jimothy's `src/linear.ts` (MIT, Andrew Prifer).
+- New option `solver`. Setting `learningRate` or `l2` without it keeps SGD, so tuned calls behave as before. The default `epochs` is 100 with Adam.
+- Saved heads keep working unchanged. Retrained heads can score differently, so pick `minConfidence` thresholds again: the CLINC150 out-of-scope threshold moved from 0.9 to 0.98.
+- Training takes longer: in Node, BANKING77 went from 84 s to 123 s on cached embeddings. In Chrome, BANKING77 took 208 s with a longest page stall of 19 ms.
+
 ## 0.1.1 (2026-11-02)
 
 ### Models

package/README.md

@@ -16,8 +16,8 @@ const r = await d.decide("my card never arrived", ["card_arrival", "refund", "ot
 r.choice; // "card_arrival" | "refund" | "other"
 ```
 
-Give it 5 examples per label and it reaches 82% on BANKING77 (77 intents); give it the full training set and it
-reaches 90.7%, with each answer taking about 8 ms (median) in Chrome. Training runs in the user's tab, so their data never
+Give it 5 examples per label and it reaches 84% on BANKING77 (77 intents); give it the full training set and it
+reaches 93.1%, with each answer taking about 8 ms (median) in Chrome. Training runs in the user's tab, so their data never
 leaves the device.
 
 ## Zero-shot
@@ -167,6 +167,15 @@ Decide Small (bge-small-en-v1.5, the default) at q8, without label names, Ryzen
 and calibration error (ECE, lower is better) on each dataset's test split, mean of 3 seeds (one run for "all"):
 [eval/results-bge-small.json](eval/results-bge-small.json).
 
+| Examples per label | 5 | 10 | 25 | all |
+|---|---|---|---|---|
+| BANKING77 (77 labels) | 83.7% | 87.5% | 90.2% | 93.1% |
+| CLINC150 (150 labels) | 90.2% | 92.8% | 94.3% | 96.3% |
+| ECE, BANKING77 | 0.050 | 0.015 | 0.013 | 0.013 |
+| ECE, CLINC150 | 0.020 | 0.015 | 0.016 | 0.008 |
+
+Those are 0.2, with label names, from [eval/results-h14.json](eval/results-h14.json). HWU64 with all data goes from 90.15% to 91.17%. Before 0.2 (SGD at every size, no label names, `eval/results-bge-small.json`):
+
 | Examples per label | 5 | 10 | 25 | all |
 |---|---|---|---|---|
 | BANKING77 (77 labels) | 82.3% | 83.8% | 88.7% | 90.7% |
@@ -189,7 +198,7 @@ queries, [eval/tiny.ts](eval/tiny.ts), measured with all-MiniLM-L6-v2), 1 exampl
 a threshold there.
 
 **In Chrome** (BANKING77, full training set, worker, cross-origin isolated, 16 threads): 7.7 ms median and 13.8 ms
-p95 per decision, 228 s to train on 10,003 examples with no page stall longer than 18 ms. all-MiniLM-L6-v2 in the
+p95 per decision before 0.2, 228 s to train on 10,003 examples with no page stall longer than 18 ms. A 0.2 run on the same machine took 208 s, with a longest stall of 19 ms and 93.05% test accuracy (runs on this laptop vary: another 0.1 run took 147 s). all-MiniLM-L6-v2 in the
 same run: 4.2 ms median, 7.7 ms p95. In Node: 3.2 ms median for a single text, 1.1 ms per text batched (MiniLM 2.0
 and 0.7). Earlier MiniLM run: [bench/results.json](bench/results.json).
 

package/dist/head.js

@@ -1,5 +1,10 @@
 const BATCH = 64;
 const CENTROID_SCALE = 20;
+// At 5 per label Adam lost 0.29 points on CLINC150, at 10 it gained on BANKING77 and CLINC150 (eval/spike). The cutoff
+// between them wasn't measured on its own.
+const ADAM_FROM = 8;
+// Without a better holdout loss for this many epochs, stop the probe run.
+const ADAM_PATIENCE = 30;
 // Lets a tab paint and handle input between training batches. A message-channel task, unlike setTimeout, isn't
 // throttled to once a second (or minute) when the tab is in the background.
 function yieldToPage() {
@@ -56,11 +61,23 @@ function batchGradient(ids, X, y, W, B, gW, gb, z, C, dim) {
         }
     }
 }
+// One Adam step from the summed gradient. Adam's settings follow jimothy's src/linear.ts (MIT, Andrew Prifer).
+function adamStep(p, g, m, v, rows, l2, lr, epoch) {
+    const c1 = 1 - 0.9 ** epoch, c2 = 1 - 0.999 ** epoch;
+    for (let k = 0; k < p.length; k++) {
+        const d = g[k] / rows + l2 * p[k];
+        m[k] = 0.9 * m[k] + 0.1 * d;
+        v[k] = 0.999 * v[k] + 0.001 * d * d;
+        p[k] -= (lr * m[k]) / c1 / (Math.sqrt(v[k] / c2) + 1e-8);
+    }
+}
 /** Fits weights, biases and temperature to embeddings `X` with label indices `y`. */
 export async function fitHead(X, y, labels, opts = {}, priors, priorWeight = 1, centroidScale = CENTROID_SCALE, smallFloor = 1, calibration) {
     // smallFloor: the lowest temperature a head with under 50 held-out answers may get; createDecider picks it per model.
     // calibration: labelled examples known to be right to fit the temperature on instead of the training examples.
-    const { epochs = 60, learningRate = 2, l2 = 1e-4, holdout = 0.1, seed = 7, onEpoch } = opts;
+    // A set learningRate or l2 was tuned for SGD, so it keeps SGD unless the solver is named.
+    const adam = opts.solver ? opts.solver === "adam" : opts.learningRate === undefined && opts.l2 === undefined && X.length / labels.length >= ADAM_FROM;
+    const { epochs = adam ? 100 : 60, learningRate = adam ? 0.05 : 2, l2 = adam ? 0 : 1e-4, holdout = 0.1, seed = 7, onEpoch } = opts;
     const dim = X[0].length, C = labels.length;
     // One flat buffer, so the inner loops read typed memory instead of an array of arrays.
     const Xf = new Float64Array(X.length * dim);
@@ -92,6 +109,33 @@ export async function fitHead(X, y, labels, opts = {}, priors, priorWeight = 1,
         let best = { acc: check ? acc() : 0, weights: head.weights.slice(), bias: head.bias.slice(), epochs: 0 };
         const gW = new Float64Array(C * dim), gb = new Float64Array(C), z = new Float64Array(C);
         const W = head.weights, B = head.bias;
+        if (adam) {
+            const mW = new Float64Array(C * dim), vW = new Float64Array(C * dim), mB = new Float64Array(C), vB = new Float64Array(C);
+            let bestLoss = Infinity;
+            for (let ep = 0; ep < nEpochs; ep++) {
+                gW.fill(0);
+                gb.fill(0);
+                for (let b = 0; b < rows.length; b += BATCH) {
+                    batchGradient(rows.slice(b, b + BATCH), Xf, y, W, B, gW, gb, z, C, dim);
+                    await yieldToPage();
+                }
+                adamStep(W, gW, mW, vW, rows.length, l2, learningRate, ep + 1);
+                adamStep(B, gb, mB, vB, rows.length, 0, learningRate, ep + 1);
+                if (!check)
+                    continue;
+                const loss = check.reduce((sum, i) => sum - Math.log(softmax(logits(head, X[i]), 1)[y[i]] + 1e-300), 0);
+                if (onEpoch)
+                    onEpoch(ep, acc());
+                if (loss < bestLoss * (1 - 1e-7)) {
+                    bestLoss = loss;
+                    best = { acc: 0, weights: W.slice(), bias: B.slice(), epochs: ep + 1 };
+                }
+                if (ep + 1 - best.epochs >= ADAM_PATIENCE)
+                    break;
+                await yieldToPage();
+            }
+            return check ? { ...head, weights: best.weights, bias: best.bias, epochs: best.epochs } : { ...head, epochs: nEpochs };
+        }
         const order = [...rows];
         for (let ep = 0; ep < nEpochs; ep++) {
             order.sort(() => rand() - 0.5);

package/dist/index.d.ts

@@ -51,11 +51,16 @@ export interface TrainOptions<L extends string = string> {
      * can still be chosen. Needs `labelNames`.
      */
     labels?: readonly L[];
-    /** Passes over the examples. Default 60; the best holdout epoch is kept. */
+    /**
+     * How to fit the head. By default `adam` from 8 examples per label (mean across labels), `sgd` below that. Setting
+     * `learningRate` or `l2` without this also picks `sgd`.
+     */
+    solver?: "adam" | "sgd";
+    /** Most passes over the examples. Default 100 with `adam`, 60 with `sgd`; the best holdout epoch is kept. */
     epochs?: number;
-    /** Default 2. */
+    /** Default 0.05 with `adam`, 2 with `sgd`. */
     learningRate?: number;
-    /** Weight decay. Default 1e-4. */
+    /** Weight decay. Default 0 with `adam`, 1e-4 with `sgd`. */
     l2?: number;
     /** Share of examples held out for early stopping and temperature fitting. Default 0.1. */
     holdout?: number;
@@ -74,7 +79,7 @@ export interface DeciderOptions {
     /**
      * A Hugging Face sentence-embedding model with ONNX weights. Default `debuffjp/decide-small` (Decide Small, 34 MB at q8).
      * Before 0.1.1 the default id was `Xenova/bge-small-en-v1.5`, the same weights; that id and heads recording it still work.
-     * Pre-release builds defaulted to `Xenova/all-MiniLM-L6-v2` (23 MB): smaller and 1.6x faster, 3 to 5 points less
+     * Pre-release builds defaulted to `Xenova/all-MiniLM-L6-v2` (23.7 MB): smaller and about 1.8x faster in Chrome, 3 to 5 points less
      * accurate. Pass it to keep using heads trained on it.
      */
     model?: string;

package/dist/index.js

@@ -103,8 +103,8 @@ export async function createDecider({ model: requested = SMALL, device, dtype =
     /**
      * Trains a head on labelled examples: mini-batch gradient descent on the embeddings, keeping the best holdout
      * epoch, then a temperature fitted on the holdout so `confidence` is calibrated. Every label needs at least a few
-     * examples. Each label starts at the mean of its examples and its name, so 5 per label already gave 79% on BANKING77 and
-     * 87% on CLINC150 (see the README).
+     * examples. Each label starts at the mean of its examples and its name, so 5 per label already gave 82.3% on BANKING77 and
+     * 89.6% on CLINC150 (see the README).
      */
     async function train(examples, opts = {}) {
         const bad = examples.findIndex((e) => typeof e?.text !== "string" || typeof e?.label !== "string" || !e.label);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "decidekit",
-  "version": "0.1.1",
+  "version": "0.2.0",
   "description": "A System One model in the browser: typed, calibrated text decisions on the user's device. Learns from an LLM when it isn't sure.",
   "license": "MIT",
   "repository": "github:debuffjp/decidekit",