widgets/softmax.js

/* Softmax tester — the detailed recipe attention and classifiers use: subtract the max (stability),
 * exponentiate each logit, sum, then divide each by the sum → a probability distribution, animated into
 * bars. Interconnected: if the paper feeds in upstream scores (io.consumes), it softmaxes the first row
 * of that score matrix and publishes the resulting weights downstream. Same math the video narrates. */
(function () {
  const k = MANIMO.kit, bus = MANIMO.bus;
  function mount(root, concept, theme) {
    const p = (concept && concept.params) || {};
    const io = (concept && concept.io) || {};
    const up = io.consumes ? bus.get(io.consumes) : null;
    const v0 = (up && up.length ? up[0].slice() : (p.logits || [2.0, 1.0, 0.1, -0.5]));
    const grid = k.grid(1, v0.length, [v0]);
    const cv = document.createElement("canvas"); cv.width = 460; cv.height = 180; cv.className = "plot";
    cv.setAttribute("role", "img"); cv.setAttribute("aria-label", "Softmax probability distribution bar chart");
    const log = k.panel("steps");
    const note = up ? k.el("div", { class: "matlbl" }, ["row 0 of S = Q·Kᵀ/√d (the scaled scores)"]) : null;

    function calc() {
      const v = grid.get()[0], mx = Math.max.apply(null, v);
      const ex = v.map((x) => Math.exp(x - mx));
      const sum = ex.reduce((a, b) => a + b, 0);
      const prob = ex.map((e) => e / sum);
      return { v, mx, ex, sum, prob };
    }
    function publish() { if (io.produces) bus.set(io.produces, [calc().prob]); }

    // steps: 1 = max, 2..1+n = exp each element, 2+n = sum, 3+n..2+2n = divide each element.
    const n = v0.length, S_MAX = 1, S_EXP = 1, S_SUM = 1 + n + 1, total = 2 * n + 2;
    function render(step) {
      const { v, mx, ex, sum, prob } = calc();
      k.clear(log); log.appendChild(k.el("div", { class: "steps-h" }, ["softmax  (" + step + " / " + total + ")"]));
      if (step >= S_MAX) k.logLine(log, "1. subtract max = " + k.fmt(mx) + "  (numerical stability)", step === S_MAX);
      const expShown = k.clamp(step - S_EXP, 0, n);
      for (let i = 0; i < expShown; i++) k.logLine(log, "2." + (i + 1) + " e^(" + k.fmt(v[i]) + "−" + k.fmt(mx) + ") = " + k.fmt(ex[i]), i === expShown - 1 && step < S_SUM);
      if (step >= S_SUM) k.logLine(log, "3. sum of exponentials = " + k.fmt(sum), step === S_SUM);
      const divShown = k.clamp(step - S_SUM, 0, n);
      for (let i = 0; i < divShown; i++) k.logLine(log, "4." + (i + 1) + " p[" + i + "] = " + k.fmt(ex[i]) + " / " + k.fmt(sum) + " = <b>" + k.fmt(prob[i]) + "</b>", i === divShown - 1);
      // bars: logits → (exp, partially) → probabilities
      const done = divShown >= n;
      bars(done ? prob : (expShown > 0 ? norm(ex) : v.map(() => 0)), done);
    }
    function bars(vals, asProb) {
      const ctx = cv.getContext("2d"), W = cv.width, H = cv.height, n2 = vals.length, bw = (W - 20) / n2;
      ctx.fillStyle = theme.bg; ctx.fillRect(0, 0, W, H);
      const mx = Math.max.apply(null, vals.concat([1e-6]));
      vals.forEach((val, i) => {
        const h = (val / mx) * (H - 36), x = 10 + i * bw;
        ctx.fillStyle = i === argmax(vals) && asProb ? theme.accent2 : theme.accent;
        ctx.fillRect(x + 4, H - 22 - h, bw - 8, h);
        ctx.fillStyle = theme.ink; ctx.font = "12px system-ui"; ctx.textAlign = "center";
        ctx.fillText(asProb ? k.fmt(val) : "", x + bw / 2, H - 26 - h);
        ctx.fillText("x" + i, x + bw / 2, H - 6);
      });
    }
    function norm(a) { const m = Math.max.apply(null, a.concat([1e-6])); return a.map((e) => e / m); }
    function argmax(a) { let m = 0; a.forEach((val, i) => { if (val > a[m]) m = i; }); return m; }
    const ctrl = k.stepper(new Array(total), render);

    function rebuild() { ctrl.reset(); publish(); }
    grid.node.addEventListener("input", rebuild);
    if (io.consumes) bus.on(io.consumes, (S) => { if (S && S.length) { grid.set([S[0]]); rebuild(); } });
    publish();

    root.appendChild(k.row([labeled("logits", grid.node), note].filter(Boolean), "mats"));
    root.appendChild(cv);
    root.appendChild(k.row([k.button("▸ Step", () => ctrl.next()), k.button("▶ Play", () => ctrl.play(880)), k.button("↺ Reset", () => ctrl.reset(), "ghost")], "ctrls"));
    root.appendChild(log);
    return k.player(ctrl, 780);
  }
  function labeled(t, node) { return MANIMO.kit.el("div", { class: "matcol" }, [MANIMO.kit.el("div", { class: "matlbl" }, [t]), node]); }
  MANIMO.registry["softmax"] = { title: "Softmax", mount };
})();