widgets/attention.js

/* Scaled dot-product attention tester — the whole pipeline, step by detailed step: raw scores S = Q·Kᵀ,
 * scale by 1/√d, softmax each row into weights, then the weighted sum ·V — one row at a time, with the
 * score and weight matrices shown as heatmaps. Q=K=V=X (self-attention). Interconnected: it reads the
 * SAME token matrix X the rest of the paper flows through (io.consumes) and publishes its output, so
 * editing X upstream re-derives the attention output here. Same pipeline the Manim `attention` walks. */
(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 toks = (p.tokens || ["the", "cat", "sat"]).slice(0, 5);
    const X0 = (up && up.length) ? up.map((r) => r.slice())
      : (p.X || toks.map((_, i) => toks.map((__, j) => (i === j ? 1 : 0))));
    const d = X0.length, f = (X0[0] || []).length, sc = 1 / Math.sqrt(f || 1);
    const X = k.grid(d, f, X0, { readonly: !!io.consumes });   // upstream-driven ⇒ read-only here
    const log = k.panel("steps");
    const sCanvas = mkCanvas(), wCanvas = mkCanvas();
    sCanvas.setAttribute("aria-label", "Attention score heatmap, S = Q·Kᵀ/√d");
    wCanvas.setAttribute("aria-label", "Attention weight heatmap, softmax of scores per row");

    function compute() {
      const x = X.get();
      const raw = x.map((qi) => x.map((kj) => k.dot(qi, kj)));
      const scores = raw.map((r) => r.map((s) => s * sc));
      const weights = scores.map((r) => k.softmax(r));
      const out = weights.map((w) => (x[0] || []).map((_, j) => k.dot(w, k.col(x, j))));
      return { raw, scores, weights, out };
    }
    function publish() { if (io.produces) bus.set(io.produces, compute().out); }

    const total = 2 + 2 * d;                                    // raw, scale, d softmax rows, d output rows
    function render(step) {
      const { raw, scores, weights, out } = compute();
      k.clear(log); log.appendChild(k.el("div", { class: "steps-h" }, ["self-attention  (" + step + " / " + total + ")"]));
      k.heatmap(sCanvas, step >= 2 ? scores : (step >= 1 ? raw : zeros(d)), theme, { values: step >= 1 });
      if (step >= 1) k.logLine(log, "1. raw scores S = Q·Kᵀ — every token scored against every token", step === 1);
      if (step >= 2) k.logLine(log, "2. scale by 1/√" + f + " = " + k.fmt(sc) + "  (stabilises gradients)", step === 2);
      const wRows = k.clamp(step - 2, 0, d);
      k.heatmap(wCanvas, weights.map((r, i) => (i < wRows ? r : Array(d).fill(0))), theme, { values: wRows > 0 });
      for (let i = 0; i < wRows; i++) k.logLine(log, "3." + (i + 1) + " softmax(row " + i + ") → weights [" + weights[i].map(k.fmt).join(", ") + "] (sum 1)", i === wRows - 1 && step <= 2 + d);
      const oRows = k.clamp(step - 2 - d, 0, d);
      for (let i = 0; i < oRows; i++) k.logLine(log, "4." + (i + 1) + " out[" + i + "] = Σ wᵢ·V = [" + out[i].map(k.fmt).join(", ") + "]", i === oRows - 1);
    }
    const ctrl = k.stepper(new Array(total), render);

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

    root.appendChild(k.row([labeled("tokens × features (Q=K=V)", X.node)], "mats"));
    root.appendChild(k.row([labeled("S = Q·Kᵀ / √d", sCanvas), labeled("softmax(S) → weights", wCanvas)], "heat"));
    root.appendChild(k.row([k.button("▸ Step", () => ctrl.next()), k.button("▶ Play", () => ctrl.play(1200)), k.button("↺ Reset", () => ctrl.reset(), "ghost")], "ctrls"));
    root.appendChild(log);
    return k.player(ctrl, 1150);
  }
  function mkCanvas() { const c = document.createElement("canvas"); c.width = 156; c.height = 156; c.className = "heat"; c.setAttribute("role", "img"); return c; }
  function labeled(t, node) { return MANIMO.kit.el("div", { class: "matcol" }, [MANIMO.kit.el("div", { class: "matlbl" }, [t]), node]); }
  function zeros(n) { return Array.from({ length: n }, () => Array(n).fill(0)); }
  MANIMO.registry["attention"] = { title: "Attention", mount };
})();