@dniskav/neuron 0.1.6 → 0.2.0

package/dist/index.mjs

@@ -475,6 +475,565 @@ var NetworkLSTM = class {
   }
 };
 
+// src/MatMul.ts
+function matMul(A, B) {
+  const rows = A.length;
+  const inner = B.length;
+  const cols = B[0].length;
+  const C = Array.from({ length: rows }, () => new Array(cols).fill(0));
+  for (let i = 0; i < rows; i++)
+    for (let k = 0; k < inner; k++) {
+      const aik = A[i][k];
+      for (let j = 0; j < cols; j++)
+        C[i][j] += aik * B[k][j];
+    }
+  return C;
+}
+function transpose(A) {
+  const rows = A.length, cols = A[0].length;
+  const T = Array.from({ length: cols }, () => new Array(rows).fill(0));
+  for (let i = 0; i < rows; i++)
+    for (let j = 0; j < cols; j++)
+      T[j][i] = A[i][j];
+  return T;
+}
+function softmax(row) {
+  const max = Math.max(...row);
+  const exps = row.map((v) => Math.exp(v - max));
+  const sum = exps.reduce((a, b) => a + b, 0);
+  return exps.map((e) => e / sum);
+}
+function softmaxBackward(dS, s) {
+  const dot = s.reduce((acc, si, i) => acc + dS[i] * si, 0);
+  return s.map((si, i) => si * (dS[i] - dot));
+}
+var WeightMatrix = class {
+  constructor(rows, cols) {
+    const limit = Math.sqrt(2 / (rows + cols));
+    this.W = Array.from(
+      { length: rows },
+      () => Array.from({ length: cols }, () => (Math.random() * 2 - 1) * limit)
+    );
+    this.opts = Array.from(
+      { length: rows },
+      () => Array.from({ length: cols }, () => new Adam())
+    );
+  }
+  // Apply pre-computed gradient (same shape as W).
+  update(dW, lr) {
+    for (let i = 0; i < this.W.length; i++)
+      for (let j = 0; j < this.W[0].length; j++)
+        this.W[i][j] = this.opts[i][j].step(this.W[i][j], dW[i][j], lr);
+  }
+};
+var EmbeddingMatrix = class {
+  constructor(vocabSize, d_model) {
+    const limit = Math.sqrt(1 / d_model);
+    this.W = Array.from(
+      { length: vocabSize },
+      () => Array.from({ length: d_model }, () => (Math.random() * 2 - 1) * limit)
+    );
+  }
+  get(idx) {
+    return [...this.W[idx]];
+  }
+  update(idx, grad, lr) {
+    for (let m = 0; m < this.W[idx].length; m++)
+      this.W[idx][m] += lr * grad[m];
+  }
+};
+
+// src/AttentionHead.ts
+var AttentionHead = class {
+  constructor(d_model, d_k, d_v) {
+    // d_v × d_model
+    this.cache = null;
+    this.d_k = d_k;
+    this.d_v = d_v;
+    this.Wq = new WeightMatrix(d_k, d_model);
+    this.Wk = new WeightMatrix(d_k, d_model);
+    this.Wv = new WeightMatrix(d_v, d_model);
+  }
+  // ── Forward ───────────────────────────────────────────────────────────────
+  // X: seqLen × d_model  →  out: seqLen × d_v
+  predict(X) {
+    const seqLen = X.length;
+    const scale = 1 / Math.sqrt(this.d_k);
+    const Q = X.map(
+      (x) => this.Wq.W.map((wq) => wq.reduce((s, w, m) => s + w * x[m], 0))
+    );
+    const K = X.map(
+      (x) => this.Wk.W.map((wk) => wk.reduce((s, w, m) => s + w * x[m], 0))
+    );
+    const V = X.map(
+      (x) => this.Wv.W.map((wv) => wv.reduce((s, w, m) => s + w * x[m], 0))
+    );
+    const scores = Array.from(
+      { length: seqLen },
+      (_, i) => Array.from(
+        { length: seqLen },
+        (_2, j) => Q[i].reduce((s, q, k) => s + q * K[j][k], 0) * scale
+      )
+    );
+    const attn = scores.map((row) => softmax(row));
+    const out = Array.from(
+      { length: seqLen },
+      (_, i) => Array.from(
+        { length: this.d_v },
+        (_2, d) => attn[i].reduce((s, a, j) => s + a * V[j][d], 0)
+      )
+    );
+    this.cache = { X, Q, K, V, scores, attn };
+    return out;
+  }
+  // ── Backward ──────────────────────────────────────────────────────────────
+  // dOut: seqLen × d_v  →  dX: seqLen × d_model
+  //
+  // Steps:
+  //   1. dV    = attn^T @ dOut
+  //   2. dAttn = dOut @ V^T            (attention weight gradients)
+  //   3. dScores = softmaxBwd(dAttn) / √d_k
+  //   4. dQ   = dScores @ K,   dK = dScores^T @ Q
+  //   5. dWq  = dQ^T @ X,      dWk = dK^T @ X,  dWv = dV^T @ X
+  //   6. dX   = dQ @ Wq  +  dK @ Wk  +  dV @ Wv
+  backward(dOut, lr) {
+    const { X, Q, K, V, attn } = this.cache;
+    const seqLen = X.length;
+    const d_model = X[0].length;
+    const scale = 1 / Math.sqrt(this.d_k);
+    const dV = Array.from(
+      { length: seqLen },
+      (_, j) => Array.from(
+        { length: this.d_v },
+        (_2, d) => attn.reduce((s, a, i) => s + a[j] * dOut[i][d], 0)
+      )
+    );
+    const dAttn = Array.from(
+      { length: seqLen },
+      (_, i) => Array.from(
+        { length: seqLen },
+        (_2, j) => dOut[i].reduce((s, d, k) => s + d * V[j][k], 0)
+      )
+    );
+    const dScores = dAttn.map(
+      (da, i) => softmaxBackward(da, attn[i]).map((v) => v * scale)
+    );
+    const dQ = Array.from(
+      { length: seqLen },
+      (_, i) => Array.from(
+        { length: this.d_k },
+        (_2, k) => dScores[i].reduce((s, ds, j) => s + ds * K[j][k], 0)
+      )
+    );
+    const dK = Array.from(
+      { length: seqLen },
+      (_, j) => Array.from(
+        { length: this.d_k },
+        (_2, k) => dScores.reduce((s, ds, i) => s + ds[j] * Q[i][k], 0)
+      )
+    );
+    const dWq = Array.from(
+      { length: this.d_k },
+      (_, k) => Array.from(
+        { length: d_model },
+        (_2, m) => dQ.reduce((s, dq, i) => s + dq[k] * X[i][m], 0)
+      )
+    );
+    const dWk = Array.from(
+      { length: this.d_k },
+      (_, k) => Array.from(
+        { length: d_model },
+        (_2, m) => dK.reduce((s, dk, i) => s + dk[k] * X[i][m], 0)
+      )
+    );
+    const dWv = Array.from(
+      { length: this.d_v },
+      (_, k) => Array.from(
+        { length: d_model },
+        (_2, m) => dV.reduce((s, dv, i) => s + dv[k] * X[i][m], 0)
+      )
+    );
+    this.Wq.update(dWq, lr);
+    this.Wk.update(dWk, lr);
+    this.Wv.update(dWv, lr);
+    const dX = Array.from(
+      { length: seqLen },
+      (_, i) => Array.from(
+        { length: d_model },
+        (_2, m) => dQ[i].reduce((s, dq, k) => s + dq * this.Wq.W[k][m], 0) + dK[i].reduce((s, dk, k) => s + dk * this.Wk.W[k][m], 0) + dV[i].reduce((s, dv, k) => s + dv * this.Wv.W[k][m], 0)
+      )
+    );
+    return dX;
+  }
+  // Attention weights from the last predict() call — useful for visualization.
+  getAttentionWeights() {
+    return this.cache ? this.cache.attn : null;
+  }
+};
+
+// src/MultiHeadAttention.ts
+var MultiHeadAttention = class {
+  // seqLen × (nHeads * d_k)
+  constructor(d_model, nHeads) {
+    // d_model × (nHeads * d_k)
+    // Cached for backward
+    this._concat = null;
+    this.nHeads = nHeads;
+    this.d_model = d_model;
+    this.d_k = Math.floor(d_model / nHeads);
+    this.heads = Array.from(
+      { length: nHeads },
+      () => new AttentionHead(d_model, this.d_k, this.d_k)
+    );
+    this.Wo = new WeightMatrix(d_model, nHeads * this.d_k);
+  }
+  // ── Forward ───────────────────────────────────────────────────────────────
+  // X: seqLen × d_model  →  out: seqLen × d_model
+  predict(X) {
+    const seqLen = X.length;
+    const headOuts = this.heads.map((h) => h.predict(X));
+    const concat = Array.from(
+      { length: seqLen },
+      (_, i) => headOuts.flatMap((ho) => ho[i])
+    );
+    const out = concat.map(
+      (c) => this.Wo.W.map((row) => row.reduce((s, w, j) => s + w * c[j], 0))
+    );
+    this._concat = concat;
+    return out;
+  }
+  // ── Backward ──────────────────────────────────────────────────────────────
+  // dOut: seqLen × d_model  →  dX: seqLen × d_model
+  backward(dOut, lr) {
+    const seqLen = dOut.length;
+    const concatD = this.nHeads * this.d_k;
+    const d_model = this.d_model;
+    const concat = this._concat;
+    const dConcat = dOut.map(
+      (do_) => Array.from(
+        { length: concatD },
+        (_, j) => this.Wo.W.reduce((s, row, k) => s + do_[k] * row[j], 0)
+      )
+    );
+    const dWo = Array.from(
+      { length: d_model },
+      (_, k) => Array.from(
+        { length: concatD },
+        (_2, j) => dOut.reduce((s, row, i) => s + row[k] * concat[i][j], 0)
+      )
+    );
+    this.Wo.update(dWo, lr);
+    const dX = Array.from(
+      { length: seqLen },
+      () => new Array(d_model).fill(0)
+    );
+    for (let h = 0; h < this.nHeads; h++) {
+      const start = h * this.d_k;
+      const dHeadOut = dConcat.map((dc) => dc.slice(start, start + this.d_k));
+      const dXh = this.heads[h].backward(dHeadOut, lr);
+      for (let i = 0; i < seqLen; i++)
+        for (let m = 0; m < d_model; m++)
+          dX[i][m] += dXh[i][m];
+    }
+    return dX;
+  }
+  // Attention weights per head from the last predict() — for visualization.
+  // Returns: nHeads × seqLen × seqLen
+  getAttentionWeights() {
+    return this.heads.map((h) => h.getAttentionWeights());
+  }
+};
+
+// src/LayerNorm.ts
+var LayerNorm = class {
+  constructor(dim) {
+    this.eps = 1e-5;
+    // Per-position cache populated during the forward pass.
+    // resetCache() must be called before each sequence forward pass.
+    this._cache = [];
+    this.gamma = new Array(dim).fill(1);
+    this.beta = new Array(dim).fill(0);
+  }
+  // Call once before forward-passing a new sequence.
+  resetCache(seqLen) {
+    this._cache = new Array(seqLen);
+  }
+  // Normalize a single position's feature vector.
+  // pos must match the position index used in the corresponding backwardOne call.
+  predictOne(x, pos) {
+    const N = x.length;
+    const mean = x.reduce((s, v) => s + v, 0) / N;
+    const vari = x.reduce((s, v) => s + (v - mean) ** 2, 0) / N;
+    const std = Math.sqrt(vari + this.eps);
+    const x_norm = x.map((v) => (v - mean) / std);
+    this._cache[pos] = { x_norm, std };
+    return x_norm.map((xn, i) => this.gamma[i] * xn + this.beta[i]);
+  }
+  // Backprop through layer norm for one position.
+  //
+  // Given dL/dy (dOut), computes dL/dx:
+  //   Let D = dOut ⊙ γ
+  //   dL/dx_i = (1/std) * (D_i − mean(D) − x_norm_i * mean(D ⊙ x_norm))
+  //
+  // Also updates γ and β via SGD:
+  //   γ_i += lr * dOut_i * x_norm_i
+  //   β_i += lr * dOut_i
+  //
+  // SGD (not Adam) for γ/β: they are aggregated across all positions in the
+  // sequence (de-facto mini-batch update), so the gradient is already smoothed.
+  backwardOne(dOut, pos, lr) {
+    const { x_norm, std } = this._cache[pos];
+    const N = dOut.length;
+    for (let i = 0; i < N; i++) {
+      this.gamma[i] += lr * dOut[i] * x_norm[i];
+      this.beta[i] += lr * dOut[i];
+    }
+    const D = dOut.map((d, i) => d * this.gamma[i]);
+    const mD = D.reduce((s, v) => s + v, 0) / N;
+    const mDxn = D.reduce((s, d, i) => s + d * x_norm[i], 0) / N;
+    return D.map((d, i) => (d - mD - x_norm[i] * mDxn) / std);
+  }
+};
+
+// src/TransformerBlock.ts
+var TransformerBlock = class {
+  constructor({ d_model, nHeads, d_ff }) {
+    // Forward caches (needed for backprop)
+    this._X = null;
+    this._attnOut = null;
+    this._h1 = null;
+    this._ff1Pre = null;
+    // pre-ReLU
+    this._ff1Out = null;
+    // post-ReLU
+    this._ff2Out = null;
+    this.d_model = d_model;
+    this.d_ff = d_ff;
+    this.attn = new MultiHeadAttention(d_model, nHeads);
+    this.norm1 = new LayerNorm(d_model);
+    this.norm2 = new LayerNorm(d_model);
+    this.ff1 = new WeightMatrix(d_ff, d_model);
+    this.ff2 = new WeightMatrix(d_model, d_ff);
+    this.b1 = new Array(d_ff).fill(0);
+    this.b2 = new Array(d_model).fill(0);
+    this.b1Opts = Array.from({ length: d_ff }, () => new Adam());
+    this.b2Opts = Array.from({ length: d_model }, () => new Adam());
+  }
+  // ── Forward ───────────────────────────────────────────────────────────────
+  // X: seqLen × d_model  →  out: seqLen × d_model
+  predict(X) {
+    const seqLen = X.length;
+    const attnOut = this.attn.predict(X);
+    this.norm1.resetCache(seqLen);
+    const h1 = X.map((x, i) => {
+      const added = x.map((v, k) => v + attnOut[i][k]);
+      return this.norm1.predictOne(added, i);
+    });
+    const ff1Pre = h1.map(
+      (h) => this.ff1.W.map((row, k) => row.reduce((s, w, m) => s + w * h[m], this.b1[k]))
+    );
+    const ff1Out = ff1Pre.map((pre) => pre.map((v) => Math.max(0, v)));
+    const ff2Out = ff1Out.map(
+      (h) => this.ff2.W.map((row, k) => row.reduce((s, w, m) => s + w * h[m], this.b2[k]))
+    );
+    this.norm2.resetCache(seqLen);
+    const out = h1.map((h, i) => {
+      const added = h.map((v, k) => v + ff2Out[i][k]);
+      return this.norm2.predictOne(added, i);
+    });
+    this._X = X;
+    this._attnOut = attnOut;
+    this._h1 = h1;
+    this._ff1Pre = ff1Pre;
+    this._ff1Out = ff1Out;
+    this._ff2Out = ff2Out;
+    return out;
+  }
+  // ── Backward ──────────────────────────────────────────────────────────────
+  // dOut: seqLen × d_model  →  dX: seqLen × d_model
+  backward(dOut, lr) {
+    const seqLen = dOut.length;
+    const d_model = this.d_model;
+    const h1 = this._h1;
+    const ff1Out = this._ff1Out;
+    const ff1Pre = this._ff1Pre;
+    const dAdded2 = dOut.map((do_, i) => this.norm2.backwardOne(do_, i, lr));
+    const dFf1Out = dAdded2.map(
+      (da) => Array.from(
+        { length: this.d_ff },
+        (_, k) => this.ff2.W.reduce((s, row, m) => s + row[k] * da[m], 0)
+      )
+    );
+    const dW2 = Array.from(
+      { length: d_model },
+      (_, m) => Array.from(
+        { length: this.d_ff },
+        (_2, k) => dAdded2.reduce((s, da, i) => s + da[m] * ff1Out[i][k], 0)
+      )
+    );
+    const db2 = Array.from(
+      { length: d_model },
+      (_, m) => dAdded2.reduce((s, da) => s + da[m], 0)
+    );
+    this.ff2.update(dW2, lr);
+    for (let m = 0; m < d_model; m++)
+      this.b2[m] = this.b2Opts[m].step(this.b2[m], db2[m], lr);
+    const dFf1Pre = dFf1Out.map(
+      (d, i) => d.map((v, k) => ff1Pre[i][k] > 0 ? v : 0)
+    );
+    const dH1_fromFf = dFf1Pre.map(
+      (dp) => Array.from(
+        { length: d_model },
+        (_, m) => this.ff1.W.reduce((s, row, k) => s + dp[k] * row[m], 0)
+      )
+    );
+    const dW1 = Array.from(
+      { length: this.d_ff },
+      (_, k) => Array.from(
+        { length: d_model },
+        (_2, m) => dFf1Pre.reduce((s, dp, i) => s + dp[k] * h1[i][m], 0)
+      )
+    );
+    const db1 = Array.from(
+      { length: this.d_ff },
+      (_, k) => dFf1Pre.reduce((s, dp) => s + dp[k], 0)
+    );
+    this.ff1.update(dW1, lr);
+    for (let k = 0; k < this.d_ff; k++)
+      this.b1[k] = this.b1Opts[k].step(this.b1[k], db1[k], lr);
+    const dH1 = Array.from(
+      { length: seqLen },
+      (_, i) => dH1_fromFf[i].map((v, m) => v + dAdded2[i][m])
+    );
+    const dAdded1 = dH1.map((d, i) => this.norm1.backwardOne(d, i, lr));
+    const dAttnOut = dAdded1;
+    const dX_skip = dAdded1;
+    const dX_fromAttn = this.attn.backward(dAttnOut, lr);
+    const dX = Array.from(
+      { length: seqLen },
+      (_, i) => Array.from(
+        { length: d_model },
+        (_2, m) => dX_fromAttn[i][m] + dX_skip[i][m]
+      )
+    );
+    return dX;
+  }
+  // Attention weights from the last predict() — for visualization.
+  getAttentionWeights() {
+    return this.attn.getAttentionWeights();
+  }
+};
+
+// src/NetworkTransformer.ts
+var NetworkTransformer = class {
+  constructor(seqLen, options = {}) {
+    const {
+      vocabSize = 10,
+      d_model = 64,
+      nHeads = 4,
+      d_ff = 128,
+      nBlocks = 4,
+      nClasses = 9
+    } = options;
+    this.seqLen = seqLen;
+    this.vocabSize = vocabSize;
+    this.d_model = d_model;
+    this.nClasses = nClasses;
+    this.tokenEmb = new EmbeddingMatrix(vocabSize, d_model);
+    this.posEmb = new EmbeddingMatrix(seqLen, d_model);
+    this.blocks = Array.from(
+      { length: nBlocks },
+      () => new TransformerBlock({ d_model, nHeads, d_ff })
+    );
+    this.outputProj = new WeightMatrix(nClasses, d_model);
+    this.outputBias = new Array(nClasses).fill(0);
+    this.outBiasOpts = Array.from({ length: nClasses }, () => new Adam());
+  }
+  // ── Forward pass ──────────────────────────────────────────────────────────
+  // tokens: seqLen integer ids  →  seqLen * nClasses logits (flattened)
+  predict(tokens) {
+    const h = this._forward(tokens);
+    return h.flatMap(
+      (hi) => this.outputProj.W.map(
+        (row, c) => row.reduce((s, w, m) => s + w * hi[m], this.outputBias[c])
+      )
+    );
+  }
+  // ── Training step (online, one sample at a time) ───────────────────────────
+  // tokens:  seqLen integer ids
+  // targets: seqLen * nClasses values (e.g. one-hot per cell)
+  // mask:    optional boolean[seqLen] — only compute loss/gradients for
+  //          positions where mask[i] = true (e.g. empty cells in Sudoku)
+  // Returns: MSE loss over the masked positions.
+  train(tokens, targets, lr, mask) {
+    const h = this._forward(tokens);
+    const logits = h.map(
+      (hi) => this.outputProj.W.map(
+        (row, c) => row.reduce((s, w, m) => s + w * hi[m], this.outputBias[c])
+      )
+    );
+    let loss = 0;
+    let count = 0;
+    const dLogits = Array.from({ length: this.seqLen }, (_, i) => {
+      if (mask && !mask[i]) return new Array(this.nClasses).fill(0);
+      count++;
+      return Array.from({ length: this.nClasses }, (_2, c) => {
+        const t = targets[i * this.nClasses + c];
+        const p = logits[i][c];
+        loss += (p - t) ** 2;
+        return 2 * (p - t);
+      });
+    });
+    if (count > 0) loss /= count * this.nClasses;
+    const dH = Array.from(
+      { length: this.seqLen },
+      (_, i) => Array.from(
+        { length: this.d_model },
+        (_2, m) => dLogits[i].reduce((s, dl, c) => s + dl * this.outputProj.W[c][m], 0)
+      )
+    );
+    const dWout = Array.from(
+      { length: this.nClasses },
+      (_, c) => Array.from(
+        { length: this.d_model },
+        (_2, m) => dLogits.reduce((s, dl, i) => s + dl[c] * h[i][m], 0)
+      )
+    );
+    const dBout = Array.from(
+      { length: this.nClasses },
+      (_, c) => dLogits.reduce((s, dl) => s + dl[c], 0)
+    );
+    this.outputProj.update(dWout, lr);
+    for (let c = 0; c < this.nClasses; c++)
+      this.outputBias[c] = this.outBiasOpts[c].step(this.outputBias[c], dBout[c], lr);
+    let dX = dH;
+    for (let b = this.blocks.length - 1; b >= 0; b--)
+      dX = this.blocks[b].backward(dX, lr);
+    for (let i = 0; i < this.seqLen; i++) {
+      this.tokenEmb.update(tokens[i], dX[i], lr);
+      this.posEmb.update(i, dX[i], lr);
+    }
+    return loss;
+  }
+  // Attention weights from every block for visualization.
+  // Returns: nBlocks × nHeads × seqLen × seqLen  (nulls if not yet run)
+  getAttentionWeights() {
+    return this.blocks.map((b) => b.getAttentionWeights());
+  }
+  // ── Internal ──────────────────────────────────────────────────────────────
+  // Shared embedding + block forward pass.
+  _forward(tokens) {
+    let h = tokens.map((id, i) => {
+      const te = this.tokenEmb.get(id);
+      const pe = this.posEmb.get(i);
+      return te.map((v, m) => v + pe[m]);
+    });
+    for (const block of this.blocks)
+      h = block.predict(h);
+    return h;
+  }
+};
+
 // src/losses.ts
 function mse(predicted, actual) {
   return predicted.reduce((sum, p, i) => sum + (actual[i] - p) ** 2, 0) / predicted.length;
@@ -499,15 +1058,22 @@ function crossEntropyDeltaRaw(predicted, actual) {
 }
 export {
   Adam,
+  AttentionHead,
+  EmbeddingMatrix,
   LSTMLayer,
   Layer,
+  LayerNorm,
   Momentum,
+  MultiHeadAttention,
   Network,
   NetworkLSTM,
   NetworkN,
+  NetworkTransformer,
   Neuron,
   NeuronN,
   SGD,
+  TransformerBlock,
+  WeightMatrix,
   crossEntropy,
   crossEntropyDelta,
   crossEntropyDeltaRaw,
@@ -516,9 +1082,13 @@ export {
   linear,
   makeElu,
   makeLeakyRelu,
+  matMul,
   mse,
   mseDelta,
   relu,
   sigmoid2 as sigmoid,
-  tanh
+  softmax,
+  softmaxBackward,
+  tanh,
+  transpose
 };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dniskav/neuron",
-  "version": "0.1.6",
+  "version": "0.2.0",
   "description": "Minimal neural network from scratch — neuron, layer, network, backpropagation. No dependencies.",
   "main": "dist/index.js",
   "module": "dist/index.mjs",